SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We illustrate the Multi-ResNet, a novel U-Net architecture, in Figure 1 where we choose  $E_{i} = \mathrm{Id}_{V_{i}}$ . Practically speaking, the Multi-ResNet simplifies its encoder to have no learnable parameters, and simply projects to  $V_{i}$  on resolution  $i$ . The latter can for the example of Haar wavelets be realised by computing a multi-level Discrete Wavelet Transform (DWT) (or equivalently average pooling) over the input data [20]. Multi-ResNets allow us to save the parameters in the encoder, and instead direct them to bolster the decoder. In our experiments in §5.1, we compare Multi-ResNets to Residual U-Nets and find that for PDE surrogate modelling and image segmentation, Multi-ResNets yield superior performance to Residual U-Nets as Haar wavelets are apparently a good choice for  $\mathcal{V}$ , while for other problems, choosing Haar wavelets is suboptimal. Future work should hence investigate how to optimally choose  $\mathcal{V}$  for a problem at hand. To this end, we will discuss natural bases for  $\mathcal{V}$  and  $\mathcal{W}$  for specific problems in the remainder of this section.
3.2 U-Nets which guarantee boundary conditions
Next, our main goal is to show how to design U-Nets which choose  $\mathcal{W}$  in order to encode constraints on the output space directly into the U-Net architecture. This renders the U-Net more efficient as it no longer needs to learn the constraints from data. We consider an example from PDE surrogate modelling, approximating solutions to PDE using neural networks, a nascent research direction where U-Nets already play an important role [6], where our constraints are given boundary conditions and the solution space of our PDE. In the elliptic boundary value problem on  $\mathbb{X} = [0,1]$  [31], the task is to predict a weak (see Appendix A) PDE solution  $u$  from its forcing term  $f$  given by
$$
\Delta u = f, \quad u (0) = u (1) = 0, \tag {2}
$$
where  $u$  is once weakly differentiable when the equation is viewed in its weak form,  $f \in L^{2}(\mathbb{X})$  and  $\Delta u$  is the Laplacian of  $u$ . In contrast to Examples 1 and 2, we choose the decoder spaces as subspaces of  $W = \mathcal{H}_0^1$ , the space of one weakly differentiable functions with nullified boundary condition, a Hilbert space (see Appendix A), and choose  $V = L^{2}(\mathbb{X})$ , the space of square integrable functions. This choice ensures that input and output functions of our U-Net are in the correct function class for the prescribed problem. We now want to choose a basis to construct the subspaces  $\mathcal{V}$  and  $\mathcal{W}$  of  $V$  and  $W$ . For  $\mathcal{V}$ , just like in Multi-ResNets in Example 2, we choose  $V_{j}$  to be the Haar wavelet space of resolution  $j$ , an orthogonal basis. For  $\mathcal{W}$ , we also choose a refinable basis, but one which is natural to  $\mathcal{H}_0^1$ . In particular, we choose  $W_{i} = \mathrm{span}\{\phi_{k,j}: j \leq k, k = 1, \dots, 2^{i - 1}\}$  where
$$
\phi_ {k, j} (x) = \phi \left(2 ^ {k} x + j / 2 ^ {k}\right), \quad \phi (x) = 2 x \cdot 1 _ {[ 0, 1 / 2)} (x) + (2 - 2 x) \cdot 1 _ {[ 1 / 2, 1 ]} (x). \tag {3}
$$
This constructs an orthogonal basis of  $\mathcal{H}_0^1$ , illustrated in Figure 4, which emulates our design choice in Section 3.1 where the orthogonal Haar wavelet basis was beneficial as  $W$  was  $L^2$ -valued. Each  $\phi_{k,j}$  obeys the regularity and boundary conditions of our PDE, and consequently, an approximate solution from our U-Net obeys these functional constraints as well.
Figure 4: Refinement of an orthogonal basis for  $\mathcal{H}_0^1 = \operatorname{span}\{\phi_{0,0},\phi_{1,0},\phi_{1,1}\}$ . We visualise the graphs of basis functions defined in (3): [Left]  $\phi_{0,0} = \phi$ , [Top Right]  $\phi_{1,0}$ , and [Bottom Right]  $\phi_{1,1}$ . When increasing resolution, steeper triangular-shaped basis functions are constructed.
These constraints are encoded into the U-Net architecture and hence need not be learned from data. This generalised U-Net design paves the way to broaden the application of U-Nets, analogous to the choice of bases for Finite Element [32] or Discontinuous Galerkin methods [33].
Figure 5: U-Nets encoding the topological structure of a problem. [Left] A refinable Haar wavelet basis with basis functions on a right triangle,  $\phi_{i,j = 0} = \mathbb{1}_{\mathrm{red}} - \mathbb{1}_{\mathrm{blue}}$ . [Right] A sphere and a Möbius strip meshed with a Delaunay triangulation [35, 36]. Figures and code as modified from [37].
3.3 U-Nets for complicated geometries
Our framework further allows us to design U-Nets which encode the geometry of the input space right into the architecture. This no longer requires to learn the geometric structure from data and enables U-Nets for complicated geometries. In particular, we are motivated by tessellations, the partitioning of a surface into smaller shapes, which play a vital role in modelling complicated geometries across a wide range of engineering disciplines [34]. We here focus on U-Nets on a triangle due to the ubiquitous use of triangulations, for instance in CAD models or simulations, but note that our design principles can be applied to other shapes featuring a self-similarity property. We again are interested in finding a natural basis for this geometry, and characterise key components of our U-Net.
In this example, neither  $W$  nor  $V$  are selected to be  $L^2(\mathbb{X})$  valued on the unit square (or rectangle)  $\mathbb{X}$ . Instead, in contrast to classical U-Net design in literature,  $W = V = L^2(\triangle)$ , where  $\triangle$  is a right-triangular domain illustrated in Figure 5 [Left]. Note that this right triangle has a self-similarity structure in that it can be constructed from four smaller right triangles, continuing recursively. A refinable Haar wavelet basis for this space can be constructed by starting from  $\phi_{i,0} = \mathbb{1}_{\mathrm{red}} - \mathbb{1}_{\mathrm{blue}}$  for  $j = 0$  as illustrated in Figure 5 [Left]. This basis can be refined through its self-similarity structure to define each subspace from these basis functions via  $W_i = V_i = \operatorname{Span}\{\phi_{k,j} : j \leq k, k = 1, \ldots, 2^{i-1}\}$  (see Appendix A for details). In §5.2, we investigate this U-Net design experimentally. In Appendix A we sketch out how developing our U-Net on triangulated manifolds enables score-based diffusion models on a sphere [16] without any adjustments to the diffusion process itself. This approach can be extended to complicated geometries such as manifolds or CW-complexes as illustrated in Figure 5 [Right].
4 Why U-Nets are a useful inductive bias in diffusion models
U-Nets are the go-to neural architecture for diffusion models particularly on image data, as demonstrated in an abundance of previous work [8, 9, 10, 11, 12, 38, 39, 40, 41]. However, the reason why U-Nets are particularly effective in the context of diffusion models is understudied. Our U-Net framework enables us to analyse this question. We focus on U-Nets over nested Haar wavelet subspaces  $\mathcal{V} = \mathcal{W}$  that increase to  $V = W = L^2([0,1])$ , with orthogonal projection  $Q_j: W \mapsto W_i$  on to  $W_i$  corresponding to an average pooling operation  $Q_i$  [42] (see Appendix C). U-Nets with average pooling are a common choice for diffusion models in practice, for instance when modelling images [8, 9, 38, 39, 40]. We provide theoretical results which identify that high-frequencies in a forward diffusion process are dominated by noise exponentially faster, and how U-Nets with average pooling exploit this in their design.
Let  $X \in W$  be an infinite resolution image. For each resolution  $i$  define the image  $X_{i} = Q_{i}X \in W_{i}$  on  $2^{i}$  pixels which can be described by  $X_{i} = \sum_{k} X_{i}^{(k)}\phi_{k}$ , where  $\Phi = \{\phi_{k} : k = 1, \dots, 2^{i}\}$  is the standard (or 'pixel') basis. The image  $X_{i}$  is a projection of the infinite resolution image  $X$  to the finite resolution  $i$ . We consider the family of denoising processes  $\{X_{i}(t)\}_{i=1}^{\infty}$ , where for resolution  $i$ , the process  $X_{i}(t) = \sum_{k} X_{i}^{(k)}(t)\phi_{k} \in W_{i}$  is initialised at  $X_{i}$  and evolves according to the denoising diffusion forward process (DDPM, [8]) at each pixel  $X_{i}^{(k)}(t) := \sqrt{1 - \alpha_{t}}X_{i}^{(k)} + \sqrt{\alpha_{t}}\varepsilon^{(k)}$  for standard Gaussian noise  $\varepsilon^{(k)}$ . We now provide our main result (see Appendix A for technical details).
Theorem 2. For time  $t \geq 0$  and  $j \geq i$ ,  $Q_{i}X_{j}(t) \stackrel{d}{=} X_{i}(t)$ . Furthermore if  $X_{i}(t) = \sum_{j=0}^{i} \widehat{X}^{(j)}(t) \cdot \widehat{\phi}_{j}$  be the decomposition of  $X_{i}(t)$  in its Haar wavelet frequencies (see Appendix C). Each component  $\widehat{X}^{(j)}(t)$  of the vector has variance  $2^{j-1}$  relative to the variance of the base Haar wavelet frequency.
Figure 6: PDE modelling and image segmentation with a Multi-ResNet. [Left,Middle] Rolled out PDE trajectories (ground-truth, prediction,  $L^2$ -error) from the Navier-Stokes [Left], and the Shallow Water equation [Middle]. Figure and code as modified from [6, Figure 1]. [Right] MRI images from WMH with overlayed ground-truth (green) and prediction (blue) mask.
Theorem 2 analyses the noising effect of a forward diffusion process in a Haar wavelet basis. It states that the noise introduced by the forward diffusion process is more prominent in the higher-frequency wavelet coefficients (large  $k$ ), whereas the lower-frequency coefficients (small  $k$ ) preserves the signal. Optimal recovery of the signal in such scenario has been investigated in [43], where soft thresholding of the wavelet coefficients provides a good  $L^2$  estimator of the signal and separates this from the data. In other words, if we add i.i.d. noise to an image, we are noising the higher frequencies faster than the lower frequencies. In particular, high frequencies are dominated by noise exponentially faster. It also states that the noising effect of our diffusion on resolution  $i$ , compared to the variance of the base frequency  $i = 0$ , blows up the higher-frequency details as  $i \to \infty$  for any positive diffusion time.
We postulate that U-Nets with average pooling exploit precisely this observation. Recall that the primary objective of a U-Net in denoising diffusion models is to separate the signal from the noise which allows reversing the noising process. In the 'learning noise' or  $\epsilon$  recovery regime, the network primarily distinguishes the noise from the signal in the input. Yet, our analysis remains relevant, as it fundamentally pertains to the signal-to-noise ratio. Through average pooling, the U-Net discards those higher-frequency subspaces which are dominated by noise, because average pooling is conjugate to projection in a Haar wavelet basis [20, Theorem 2]. This inductive bias enables the encoder and decoder networks to focus on the signal on a low enough frequency which is not dominated by noise. As the subspaces are coupled via preconditioning, the U-Net can learn the signal which is no longer dominated by noise, added on each new subspace. This renders U-Nets a computationally efficient choice in diffusion models and explains their ubiquitous use in this field.
We conduct three main experimental analyses: (A) Multi-ResNets which feature an encoder with no learnable parameters as an alternative to classical Residual U-Nets, (B) Multi-resolution training and sampling, (C) U-Nets encoding the topological structure of triangular data. We refer to Appendix B.4 for our Ablation Studies, where a key result is that U-Nets crucially benefit from the skip connections, hence the encoder is successful and important in compressing information. We also analyse the multi-resolution structure in U-Nets, and investigate different orthogonal wavelet bases. These analyses are supported by experiments on three tasks: (1) Generative modelling of images with diffusion models, (2) PDE Modelling, and (3) Image segmentation. We choose these tasks as U-Nets are a go-to and competitive architecture for them. We report the following performance metrics with mean and standard deviation over three random seeds on the test set: FID score [44] for (1), rollout mean-squared error (r-MSE) [6] for (2), and the Sørensen-Dice coefficient (Dice) [45, 46] for (3). As datasets, we use MNIST [47], a custom triangular version of MNIST (MNIST-Triangular) and CIFAR10 [48] for (1), Navier-stokes and Shallow water equations [49] for (2), and the MICCAI 2017 White Matter Hyperintensity (WMH) segmentation challenge dataset [50, 51] for (3). We provide our PyTorch code base at https://github.com/FabianFalck/unet-design. We refer to Appendices B, and D for details on experiments, further experimental results, the datasets, and computational resources used.
Table 1: Quantitative performance of the (Haar wavelet) Multi-ResNet compared to a classical (Haar wavelet) Residual U-Net on two PDE modelling and an image segmentation task.
Dataset	Neural architecture	# Params.	r-MSE ↓ / Dice ↑
Navier-stokes128 × 128	Residual U-Net	34.5 M	0.0057 ± 2 · 10-5
Multi-ResNet, no params. added in dec. (ours)	15.7 M	0.0107 ± 9 · 10-5
Multi-ResNet, saved params. added in dec. (ours)	34.5 M	0.0040 ± 2 · 10-5
Shallow water96 × 192	Residual U-Net	34.5 M	0.1712 ± 0.0005
Multi-ResNet, no params. added in dec. (ours)	15.7 M	0.4899 ± 0.0156
Multi-ResNet, saved params. added in dec. (ours)	34.5 M	0.1493 ± 0.0070
WMH200 × 200	Residual U-Net	2.2 M	0.8069 ± 0.0234
Multi-ResNet, no params. added in dec. (ours)	1.0 M	0.8190 ± 0.0047
Multi-ResNet, saved params. added in dec. (ours)	2.2 M	0.8346 ± 0.0388
5.1 The role of the encoder in a U-Net
In §3.1 we motivated Multi-ResNets, Residual U-Nets with identity operators as encoders over Haar wavelet subspaces  $\mathcal{V} = \mathcal{W}$  of  $V = W = L^2 (\mathbb{X})$ . We analysed the role of the encoder as learning a change of basis map and found that it does not need to do any work, if  $\mathcal{V}$ , the initial basis, is chosen optimally for the problem at hand. Here, we put exactly this hypothesis derived from our theory to a test. We compare classical (Haar wavelet) Residual U-Nets with a (Haar wavelet) Multi-ResNet. In Table 1, we present our results quantified on trajectories from the Navier-stokes and Shallow water PDE equations unrolled over several time steps and image segmentation as illustrated in Figure 6. Our results show that Multi-ResNets have competitive and sometimes superior performance when compared to a classical U-Net with roughly the same number of parameters. Multi-ResNets outperform classical U-Nets by  $29.8\%$ ,  $12.8\%$  and  $3.4\%$  on average over three random seeds, respectively. In Appendix B.1, we also show that U-Nets outperform FNO [52], another competitive architecture for PDE modelling, in this experimental setting.
Figure 7: Preconditioning enables multi-resolution training and sampling of diffusion models.
For the practitioner, this is a rather surprising result. We can simplify classical U-Nets by replacing their parameterised encoder with a fixed, carefully chosen hierarchy of linear transformation as projection operators  $P_{i}$ , here a multi-level Discrete Wavelet Transform (DWT) using Haar wavelets, and identity operators for  $E_{i}$ . This 'DWT encoder' has no learnable parameters and comes at almost no computational cost. We then add the parameters we save into the decoder and achieve competitive and—on certain problems—strictly better performance when compared with a classical U-Net. However, as we show for generative modelling with diffusion models in Appendix B.1, Multi-ResNets are competitive with, yet inferior to Residual U-Nets, because the initial basis which  $\mathcal{V}$  imposes is suboptimal and the encoder would benefit from learning a better basis. This demonstrates the strength of our framework in understanding the role of the encoder and when it is useful to parameterise it, which depends on our problem at hand. It is now obvious that future empirical work should explore how to choose  $\mathcal{V}$  (and  $\mathcal{W}$ ) optimally, possibly eliminating the need of a parameterised encoder, and also carefully explore how to optimally allocate and make use of the (saved) parameters in Multi-ResNets [53, 54].
5.2 Staged training enables multi-resolution training and inference
Having characterised the self-similarity structure of U-Nets in §2, a natural idea is to explicitly train the U-Net  $\mathcal{U}$  on resolution  $i - 1$  first, then train the U-Net on resolution  $i$  while preconditioning on  $U_{i - 1}$ , continuing recursively. Optionally, we can freeze the weights of  $\mathcal{U}_{i - 1}$  upon training on resolution  $i - 1$ . We formalise this idea in Algorithm 1. Algorithm 1 enables training and inference of U-Nets on multiple resolutions, for instance when several datasets are available. It has two additional advantages. First, it makes the U-Net modular with respect to data. When data on a higher-resolution is available, we can reuse our U-Net pretrained on a lower-resolution in a principled way. Second, it enables to checkpoint the model during prototyping and experimenting with the model as we can see low-resolution outputs of our U-Net early.
Figure 8: U-Nets encode the geometric structure of data.
In Figure 7 we illustrate samples from a DDPM diffusion model [8] with a Residual U-Net trained with Algorithm 1. We use images and noise targets on multiple resolutions (CIFAR10/MNIST:
$\{4 \times 4, 8 \times 8, 16 \times 16, 32 \times 32\}$  as inputs during each training stage. We observe high-fidelity samples on multiple resolutions at the end of each training stage, demonstrating how a U-Net trained via Algorithm 1 can utilise the data available on four different resolutions. It is also worth noting that training with Algorithm 1 as opposed to single-stage training (as is standard practice) does not substantially harm performance of the highest-resolution samples: (FID on CIFAR10: staged training:  $8.33 \pm 0.010$ ; non-staged training:  $7.858 \pm 0.250$ ). We present results on MNIST, Navier-Stokes and Shallow water, with Multi-ResNets, and with a strict version of Algorithm 1 where we freeze  $E_i^\theta$  and  $D_i^\theta$  after training on resolution  $i$  in Appendix B.2.
5.3 U-Nets encoding topological structure
In §3.3, we showed how to design U-Nets with a natural basis on a triangular domain, which encodes the topological structure of a problem into its architecture. Here, we provide proof-of-concept results for this U-Net design. In Figure 8, we illustrate samples from a DDPM diffusion model [8] with a U-Net where we choose  $\mathcal{V}$  and  $\mathcal{W}$  as Haar wavelet subspaces of  $W = V = L^2 (\triangle)$  (see §3.3), ResNet encoders and decoders and average pooling. The
model was trained on MNIST-Triangular, a custom version of MNIST with the digit and support over a right-angled triangle. While we observe qualitatively correct samples from this dataset, we note that these are obtained with no hyperparameter tuning to improve their fidelity. This experiment has a potentially large scope as it paves the way to designing natural U-Net architectures on tessellations of complicated geometries such as spheres, manifolds, fractals, or CW-complexes.
Algorithm 1 Multi-resolution training and sampling via preconditioning.
Require: Boolean FREEZE.
1: for i←{1,...,J} do
2: if i>1 then
3: Precondition on Ui-1.
4: Train Ui
5: if FREEZE is True then
6: Freeze Eiθ and Diθ (fix parameters).
We provided a framework for designing and analysing U-Nets. Our work has several limitations: We put particular emphasis on Hilbert spaces as the decoder spaces. We focus on orthogonal wavelet bases, in particular of  $L^2(\mathbb{X})$  or  $L^2(\triangle)$ , while other bases could be explored (e.g. Fourier frequencies, radial basis functions). Our framework is motivated by subspace preconditioning, with requires the user to actively design and choose which subspaces they wish to precondition on. Our analysis of signal and noise concentration in Theorem 2 has been conducted for a particular, yet common choice of denoising diffusion model, with one channel, and with functions supported on one spatial dimension only, but can be straightforwardly extended with the use of a Kronecker product. Little to no tuning is performed how to allocate the saved parameters in Multi-ResNet in §5.1. We design and empirically demonstrate U-Nets on triangles, while one could choose a multitude of other topological structures. Lastly, future work should investigate optimal choices of  $\mathcal{U}$  for domain-specific problems.
Christopher Williams acknowledges support from the Defence Science and Technology (DST) Group and from a ESPRC DTP Studentship. Fabian Falck acknowledges the receipt of studentship awards from the Health Data Research UK-The Alan Turing Institute Wellcome PhD Programme (Grant Ref: 218529/Z/19/Z), and the Enrichment Scheme of The Alan Turing Institute under the EPSRC Grant EP/N510129/1. Chris Holmes acknowledges support from the Medical Research Council Programme Leaders award MC_UP_A390_1107, The Alan Turing Institute, Health Data Research, U.K., and the U.K. Engineering and Physical Sciences Research Council through the Bayes4Health programme grant. Arnaud Doucet acknowledges support of the UK Defence Science and Technology Laboratory (Dstl) and EPSRC grant EP/R013616/1. This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative. Saifuddin Syed and Arnaud Doucet also acknowledge support from the EPSRC grant EP/R034710/1.
The authors report no competing interests.
This research is supported by research compute from the Baskerville Tier 2 HPC service. Baskerville is funded by the EPSRC and UKRI through the World Class Labs scheme (EP/T022221/1) and the Digital Research Infrastructure programme (EP/W032244/1) and is operated by Advanced Research Computing at the University of Birmingham.