SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 7: Skip connections in U-Nets are crucial. We compare two Multi-ResNets, (Haar wavelet) Residual U-Nets and (Haar wavelet) Multi-ResNets, with and without skip connections, focussing on the two PDE modelling datasets (Navier-Stokes and Shallow water).
Neural architecture	# Params.	r-MSE ↓
Navier-stokes 128 × 128	Residual U-Net w/ skip con.	34.5 M	0.0057 ± 2 · 10-5
Residual U-Net w/o skip con.	34.5 M	0.0078 ± 10-5
Multi-ResNet w/ skip con.	34.5 M	0.0040 ± 2 · 10-5
Multi-ResNet w/o skip con.	34.5 M	0.0831 ± 0.1080
Shallow water 128 × 128	Residual U-Net w/ skip con.	34.5 M	0.1712 ± 0.0005
Residual U-Net w/o skip con.	34.5 M	0.4950 ± 0.02384
Multi-ResNet w/ skip con.	34.5 M	0.1493 ± 0.0070
Multi-ResNet w/o skip con.	34.5 M	0.6302 ± 0.01025
B.4.2 Ablation 2: The importance of preconditioning in U-Nets
In this section, we are interested in analysing the importance of preconditioning across subspaces in U-Nets. To this end, we analyse whether and how much a U-Net benefits from the dependency in the form of preconditioning between its input and output spaces  $\mathcal{V}$  and  $\mathcal{W}$ . In Table 8, we compare a Residual U-Net with multiple subspaces with a ResNet with only one subspace on Navier-Stokes and Shallow water, respectively. We obtain the ResNet on a single subspace from the Residual U-Net by replacing  $P_{i}$  as well as the (implicit) embedding operations with identity operators, and additionally feed zeros across the skip connections. This is because the function of the skip connections is superfluous due to them not being able to feed a compressed representation of the input. We note that the performance of the ResNet can potentially be improved by allocating the number of channels different to the increasing and decreasing choice in a U-Net, but we have not explored this. We further train on Shallow water for 10 epochs evaluating on the test set, and on Navier-Stokes for approximately  $400\mathrm{K}$  iterations evaluating on the validation set. This is due to significantly longer running times of the ResNets constructed as outlined above. We note that performance decreases very slowly after this point, and the trend of this experiment is unambiguously clear, not requiring extra training time.
The results in Table 8 show that preconditioning via the U-Net's self-similarity structure is a key reason for their empirical success. The Residual U-Net on multiple subspaces outperforms the ResNet
on a single subspace by a large margin. This also justifies the focus of studying preconditioning in this work.
Table 8: On the effect of subspace preconditioning vs. a plain ResNet. Quantitative performance of the (Haar Wavelet) Residual U-Net over multiple input and output subspaces compared to a ResNet on a single subspace, both trained on Navier-Stokes and Shallow water.
Dataset	Neural architecture	# Params.	r-MSE ↓
Navier-stokes	Residual U-Net (multiple subspaces)	34.5 M	0.0086 ± 9.7 · 10-5
128 × 128	ResNet (one subspace)	34.5 M	0.3192 ± 0.0731
Shallow water	Residual U-Net (multiple subspaces)	34.5 M	0.3454 ± 0.02402
96 × 192	ResNet (one subspace)	34.5 M	2.9443 ± 0.2613
B.4.3 Ablation 3: On the importance of the wavelet transform in Multi-ResNets
In Table 9 we analyse the effect of different wavelets for the DWT we compute in  $P_{i}$ . We selected examples from different wavelet families, all of them being orthogonal. A good overview of different wavelets which can be straight-forwardly used to compute DWT can be found on this URL: https://wavelets.pybytes.com/. On Navier-Stokes we analyse the effect of choosing Daubechies 2 and 10 wavelets systematically over different random seeds. On Shallow water we explore many different Wavelet families generally without computing random seeds. We find that the Wavelets we choose for  $P_{i}$  have some, but rather little impact on model performance. We lastly note that the downsampling operation has also been studied in [60] in the context of score-based diffusion models.
Table 9: On the importance of choosing different wavelets for the Discrete Wavelet Transforms (DWT) in Multi-ResNets. We report quantitative performance on the test set of Navier-Stokes and Shallow water.
Dataset	Neural architecture	# Params.	r-MSE ↓
Navier-stokes 128 × 128	Multi-ResNet w/ Haar wavelet DWT	34.5 M	0.0040 ± 2 · 10-5
Multi-ResNet w/ Daubechies 2 DWT	34.5 M	0.0041 ± 3 · 10-5
Multi-ResNet w/ Daubechies 10 DWT	34.5 M	0.0068 ± 1.1 · 10-4
Shallow water 96 × 192	Multi-ResNet w/ Haar wavelet DWT	34.5 M	0.1493 ± 0.0070
Multi-ResNet w/ Daubechies 2 DWT	34.5 M	0.1081
Multi-ResNet w/ Daubechies 10 DWT	34.5 M	0.1472
Multi-ResNet w/ Coiflets 1 DWT	34.5 M	0.1305
Multi-ResNet w/ Discrete Meyer DWT	34.5 M	0.1237
B.5 On the importance of preconditioning in residual learning: Synthetic experiment
We begin by providing a second example with  $w(v) = v^3$  in Figure 20, corresponding to Figure 2 in §1. Note that in comparison to the problem  $w(v) = v^2$  in Figure 2, the effect of the two preconditioners swaps in Figure 20: Using  $R^{\mathrm{pre}}(v) = v$  as the preconditioner produces an almost perfect approximation of the ground-truth function, while choosing  $R^{\mathrm{pre}}(v) = |v|$  results in deteriorate performance. This is because  $R^{\mathrm{pre}}(v) = v$  is a 'good' initial guess for  $w(v) = v^3$ , but a very 'poor' guess for  $w(x) = v^2$ , and vice versa for  $R^{\mathrm{pre}}(v) = |v|$ . This illustrates that the choice of using a good preconditioner is problem-dependent, and can be crucial to solve an optimisation problem at hand.
Figure 20: The importance of preconditioning. We learn the ground-truth functions  $w(v) = \{v^3\}$  using a ResNet  $R^{\mathrm{res}}(v) = R_{\ell}^{\mathrm{pre}}(v) + R(v)$  with preconditioners [Left]  $R_1^{\mathrm{pre}}(v) = v$  and [Right]  $R_2^{\mathrm{pre}}(v) = |v|$ .
Having discussed the results of our synthetic experiment (see Figs. 2 and 20) which demonstrates the importance of preconditioning, we here provide further details on the experimental setting. We train a ResNet  $R^{\mathrm{res}}(v) = R_{\ell}^{\mathrm{pre}}(v) + R(v)$  over a grid of values  $\{(v_i, w_i)\}_{i=1}^N$  where  $v_i \in [-1,1]$ ,  $w_i = w^{(k)}(v_i)$ ,  $N = 50$ , and  $w^{(1)}(v) = v^2$ ,  $w^{(2)}(v) = v^3$ ,  $R_{\ell}^{\mathrm{pre}}(v) = v$ ,  $R_{\ell}^{\mathrm{pre}}(v) = |v|$  depending on the experiment scenario indicated by the superscripts  $(k,\ell)$ , respectively. We do not worry about generalisation in this example and use the same train and test distribution. This functional relationship is easy to learn by any neural network  $R$  of sufficient size, which would overshadow the effect of preconditioning. To construct a ResNet which is 'just expressive enough' to learn  $w^{(k)}$ , we constrain the expressivity of  $R$  in two ways: First, we choose  $R$  to be a small fully-connected network with 3 linear layers (input, hidden, output) with  $[1,20,1]$  neurons, respectively, interleaved with the ReLU activation function. Second, inspired by invertible ResNets [61], we normalize the weights  $\theta_j$  of  $R$  such that  $\| \theta_j \| < 1$  for all  $j$  where  $\| \cdot \|$  is the Frobenius norm, and repeatedly apply  $R$  with  $D = 100$  times via weight-sharing. In practice,  $\| \theta_j \|$  will be close to 1, and as  $\| \theta_j \|$  intuitively measures the 'step size' of a neural network, we constrain the neural network to '100 steps of unit length'. As an alternative, one could limit the maximum number of training steps, which would indirectly constrain the model's expressiveness. We note that our reported results are robust across a variety of hyperparameter combinations leading to qualitatively the same conclusions.
C Background
C.1 Related work
U-Nets. The U-Net [1] is a go-to, state-of-the-art architecture across various research domains and tasks. Among the most influential U-Net papers are U-Net++ [3] and Attention U-Net [4] for image segmentation, 3D U-Net [62] for volumetric segmentation, U-Nets in diffusion models, starting with the first high-resolution demonstration in DDPM [8], the nnU-Net [63, 64], a U-Net which automatises parts of its design, a probabilistic U-Net [65], conditional U-Nets [66], U-Nets for cell analysis detection and counting [67], and U-Nets for road extraction [68]. A large number of adaptations and variants of the U-Net exist. We refer to [69] for an overview of such variants in the context of image segmentation for which the U-Net was first invented. In particular, many U-Net papers use a ResNet architecture. [70] find that a key improvement for the seminal U-Net [1] is to use residual blocks instead of feed-forward convolutional layers. Beyond their use on data over a squared domain, there exist custom implementations on U-Nets for other data types, for instance on the sphere [17], on graphs [18, 71] or on more general, differentiable 3D geometries [72]. However, we note that a framework which unifies their designs and details the components for designing U-Nets on complicated data types and geometries is lacking. This paper provides such a framework.
Our work directly builds on and is motivated by the paper by Falck & Williams [20] which first connected U-Nets and multi-resolution analysis [73]. This paper showed the link between U-Nets and wavelet approximation spaces, specifically the conjugacy of Haar wavelet projection and average pooling in the context of U-Nets, which our work crucially relies on. The design of U-Nets and their connection to wavelets has also been studied in [25, 26]. Falck & Williams further analysed the regularisation properties of the U-Net bottleneck under specific assumptions restricted to the analysis of auto-encoders without skip connections, and argued by recursion what additional information is carried across each skip without this being rigorously defined. Our work in this paper is however augmenting their work in various ways. We list five notable extensions: First, we provide a unified definition of U-Nets which goes beyond and encompasses the definition in [20], and highlights the key importance of preconditioning in U-Nets as well as their self-similarity structure. Importantly, this definition is not limited to orthogonal wavelet spaces as choices for  $\mathcal{V}$  and  $\mathcal{W}$ , hence enabling the use of U-Nets for a much broader set of domains and tasks. Second, based on this definition, our framework enables the design of novel architectures given a more thorough understanding of the various components in a U-Net. This is demonstrated on the elliptic PDE problem and data over a triangular domain. Third, we analyse the usefulness of the inductive bias of U-Nets in diffusion models, a novel contribution. Fourth, our experiments with Multi-ResNets and multi-resolution training and sampling, as well as on triangular data are novel. Fifth, [20] focusses particularly on the application of U-Nets in hierarchical VAEs, which this work is not interested in in particular. In summary, while [20] provided crucial components and ideas in U-Nets, our work is focused on a framework for designing and analysing general U-Nets, which enables a wide set of applications across various domains.
Miscellaneous. We further briefly discuss various unrelated works which use similar concepts as those in this paper. Preconditioning, initialising a problem with prior information so to facilitate solving it, is a widely used concept in optimisation across various domains [74, 75, 76, 77, 13, 14]. In particular, we refer to its use in multi-level preconditioning and Garlekin methods [78, 79, 80]. Preconditioning is also used in the context of score-based diffusion models [81]. Most notably, preconditioning is a key motivation in the seminal paper on ResNets [15]. While preconditioning is a loosely defined term, we use it in the context of this literature and its usage there.
The concept of 'dimensionality reduction' is widely popular in diffusion models beyond its use in U-Nets. For instance, Stable Diffusion and Simple Diffusion both perform a diffusion in a lower-dimensional latent space and experience performance gains [11, 60]. Stable Diffusion in particular found that their model learns a "low-dimensional latent space in which high-frequency, imperceptible details are abstracted away" [11]. It is this intuition that the U-Net formalises. Simple Diffusion also features a multi-scale loss resembling the staged training in Algorithm 1. Another paper worth pointing out learns score-based diffusion models over wavelet coefficients, demonstrating that wavelets and their analysis can be highly useful in diffusion models [42].
C.2 Hilbert spaces
A Hilbert space is a vector space  $W$  endowed with an inner-product  $\langle \cdot, \cdot \rangle$  that also induces a complete norm  $\| \cdot \|$  via  $\| w \|^2 = \langle w, w \rangle$ . Due to the inner-product, we have a notion of two vectors  $w_1, w_2 \in W$  being orthogonal if  $\langle w_1, w_2 \rangle = 0$ .
In Section 3 we have paid special attention to two specific Hilbert spaces:  $L^2([0,1])$  and  $\mathcal{H}_0^1([0,1])$ .
1. the space  $L^2([0,1])$  has as elements square-integrable functions and has an inner-product given by
$$
\langle f, g \rangle_ {L ^ {2}} = \int_ {0} ^ {1} f (x) g (x) d x; \text {a n d}, \tag {34}
$$
2. the space  $\mathcal{H}_0^1 ([0,1])$  has as elements once weakly differentiable functions which vanish at both zero and one, with inner-product given by
$$
\langle f, g \rangle_ {\mathcal {H} _ {0} ^ {1}} = \int_ {0} ^ {1} f ^ {\prime} (x) g ^ {\prime} (x) d x, \tag {35}
$$
where  $f', g' \in L^2([0,1])$  are the weak derivatives of  $f$  and  $g$  respectively.
If we have a sequence  $\{\phi_k\}_{k=0}^{\infty}$  of elements of  $W$  which are pairwise orthogonal and span  $W$ , then we call this an orthogonal basis for  $W$ . Examples of orthogonal bases for both  $L^2([0,1])$  and  $\mathcal{H}_0^1([0,1])$  are given in Section C.3.
C.3 Introduction to Wavelets
Wavelets are refinable basis functions for  $L^2([0,1])$  which obey a self-similarity property in their construction. There is a 'mother' and 'father' wavelet  $\phi$  and  $\psi$ , of which the children wavelets are derivative of the mother wavelet  $\phi$ . For instance, in the case of Haar wavelets with domain [0, 1] we have that  $\psi(x) = 1$  and the mother and children wavelets given by
$$
\phi_ {k, j} (x) = \phi (2 ^ {k} x + j / 2 ^ {k}), \quad \phi (x) = 1 _ {[ 0, 1 / 2)} (x) - 1 _ {[ 1 / 2, 1 ]} (x).
$$
Under the  $L^2$ -inner product, the family  $\{\phi_{k,j}\}_{j = 0,k = 1}^{2^j,\infty}$  forms an orthogonal basis of  $L^2([0,1])$ . Further, if we define the functions
$$
\tilde {\phi} _ {k, j} (x) = \int_ {0} ^ {x} \phi_ {k, j} (y) d y, \tag {36}
$$
then each of these functions is in  $\mathcal{H}_0^1 ([0,1])$ , and is further orthogonal with respect to the  $\mathcal{H}_0^1$ -inner-product, bootstrapped from the orthogonality of the Haar wavelets.
The discrete encoding of the Haar basis can be given by the kronecker product
$$
H _ {i} = \left( \begin{array}{c c} 1 & 1 \\ 1 & - 1 \end{array} \right) \otimes H _ {i - 1}, \quad H _ {0} = \left( \begin{array}{c c} 1 & 1 \\ 1 & - 1 \end{array} \right). \tag {37}
$$
To move between the pixel basis and the Haar basis natural to average pooling, we use the mapping defined on resolution  $i$  by  $T_{i}:V_{i}\mapsto V_{i}$  through
$$
T _ {i} \left(v _ {i}\right) = \Lambda_ {i} H _ {i} v _ {i}, \tag {38}
$$
where  $H_{i}$  is the Haar matrix above and  $\Lambda_{i}$  is a diagonal scaling matrix with entries
$$
\left(\Lambda_ {i}\right) _ {k, k} = 2 ^ {- i + j - 1}, \quad \text {i f} k \in \{2 ^ {j - 1}, \dots , 2 ^ {j} \}. \tag {39}
$$
To get this, we simply identify how to represent a piecewise constant function from its pixel by pixel function values to be the coefficients of the Haar basis functions. For example, in one dimension for an image with four pixels we can describe the function as the weighted sum of the four basis functions [Right], that take values  $\pm 1$  or zero. The initial father wavelet is set to be the average of the four pixel values, and the coefficients of the mother and children wavelets are chosen to be the local deviance's of the averages from these function values, as seen in Figure 21.
For further details on wavelets, we refer to textbooks on this topic [73, 29].
Figure 21: Modelling a 1D image with four pixel values as the weighted sum of Haar wavelet frequencies. The coefficients are such that the local averages of pixel values give the Haar wavelet at a lower frequency, hence average pooling is conjugate to basis truncation here.
C.4 Images are functions
An image, here a gray-scale image with squared support, can be viewed as the graph of a function over the unit square  $[0,1]^2$  [20]. We visualise this idea in 22, referring to [20, Section 2] and [19] for a more detailed introduction. Many other signals can likewise be viewed as It is hence natural to construct our U-Net framework on function spaces, and have  $E_{i}$  and  $D_{i}$  be operators on functions.
Figure 22: Images modelled as functions.
D Code, computational resources, datasets, existing assets used
Code. We provide our code base as well as instructions to reproduce the key results in this paper under MIT License at https://github.com/FabianFalck/unet-design. Our code base uses code from four Github repositories which are sub folders. For our diffusion experiments on MNIST and MNIST-Triangular, we directly build on top of the repository https://github.com/JTT94/torch_ddpm. For our diffusion experiments on CIFAR10, we directly build on top of the repository https://github.com/w86763777/pytorch-ddpm. For our PDE experiments on Navier-Stokes and Shallow water, we use the repository https://github.com/microsoft/pdearena [6]. For our image semgentation experiments on WMH, we are inspired by the repository
https://github.com/hongweilibran/wmh_ibbmTum [51], but write the majority of code ourselves. We note that the MIT License only pertains to the original code in our repository, and we refer to these four repositories for their respective licenses.
We extended each of these repositories in various ways. We list key contributions below:
- We implemented several Residual U-Net architectures in the different repositories.
- We implemented Multi-ResNets in each repository.
- We implemented the staged training Algorithm (see Algorithm 1), as well as its strict version which freezes parameters.
- We implemented logging with weights & biases, as well miscellaneous adjustments for conveniently running code, for instance with different hyperparameters from the command line.
Datasets. In our experiments, we use the following five datasets: MNIST [47], MNIST-Triangular, a custom version of MNIST where data is supported on a triangle rather than a square, CIFAR10 [48], Navier-stokes, Shallow water [49], and the MICCAI 2017 White Matter Hyperintensity (WMH) segmentation challenge dataset [50, 51]. These five datasets—with the exception of MNIST-Triangular—have in common that data is presented over a square or rectangular domain, possibly with several channels, and of varying resolutions. We refer to the respective repositories above where these datasets have already been implemented. For MNIST-Triangular, we provide our custom implementation and dataset class as part of our code base, referring to Appendix B.3 on how it is constructed. It is also worth noting that Navier-Stokes and Shallow Water require considerable storage. On our hardware, the two datasets take up approximately 120 GB and 150 GB unzipped, respectively.
Computational resources. This work made use of two computational resources. First, we used two local machines with latest CPU hardware and one with an onboard GPU for development and debugging purposes. Second, we had access to a large compute cluster with A100 GPU nodes and appropriate CPU and RAM hardware. This cluster was shared with a large number of other users.
To reproduce a single run (without error bars) in any of the experiments, we provide approximate run times for each dataset using the GPU resources: On MNIST and MNIST-Triangular, a single run takes about 30 mins. On CIFAR10, a single run takes several days. On Navier-Stokes and Shallow water, a single run takes approximately 1.5 days and 1 day, respectively.
Existing assets used. Our code base uses the following main existing assets: Weights&Biases [82] (MIT License), PyTorch [83] (custom license), in particular the torchvision package, pytorchwavelets [84] (MIT License), PyWavelets [85] (MIT License), pytorch_lightning [86] (Apache License 2.0), matplotlib [87] (TODO), numpy [88] (BSD 3-Clause License), tensorboard [89] (Apache License 2.0), PyYaml [90] (MIT License), tqdm [91] (MPLv2.0 MIT License), scikit-learn and sklearnn [92] (BSD 3-Clause License), and pickle [93] (License N/A). We further use Github Copilot for the development of code, and in few cases use ChatGPT [94] as a writing aid.