A Unified Framework for U-Net Design and Analysis
Authors: Christopher Williams
Fabian Falck
George Deligiannidis
Christopher C. Holmes
Arnaud Doucet
Saifuddin Syed
Venue: NeurIPS 2023 poster (2023)
Keywords: U-Net
ResNet
Multi-ResNet
Generalised U-Net
Wavelets
Diffusion models
Generative modelling
PDE Modelling
Image Segmentation
TL;DR: We present a framework for designing and analysing general U-Nets.

Abstract: U-Nets are a go-to neural architecture across numerous tasks for continuous signals on a square such as images and Partial Differential Equations (PDE), however their design and architecture is understudied. In this paper, we provide a framework for designing and analysing general U-Net architectures. We present theoretical results which characterise the role of the encoder and decoder in a U-Net, their high-resolution scaling limits and their conjugacy to ResNets via preconditioning. We propose Multi-ResNets, U-Nets with a simplified, wavelet-based encoder without learnable parameters. Further, we show how to design novel U-Net architectures which encode function constraints, natural bases, or the geometry of the data. In diffusion models, our framework enables us to identify that high-frequency information is dominated by noise exponentially faster, and show how U-Nets with average pooling exploit this. In our experiments, we demonstrate how Multi-ResNets achieve competitive and often superior performance compared to classical U-Nets in image segmentation, PDE surrogate modelling, and generative modelling with diffusion models. Our U-Net framework paves the way to study the theoretical properties of U-Nets and design natural, scalable neural architectures for a multitude of problems beyond the square.

A Unified Framework for U-Net Design and Analysis
Christopher Williams +1 Fabian Falck ,,1, George Deligiannidis
Chris Holmes 1,2 Arnaud Doucet
Saifuddin Syed
University Oxford "The Alan Turing Institute
(williams, fabiar an fal ck, deligian, cholmes, doucet, saifuddin syed) eyyettssttts estats. ox ac
Abstract
U-Nets are go-to neural architecture across numerous tasks for continuous signals on square such as images and Partial Differential Equations (PDE), however their design and architecture is understudied. In this paper, we provide a framework for designing and analysing general U-Net architectures. We present theoretical results which characterise the role of the encoder and decoder in U-Net, their high-resolution scaling limits and their conjugacy to ResNets via precondition- ing. We propose Musti-ResNets, U-Nets with simplified, wavelet-based encoder without learnable parameters. Further, we show how to design novel U-Net ar- chitectures which encode function constraints, natural bases, the geometry of the data. In diffusion models, our framework enables us to identify that high- frequency information is dominated by noise exponentially faster, and show how U-Nets with average pooling exploit this. In our experiments, we demonstrate how Multi-ResNets achieve competitive and often superior performance compared classical U-Nets in image segmentation, PDE surrogate modelling, and generative modelling with diffusion models. Our U-Net framework paves the way study the theoretical properties of U-Nets and design natural, scalable neural architectures for multitude of problems beyond the square.
Introduction
U-Nets (see Figure 1) are central architecture in deep learning for continuous signals. Across many tasks as diverse as image segmentation [1, 2, 3, 4, 5], Partial Differential Equation (PDE) surrogate modelling [6, 7] and score-base based diffusion models [8, 9, go-to architecture yielding state-of-the-
U2(22)
U/
V2
10, 11, 12], U-Nets art performance. In spite of their enormous success, a framework for U-Nets which characterises for instance the specific role of the encoder and decoder in a U-Net or which spaces these operate on is lacking. In this work, provide such a framework for
W
D2
X--- Soep
U-Nets. This allows us to design U-Nets for data beyond a square domain, and enable problem, for instance a natural basis, functional constraints, or knowledge about its topology, into the neural architecture.
W
Kikpp connection
incorporate prior knowledge about
Po
U
Burleenwck
The importance preconditioning. We begin by illustrating the importance of the design principle of U-Nets: precon- ditioning. Preconditioning informally means that initialising an optimisation problem with a 'good" solution greatly benefits learn- ing [13, 14]. Consider a synthetic example using ResNets [15] which are natural in the context of
Figure 1: A resolution 2 U-Net (Def. 1). E, Idy, this is Multi-ResNet (see Def. 3).
Equua contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

U-Nets as we will show in $2.3: we are interested in learning ground-truth mapping w. V + W and w(()) = v² over V [-1,1 and W = using a ResNet R (v) = R (2) ) + R(v) where Rore, R: V + W. In Figure 2 [Left] we learn standard ResNet with [PPr(()) =
grid of values from V W, i.e. with inputs v2 E V and regression labels w; = w(()) of. In contrast, we train ResNet with Rote(u) = |w| [Right] with the same number of parameters and iterations. Both networks have been purposely designed to weakly expressive (see Appendix B.5 for details). The standard ResNet [Left] makes a poor approximation of the function, whilst the other ResNets [Right] approximation is nearly perfect. This is because (V = |0] is a 'good' initial guess or preconditioner for w(v) v², but (*) encoding good preconditioning into a neural network architecture and motivates us studying how preconditioning is used in U-Net design.
is 'bad' one. This shows the importance of
In this paper, we propose a mathematical framework for designing and analysing general U-Net architec- tures. We begin with comprehensive definition of a U-Net which characterises its components and iden- tifies its self-similarity structure which established via preconditioning. Our theoretical results delineate the role of the encoder and decoder and identify the subspaces they operate on. We then focus on ResNets as natural building blocks for U-Nets that enable flex- ible preconditioning from lower-resolutions inputs. Figure 2:
wit of sparddmm
wit of gpanntrutt Rb0 lopp odmatoon (pwooddiioeed
Mar semdimo Amfo wip.-eeeetetrtroeooo
The importance of preconditioning. Our U-Net framework paves the way to designing U-Nets which can model distributions over com- plicated geometries beyond the square, for instance CW-complexes the sphere [16], without any changes to the diffusion model itself (see Appendix A). It allows us to enforce problem-specific constraints through the architecture, such as boundary conditions of a PDE or a natural basis for problem. We also analyse why U-Nets with average pooling are a natural inductive bias in diffusion models.
manifolds, or diffusions on
More specifically, our contributions are as follows: (a) We provide the first rigorous definition of U-Nets, which enables to identify their self self-similarity structure, high-resolution scaling limits, and conjugacy to ResNets via preconditioning. (b) We present Multi-ResNets, a novel class of U-Nets over hierarchy of orthogonal wavelet spaces of (X), for compact domain X, with no learnable parameters in its encoder. In our experiments, Multi-ResNets yield competitive and often superior results when compared to classical U-Net in PDE modelling, image segmentation, and generative modelling with diffusion models. We further show how to encode problem-specific information into U-Net. In partieular, we design U-Nets incorporating a natural basis for a problem which enforces boundary conditions on the elliptic PDE problem, and design and demonstrate proof-of-concept experiments for U-Nets with a Haar wavelet basis over a triangular domain. (c) In the context of diffusion models, we analyse the forward process in a Haar wavelet basis and identify how high- frequency information dominated by noise exponentially faster than lower-frequent terms. We show how U-Nets with average pooling exploit this observation, explaining their go-to usage.
U-Nets: Neural networks via subspace preconditioning
The goal of this section is to develop mathematical framework for U-Nets which introduces the fundamental principles that underpin its architecture and enables us to design general U-Net architectures. All theoretical results are proven in Appendix A. We commence by defining the U-Net.
2.1 Anatomy of U-Net
Definition U-Net.
Let V and W be measurable spaces. A U-Ner U = D,PUU,
comprises six components:
1. Encoder subspaces:
(V are nested subsets of V such that lim V1 = V.
2. Decoder subspaces: W = =W are nested subsets of W such that lim W. = W.
3. Encoder operators: E = (E,) where E, V; + V; denoted Ex(v;) = 0z
Decoders operators: D = =D where D, : W x V + W; at resolution denoted D(w.-ilu:). The i component is called the skip connection.

---

5. Projection operators: P = (P.) where P. : V + V,, such that Pi(e,) = for &s E V,,
6. Bottleneck: Ug is the mapping Uo Vo + Wo. enabling a compressed representation of the input.
The U-Net of resolution i is the mapping U, : V2 + W; defined through the recursion (see Figure 3); Ui(1:) = Di(U,-1 (6:))
i= 1, 2,
We illustrate the Definition of a U-Net in Figure 1. Definition includes a wide array of commonly used U-Net architectures, from the seminal U-Net [1]. through to modern adaptations used in large-scale diffusion models [8, 9, 10, 11, 12], operator learning U-Nets for PDE modelling
(i((e)
E
Vi
[6, 7]. and custom designs on the sphere or graphs [17. 18]. Our framework also comprises models with multiple channels (for instance RGB in images) by choosing V; and W, for a given resolution i be product spaces V/2 = k=1 V and W = @
W.
Dj
M
M 0 Wi,, for M channels k=1
W:+:
Un-
Vi-
when necessary".. Remarkably, despite their widespread use, to the best of our knowledge, our work presents the first formal definition of a U-Net. Definition not only expands the scope of U-Nets beyond problems confined to squared domains, but also naturally incorporates problem-speciee information such as a natural basis, boundary conditions or topological structure as we will show in $3.2 and 3.3. This paves the way for designing inherently scalable neural network architectures capable of handling complicated geometries, for instance manifolds or CW-complexes. In the remainder of the section, we discuss and characterise the components of U-Net.
Figure 3: Recursive struc- ture of a U-Net.
Encoder and decoder subspaces. We begin with the spaces in V and W which a U-Net acts Current literature views U-Nets as learnable mappings between input and output tensors. In contrast, this work views U-Nets and their encoder and decoder as operators on spaces that must be chosen suit task. In order to perform computations in practice, we must restrict ourselves to subspaces Vú and W; of the potentially
spaces V and W. For instance, if our U-Net is modelling mapping between images, which can be viewed as bounded functions over squared domain X = [o. 1]? it is convenient to choose V and W as subspaces of L (X), the space of square-integrable functions [19, 20] (see Appendix C). Here, a data point wi in the decoder space W; is represented by the coefficients G.j where w; Σ, Ci.Eii and {ei.):: is a (potentially orthogonal) basis for W. We will consider precisely this case in $3.1. The projected images on the subspace Vi, W; are still functions, but piece-wise constant ones, and we store the values these functions obtain as 'pixel' tensors in our computer [20] (see Figure 22 in Appendix A).
Role of encoder, decoder, projection operators. In spite of their seemingly symmetric nature and contrast to common understanding, the roles of the encoder and decoder in a U-Net are fundamentally different from each other. The decoder resolution i learns the transition from Wi-1 to W; while incorporating information from the encoder on V/2 via the skip connection. The encoder Eg can be viewed change of basis mapping on the input space Vi at resolution i and is not directly tied to the approximation the decoder makes on W, This learned change of basis facilitates the decoder's 's approximation on W. In 13.1, we will further extend understanding of the encoder and discuss its implications for designing U-Nets. The projection operators serve to extract a compressed input. They are selected to suit task-specific needs, such as pooling operations (e.g. average pooling, max pooling) in the context of images, or orthogonal projections if V is Hilbert space. Note that there is no embedding operator', the operator facilitating the transition explicitly defined as it is invariably the natural inclusion of W into W..
higher-resolution space,
Self-similarity of U-Nets via preconditioning. The key design principle of U-Net is precondi- tioning. The U-Net of resolution i - 1, U,-1 makes an approximation on W;-1 which is input of and preconditions Uj. Preconditioning facilitates the transition from W to W; for the decoder. In the encoder, the output of P.- E, is the input of U--1- When our underlying geometry is refinable (such square), we may use refinable set of basis functions. In the standard case of U-Net with average pooling on a square domain, our underlying set of basis functions are (Haar) wavelets (see $2.3) - refinable basis functions defined on a refinable geometry. This preconditioned design of
e 'Implementation wise this corresponds to taking the Kronecker product across the spaces Vi,1 and W.... respectively.
In practice, if we for instance learn a transposed convolution operation, it can be equivalently expressed as a standard convolution operation composed with the natural inclusion operation.
3

---

U-Nets reveals self-similarity structure (see Figure 3) inherent to the U-Net when the underlying space has a refinable geometry. This enables both an efficient multi-resolution approximation of U-Nets [20] and makes them modular, as a U-Net on resolution U-Net on resolution i. We formalise this notion of preconditioning in Proposition 1 in $2.3.
is a coarse approximation for
2.2 High-resolution scaling behavior of U-Nets
Given equation (I), it is clear that the expressiveness of a U-Net 24 is governed by the expressiveness of its decoder operators D. If each D; is a universal approximator [21], then the corresponding U-Nets U, likewise have this property. Assuming we can represent any mapping U² Vú + W as a U-Net, our goal now to comprehend the role of increasing the resolution in the design of U4 and to discern whether any function U V + W can be represented as high nigh-resolution limit lim;-+o Uj of a U-Net. We will explore this question in the context of regression problems of increasing resolution.
To obtain a tractable answer, we focus on choosing W as a Hilbert space, that is W equipped with an inner product. This allows to define W an increasing sequence of orthogonal subspaces of W. Possible candidates for orthogonal bases include certain Fourier frequencies, wavelets (of given order) radial basis functions. The question of which basis is optimal depends on our problem and data at hand: some problems may be hard in one basis, but easy in another. In $4, Haar wavelets are a convenient choice. Let define S as infinite resolution data in V x W and S, the finite resolution projection in V2 x W; comprising of (1,, w:) = (P) Q(ww) for each (v, w) E S. Here, P. is the U-Net projection onto Vú, and Q: W + W; the orthogonal projection onto W,. Assume U,² and U are solutions to the finite and infinite resolution regression problems ming, eJ; Ci(U; (S,) and minger L(U)S) (S) respectively, where F; and F represent the sets of measurable functions mapping V; + W, and V + W. Let L; and L denote the I losses:
Σ C.UU,SS)) = (1EL.))ESS
||1; - U.(ril. C(U]S)? = S
Σ
- Uv))1.
(w,e)ES
The following result analyses the relationship between U/ and U as i - co where Liys and bi are the losses above conditioned on resolution j (see Appendix A).
Theorem 1. Suppose Uü and U solutions of the T.2 regression problems above. Then, Ca)(U;) S (1)(U*) with equality as i + co. Further, if QUU* is Vi-measurable, then U,+ = QUU minimises Li
Theorem states that solutions of the finite resolution regression problem converge to solutions of the infinite resolution problem. It also informs us how choose VÚ2 relative W.. we have a V where for the decoders on W, the W. component of U*, QUU*, relies solely on the input up to resolution i, then the prediction from the infiniteddimensional U-Net projected to W; can be made by the U-Net of resolution i. The optimal choice of V, must be expressive enough to encode the information necessary to learn the W; component of U This suggests that if Vú lacks expressiveness, we risk efficiency in learning the optimal value of U.. However, if V2 is too expressive, no additional information is gained, and we waste computational effort. Therefore, we should choose V; to encode the necessary information for learning information on resolution i. For example, when modelling images, if we are interested in low-resolution features on W;, high-resolution information is extraneous, as we will further explore in $4 in the context of diffusion models.
2.3
U-Nets are conjugate to ResNets
Next, our goal is to understand why ResNets are a natural choice in U-Nets. We will uncover conjugacy between U-Nets and ResNets. We begin by formalising ResNets in the context of U-Nets.
Definition ResNet, Residual U-Nets. Given measurable space X and a vector space Y.
mapping R X + Y is defined a ResNet preconditioned on Rere X + Y if R(xx) (pre(x)) + Rres (3r), where (e= = R(x)= - Rere (z) is the residual of R. A Residual U-Nel is a U-Net 1 where W. y are sequences of vector spaces, and the encoder and decoder operators E, DD are ResNets preconditioned on E (o:) = v2 and D. (w;-1 (v:) = w;-1, respectively.
A preconditioner initialises ResNet, then the ResNet learns the residual relative to it. The difficulty of training a ResNet scales with the deviation from the preconditioner, as we saw in our synthetic experiment in $1. In U-Nets, ResNets commonly serve encoder and decoder operators. In encoders,

---

preconditioning on the identity on V; allow the residual encoder E to learn a change of basis for V,, which we will discuss in more detail in $3.1. In decoders, preconditioning on the identity on the lower resolution subspace - allows the residual decoder D to learn from the lower resolution and the skip connection. Importantly, ResNets can compose to form new ResNets, which, combined with the recursion (1), implies that residual U-Nets are conjugate to ResNets.
Proposition 1. If u is residual U-Net, then U₃ is a ResNet preconditioned U (v:) U₁- U--1(8;-1), where 0,-1 = (..))))
Proposition states that a U-Net at resolution i ResNet preconditioned on U-Net of lower resolution. This suggests that U learns the information arising from the resolution increase. We will discuss the specific case E; Idy, and (u; enables us to interpret lim U; from $2.2 as a ResNet's "high-resolution" scaling limit. This is a new scaling regime for ResNets, different to time scaled Neural ODEs [22], and warrants further exploration in future work. Finally, we provide an example of the common Residual U-Net.
U, (P.- (2;)) in $3.1. Proposition also
Example 1. Haar Wavelet Residual U-Net. A Haar wavelet residual U-Net U is residual U-Net where: V = W = (x))), V = W; are multi-resolution Haar wavelet spaces (see Appendix C.3), and P. projy; is the orthogonal projection.
We design Example 1 with images in mind, noting that similar data over a squared domain such as PDE (see $5) are also applicable to this architecture. We hence choose Haar wavelet [23] subspaces of (X). the space of square integrable functions and a Hilbert space, and use average pooling as the projection operation P. [20]. Haar wavelets will in particular be useful to analyse why U-Nets are good inductive bias in diffusion models (see $4). U-Net design and their connection to wavelets has also been studied in [24, 25, 26).
3 Generalised U-Net design
In this section, we provide examples of different problems for which our framework can define a natural U-Net architecture. Inspired by Galerkin subspace methods [27], our goal is to use our framework to generalise the design of U-Nets beyond images over square. Our framework also enables us to encode problem-specifis information into the U-Net, such as a natural basis or boundary conditions, which it no longer needs to learn from data, making the U-Net model more efficient.
3.1 Multi-ResNets
closer look at the U-Net encoder. To characterise U-Net in Definition must in particular choose the encoder subspaces V. This choice depends on our problem at hand: for instance, if the inputs are images, choosing Haar wavelet subspaces is most likely favourable, because we can represent and compress images in a Haar wavelet basis well, noting that more complex (orthogonal) wavelet bases are possible [28. 29]. What if we choose V/ unfavourably? This is where the encoder comes in. While the encoder subspaces V define an initial basis for our problem, the encoder learns a change of basis map to a new, implicit basis é (v( which is more favourable. This immediately follows from Eq. (1) since U; acts on Vú through 6 The initial subspaces V can hence be viewed as a prior for the input compression task which the encoder performs.
Given our initial choice of the encoder subspaces V, the question whether and how much work the encoders & have to do depends on how far away our choice is from the optimal choice 1> for our problem. This explains why the encoders Es are commonly chosen to ResNets preconditioned on the identity E.Pre = Idy;, allowing the residual encoder ETT to learn a change of basis. If we had chosen the optimal sequence of encoder subspaces, the residual operator would not have to do any work; leaving the encoder equal to the precondition E; Idy, It also explains why in practice, encoders are in some cases chosen significantly smaller than the decoder [30]. as ResNet encoder need not do much work given good initial choice of V. It is precisely this intuition which motivates our second example of U-Net, the Multi-Res(olution Res(idual) Network
Definition 3. Musti-
Netti-ResNet is a residual U-Net with encoder F, Idy,
Example 2. Haar Wavelet Multi-ResNets A Haar Wavelet Multi-ResNer is a Haar Wavelet Residual U-Net with encoder E, Idy.
5

---

We illustrate the Multi-ResNer a novel U-Net architecture, in Figure where we choose E; Idy, simplifies its encoder to have learnable parameters, and
Practically speaking, the simply projects to Vi 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-Res 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 V, while for other problems, choosing Haar wavelets is suboptimal. Future work should hence investigate how to optimally choose for a problem at hand. To this end, we will discuss natural bases for V and W for specific problems in the remainder of this section.
to Residual
3.2 U-Nets which guarantee boundary conditions
Next, our main goal show how design U-Nets which choose 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 X = [0, 1] [31]. the task is to predict a weak (see Appendix A) PDE solution u from its forcing term given by
Au f:
20)) u(1) 0,
where u is once weakly differentiable when the equation is viewed its weak form, f E (X) and Au is the Laplacian of u. In contrast Examples and 2, we choose the decoder spaces as subspaces of W = =]., the space of one weakly differentiable functions with nullified boundary condition, Hilbert space (see Appendix A), and choose V L (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 choose basis to construct the subspaces V and W of V and W. For V, just like in Multi-Re ResNets in Example 2, we choose Vj to be the Haar wavelet space of resolution j. an orthogonal basis. For W. we also choose refinable basis, but one which natural to H-. In particular, we choose W. span( ØR.j j S k k- = ,,...,2" 21 where
(3) This constructs an orthogonal basis of Hg. illustrated in Figure 4, which emulates our design choice in Section 3.1 where the orthogonal Haar wavelet basis was beneficial as W was valuude Each oR.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.
(x) ((2'x + j/2*),
<(x) 2x (x) + (2 2x) (1/2
0.5
0.5
1
0.5
0.5
0.5
0.5
1
Figure 4: Refinement of orthogonal basis for H4] span (do,o 01.0, $1,1). We visualise the graphs of basis functions defined in (3): [Left] ø0,0 = o, [Top Right] ø1,0, and [Bottom Right] ø1,1- When increasing resolution, steeper saanuula-shaape basis functions 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 U-Nets, analogous to the choice of bases for Finite Element [32] or Discontinuous Galerkin methods [33].

---

Wo
Wi
Figure 5: U-Nets encoding the topological structure of a problem. [Left] A refinable Haar wavelet basis with basis functions on a right triangle. Φt,5=0 Ledd L [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 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, of a surface into smaller shapes, which play a vital role in modelling complicated geometries across wide range of engineering disciplines [34]. We here focus on U-Nets ubiquitous use of triangulations, for instance in CAD models or simulations, but note that our design principles can be applied other shapes featuring a self-similarity property. We again are interested in finding natural basis for this geometry, and characterise key components of our U-Net.
are motivated by tessellations, the partitioning
triangle due to the
In this example, neither W nor V are selected to be (X) valued the unit square (or rectangle) X. Instead, in contrast to classical U-Net design in literature, W = V = L2()), where A is a trihntrrangular domain illustrated in Figure [Left]. Note that this right triangle has self-similarity structure in that it can be constructed from four smaller right triangles, continuing recursively. A refinable Haar wavelet basis for this space be constructed by starting from di.n red 1 biue for j = O as illustrated in Figure 5 [Left]. This basis can be refined through its self-similarity structure to define each subspace from these basis funetions via W; V,
Opan(os.. : j k k 1, ...... 1.....22-1] (see Appendix A for details). In $5.2, 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 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 [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 demon- strated in 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 to analyse this question. We focus
U-Nets over nested Haar wavelet subspaces V = W that increase to V = W = [[0 1)), with orthogonal projection Qj W + W; on to W; corresponding 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 fighfrrequencies in forward diffusion process are dominated by noise exponentially faster, and how U-Nets with average pooling exploit this in their design.
an average pooling operation Q: [42] (see Appendix C). U-Nets with
Let X E W be an infinite resolution image. For each resolution define the image X; = Qi E W on 2' pixels which can be described by X, = E. (kk) ( ΦE. where + = { : k = 1,
1, ...., 2 is the standard (or 'pixel') basis. The image X; is a projection of the infinite resolution image to the finite resolution i. We consider the family of denoising processes {t) where for resolution i,
the process (t) = Σn x (t) E W; is initialised at X, and evolves according to the denoising diffusion forward process (DDPM, [8]) at each pixel x((²) (1) VI- xxx)(1) + (arg(()) for standard Gaussian noise ((k), We now provide our main result (see Appendix A for technical details).
Theorem 2. For time t 2 and j 2 i, (t),(t & X,(t). Furthermore if (i) Ei (t). be the decomposition of X.(t) in its Haar wavelet frequencies (see Appendix C). Each component ( of the vector has variance 2 relative to the variance of the base Haar wavelet frequency.
7

---

Ground COUP
EnTe Fredcco
Figure 6: PDE modelling and image segmentation with a Multi-ResNet. [iddle Rolled out PDE trajectories (ground-truth, prediction, r.--error) from the Navier-Stokes [Left], and the Shallow Water equation [Middle]. Figure and code as modified from [6, Figure 1J. [Right] MRI images from WMH with overlayed ground-truth (green) and prediction (blue) mask.
Theorem analyses the noising effect of a forward diffusion process in Haar wavelet basis. states that the noise introduced by the forward diffusion process is more prominent in the higher-frequeney wavelet coefficients (large k), whereas the lower-freyyency 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 good 2 estimator the signal and separates this from the data. In other words, if we add i.i.d. noise to image, the lower frequencies. In particular, high frequencies are dominated by noise exponentially faster. also states that the noising effect of our diffusion on resolution i, compared to the variance of the base frequency = 0, blows up the
noising the higher frequencies faster than
details as i - oc for any positive diffusion time.
We postulate that U-Nets with average pooling exploit precisely this observation. Recall that the primary objective U-Net in denoising diffusion models to separate the signal from the noise which allows reversing the noising process. In the 'learning noise' or recovery regime, the network primarily distinguishes the noise from signal in the input. Yet, our analysis remains relevant, as fundamentally pertains to the signal-to-noiss ratio. Through average pooling, the U-Net discards those gigher-frequecyy subspaces which are dominated by noise, because average pooling is conjugate to projection in 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 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 in this field.
5 Experiments
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 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 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 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-DDi coefficient (Dice) [45, 46] for (3). As datasets, we use MNIST [47]. custom triangular version of MNIST (MNIST- -EIISSTTTTT angular) 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 Py Torch code base at https:/github..com wwt.e/oooaaaaaaaaaiaaaiaii Fale//unet-design. We refer Appendices B, and D for details on experiments, further experimental results, the datasets, and computational resources used.
8

---

Table Quantitative performance of the (Haar wavelet) Multi-ResNet compared a classical (Haar wavelet) Residual U-Net on two PDE modelling and an image segmentation task.
Dataset
Neural architecture
# Params.
r-MSE // Dicc t
Navier-
stokes
128 x28
Residual U-Net
34.5 M
0.0057 + 2- 10
Multi-ResNer no params, added in dec. (ours)
15.7 M
0.0107 + 9- 10-5
Multi- RestiR saved params, added in dec. (ours)
34,5 M
0.0040 + 2 10
Shallow
water
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
96
Multi-Res saved params. added in dec. (ours)
34.5 M
0.1493 + 0.0070
200 200
Residual U-Net
2.2 M
0.8069 + 0.0234
WMH
Mustt,-ReNeet, 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 U-Net
Residual U- Nets with identity operators as encoders over Haar wavelet subspaces V W of V W (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 V. the initial basis, is chosen optimally for the problem hand. Here, we put exactly this hypothesis derived from our theory to test. We compare classical (Haar wavelet) Resid- ual U-Nets with (Haar wavelet)
In $3.1 we motivated
CIFAR10
In Table 1, we present our results quantified on tra- jectories from the Navier-stoke and Shallow water PDE equations unrolled over several time steps and image segmentation as illustrated in Fig- ure 6. Our results show that Multi-ResNett have competitive and sometimes superior performance when compared to a classical U-Net with roughly the same number of parameters. Multi-ResNed outper- form classical U-Nets by 29.8%, 12.8% and 3.4% on average over three random seeds, respectively. In Appendix B.I, we also show that U-Nets outperform FNO [52], another competitive architecture for PDE modelling, in this experimental setting.
32 x 32 (at M iler.)
16 x 16 (al 150 K iler.)
8 x 8 (at 100 K iiter))
4 4 (at 50 iter))
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;, here multi-level Discrete Wavelet Transform (DWT) using Haar wavelets, and identity operators for Eg. This 'DWT encoder' has no learnable parameters and comes at almost computational cost. We then add the parameters save into the decoder and achieve competitive and-on on certain mooblems-strictt, better performance when compared with a classical show for generative modelling with diffussin models in Appendix B. Multi-ResNNts are competitive with, yet inferior to Residual U-Nets, because the initial basis which V imposes is suboptimal and the encoder would benefit from learning better basis. This demonstrates the strength of our framework in understanding the role of the encoder and when parameterise it, which depends should explore how to choose V (and W) optimally, possibly eliminating the need of a parameterised encoder, and also carefully explore how optimally allocate and make of the (saved) parameters in Nestss [53, 54].
U-Net. However,
useful to our problem hand. It is now obvious that future empirical work

---

5.2 Staged training enables multi-resolution training and inference
Having characterised the self selilaimiy structure of U-Nets in $2, a natural idea is to explicitly train the U- on resolution i first, then train the U-Net on resolution while precon- ditioning on U continuing recursively. Optionally, we can freeze the weights of 4 upon training on resolution i - 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 ad- ditional advantages. First, it makes the U-Net modular with respect to data. When data higher-resolution is available, we can reuse our U-Net pretrained on lower-resolution in a principled way. Second, it enables to checkpoint the model during prototyping and experimenting with the model as we
3
8
A
Z
see low-resolution outputs of our U-Net early.
In Figure we illustrate samples from DDPM diffusion model [8] with a Residual U-Net trained with Algorithm 1. We use im- ages and noise targets on multiple resolutions (CIFAR10/MIITT:
Figure 8: U-Nets encode the geo- metric structure of data.
{4 x 4, 4,8 x 8, 16 x 16, 32 x 32}) 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 can utilise the data available on four different resolutions. It is also worth noting that training with Algorithm as opposed to single-stage training (as is standard practice) does not substantially harm performance of the hrgsusttresolution samples: (FID CIFAR10: staged training: 8.33 + 0.010; non-staged training: 7.858 + 0.250). We present results on MNIST, Navier-Stokes and Shallow water, with Multi-ResNets, and with strict version of Algorithm 1 where we freeze Ey and D after training on resolution in Appendix B.2.
5.3 U-Nets encoding topological structure
Algorithm Multi-rcsolution training and sam- pling via preconditioning.
In $3.3, showed how to design U-Nets with natural basis on a triangular domain, which encodes the topological structure of a problem into its archi- results
Require: Boolean FREEZE. for (I.....J) do if then
1:
tecture. Here, we provide proof-of-concey for this U-Net design. In Figure 8, we illustrate samples from a DDPM diffusion model [8] with a U-Net where we choose V and W as Haar wavelet subspaces of W = V = (A) (see $3.3), ResNet encoders and decoders and average pooling. The model was trained on - over 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 designing natural U-Net architectures on tessellations of complicated geometries such as spheres, manifolds, fractals, or CW-complexes.
2:
3:
Precondition on U--1-
4: Train
5:
FREEZE is True then
6:
Freeze E. and Di (fix parameters).
a custom version of MNIST with the digit and support
6 Conclusion
We provided a framework for designing and analysing U-Nets. Our work has several limitations: We put particular emphasis Hilbert spaces as the decoder spaces. We focus orthogonal wavelet bases, in particular of (X) or (A),, while other bases could be explored (c.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 extended with the use of Kronecker product. Little to
dimension only, but be no tuning is performed how to allocate the saved parameters in empirically demonstrate U-Nets triangles, while one could choose multitude of other topological structures. Lastly, future work should investigate optimal choices of 24 for domain-specine problems.
in $5.1. We design and
10

---

Acknowledgments and Disclosure of Funding
Christopher Williams acknowledges support from the Defence Science and Technology (DST) Group and from 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/7/19/2Z) and the Enrichment Scheme of The Alan Turing Institute under the EPSRC Grant EP/N5 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/RO13116 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 competing interests.
This research is supported by research compute from the Baskerville Tier 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 (FP/W032244/1) and operated by Advanced Research Computing at the University of Birmingham.
References
[1] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical Image Computing and Computer-assisted Intervention, pages 234-241. Springer, 2015.
[2] Nahian Siddique, Sidike Paheding, Colin P Elkin, and Vijay Devabhaktuni. U-net and its variants for medical image segmentation: A review theory and applications. Ieee Access, 9:82031-82057, 2021.
[3] Zongwei Zhou, Md Mahfuzur Rahman Siddiquee, Nima Tajbakhsh, and Jianming Liang. Unet++: A nested u-net architecture for medical image segmentation. In Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support: 4th International Workshop, DLMIA 2018, and 8th International Workshop, ML-CDS 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 20, 2018, Proceedings 4, pages 3-11. Springer, 2018
[4] Ozan Oktay, Jo Schlemper, Loic Le Folgoe, Matthew Lee, Mattias Heinrich, Kazunari Misawa, Kensaku Mori, Steven McDonagh, Nils Y Hammerla, Bernhard Kainz, et Attention u-net: Learning where to look for the pancreas. arXiv preprint 10XXv:/8044.3333 2018.
[5] Zoe Landgraf, Fabian Falck, Michael Bloesch, Stefan Leutenegger, and Andrew Davison. Comparing view-based and map-baser semantic labelling in real-time slam. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 6884-6890. IEEE, 2020.
[6] Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatfotemporll-sccll generalized pde modeling. arXiv preprint 2rXiv:2209.156166 2022.
[7] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 396-16611 2022.
[8] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[9] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[10] Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models, arXiv preprint 210.::2002009222.2 2021.
11

---

[11] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[12] Valentin De Bortoli, James Thomton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling, In Advances in Neural Information Processing Systems, volume 34, 2021.
[13] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-
276. 1998.
[14] Shun-ichi Amari, Jimmy Ba, Roger Grosse, Xuechen Li, Atsushi Nitanda, Taiji Suzuki, Denny Wu, and Ji Xu. When does preconditioning help or hurt generalization? arXiv preprint arXiv:2006. 10732, 2020.
[15] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[16] Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee whye Teh, and Arnaud Doucet. Riemannian score-based generative modeling. arXiv preprint arXiv:2202.02763. 2022.
[17] Fenqiang Zhao, Zhengwang Wu, Li Wang, Weili Lin, John H Gilmore, Shunren Xia, Dinggang Shen, and Gang Li. Spherical deformable u-net: Application to cortical surface parcellation and development prediction. IEEE transactions on medical imaging, 40(4):1217-1228, 2021.
[18] Hongyang Gao and Shuiwang Ji. Graph u-nets. In international conference on machine learning, pages 2083-2092. PMLR, 2019.
[19] Emilien Dupont, Yee Whye Teh, and Arnaud Doucet. Generative models as distributions of functions. arXiv preprint arXiv:2102.04770 2021.
[20] Fabian Falck, Christopher Williams, Dominic Danks, George Deligiannidis, Christopher Yau, Chris C Holmes, Arnaud Doucet, and Matthew Willetts. A multi-resolution framework for u-nets with applications to hierarchical vaes. volume 35, pages 15529-15544, 2022.
[21] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 255):359-366, 1989.
[22] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[23] Alfred Haar. Zur theorie der orthogonalen funktionenssysteme. Georg-August- Ialtuuuus--nnivrsttte Universitat, Gottingen., 1909.
[24] Jong Chul Ye, Yoseob Han, and Eunju Cha. Deep convolutional framelets: general deep learning framework for inverse problems. SIAM Journal Imaging Sciences, 1102)9911-0488 2018.
[25] Z. Ramzi, K. Michalewicz, JL. Starck, et al. Wavelets in the deep learning era. J Math Imaging Vis, 65:240-251, 2023.
[26] Liu, H. Zhang, Zhang, L. Lin, and W. Zuo. Multi-level wavelet-enn for image restoration. In Proceedings of the IEEE conference computer vision and pattern recognition workshops, pages 773-782, 2018
[27] William Reed and Thomas Hill. Triangular mesh methods for the neutron transport equation. Technical report, Los Alamos Scientific Lab., N. Mex.(USA), 1973.
[28] Kamrul Hasan Talukder and Koichi Harada. Haar wavelet based approach for image compression and quality assessment of compressed image. arXiv preprint 1rXiv:1010.4008 2010.
[29] Stéphane Mallat. A Wavelet Tour of Signal Processing. Elsevier, 1999.
12

---

[30] Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. arXiv preprint arXiv:2011.10650, 2020.
[31] Shmuel Agmon. Lectures on elliptic boundary value problems, volume 369. American Mathematical Soc., 2010.
[32] Susanne C Brenner, Ridgway Scott, and L Ridgway Scott. The mathematical theory of finite element methods, volume 3. Springer, 2008.
[33] Bernardo Cockburn, George E Karniadakis, and Chi-Wang Shu. The development of discontinu- ous Galerkin methods. Springer, 2000.
[34] Barry Boots, Kokichi Sugihara. Sung Nok Chiu, and Atsuyuki Okabe. Spatial tessellations: concepts and applications of voronoi diagrams. 2009.
[35] Boris Delaunay. Sur la sphère vide. Izv. Akad. Nauk SSSR, Otdelenie Matematicheskii i Estestvennyyk Nauk, 7(793-800):1-2, 1934.
[36] Der-Tsai Lee and Bruce Schachter. Two algorithms for constructing Delaunay triangulation. International Journal of Computer & Information Sciences, 9(3):219-442, 1980.
[37] Aryan Verma. Triangulations Using Matplotlib. masppoodit///ttpplloobbbtttaaare
[38] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011. 13456, 2020.
[39] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. In Advances in Neural Information Processing Systems, volume 34, 2021.
[40] Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational Diffusion Models. In Advances in Neural Information Processing Systems, volume 34, 2021.
[41] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint vrXiv:2205.11487 2022.
[42] Florentin Guth, Simon Coste, Valentin De Bortoli, and Stephane Mallat. Wavelet score-based generative modeling. arXiv preprint 0rXiv:2208.05003 2022.
[43] David L Donoho. De-noising by soft-thresholding IEEE Iransactions on information theory, 41(3):613 627, 1995.
[44] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[45] Thorvald A Sorensen. A method of establishing groups of equal amplitude in plant sociology based on similarity species content and application to analyses of the vegetation on danish commons. Biol. Skar., 5:1-34, 1948.
[46] Lee R Dice. Measures of the amount of ecologic association between species. Ecology, 26(3):297-302, 1945.
[47] Yann LeCun, Corinna Cortes, and C. J. Burges. MNIST handwritten digit database. http: //yann. lecun. com/ exdb/mmnst// 2010.
[48] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[49] Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatiotemporllacca generalized pde modeling, arXiv preprint 2rXiv:2209.16616 2022.
13

---

[50] Hugo J Kuijf, Matthijs Biesbroek, Jeroen De Bresser, Rutger Heinen, Simon Andermatt, Mariana Bento, Matt Berseth, Mikhail Belyaev, M Jorge Cardoso, Adria Casamitjana, al. Standardized assessment of automatic segmentation of white mattrr hyperintensities and results of the wmh segmentation challenge. IEEE transactions on medical imaging, 38(1)):5556-2668 2019.
[51] Hongwei Li, Gongfa Jiang, Jianguo Zhang, Ruixuan Wang, Zhaolei Wang, Wei-Shi Zheng, and Bjoern Menze. Fully convolutional network ensembles for white matter hyperintensities segmentation mr images. Neurolmage, 183:650-665, 2018.
[52] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differen- tial equations. arXiv preprint arXiv:2010.08880 2020.
[53] William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
[54] Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. arXiv preprint 2rXiv:22/2.11972 2022.
[55] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[56) Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jaejun Yoo. Reliable fidelity and diversity metrics for generative models. In International Conference on Machine Learning, pages 7176-7185. PMLR, 2020.
[57] Mehdi SM Sajjadi, Olivier Bachem, Mario Lucic, Olivier Bousquet, and Sylvain Gelly. As- sessing generative models via precision and recall. Advances in neural information processing systems, 31, 2018.
[58] Michael F Barnsley. Fractals everywhere. Academic press, 2014.
[59] Christoph Bandt, Nguyen Hung, and Hui Rao. On the open set condition for self-similar fractals. Proceedings of the american mathematical society, 134(5):1369-15 2006.
[60] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint 2rXi::2301.1199, 2023.
[61] Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Henrik Jacobsen. Invertible residual networks. In International Conference on Machine Learning, pages 573-582 PMLR, 2019.
[62] Özgün Çiçek, Ahmed Abdulkadir, Soeren S Lienkamp, Thomas Brox, and Olaf Ronneberger. 3d u-net: learning dense volumetric segmentation from sparse annotation. Medical Image Computing and Cossuter-Assisted intervention-MICC 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part 19, pages 424 432. Springer, 2016.
[63] Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein. nnu-net: a self-configuring method for deep learning-base biomedical image segmentation. Nature methods, 18(2)203-211, 2021.
[64] Fabian Isensee, Jens Petersen, Andre Klein, David Zimmerer, Paul F Jaeger, Simon Kohl, Jakob Wasserthal, Gregor Koehler, Tobias Norajitra, Sebastian Wirkert, et al. nnu-net: adlf-adapting framework for u-net-based medical image segmentation. arXiv preprint arXiv: 100..044 10486, 2018.
[65] Simon Kohl, Bernardino Romera- Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph R Ledsam, Klaus Maier-Hein, SM Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. A probabilistic u-net for segmentation of ambiguous images. Advances neural information processing systems, 31, 2018
14

---

[66] Patrick Esser, Ekaterina Sutter, and Björn Ommer. A variational u-net for conditional appearance and shape generation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 8857-8866, 2018
[67] Thorsten Falk, Dominic Mai, Robert Bensch, Özgün Çięek, Ahmed Abdulkadir, Yassine Marrakchi, Anton Böhm, Jan Deubner, Zoe Jäckel, Katharina Seiwald, et al. U-net: deep learning for cell counting. detection, and morphometry. Nature methods, 16(1):67-70, 2019.
[68] Zhengxin Zhang, Qingjie Liu, and Yunhong Wang. Road extraction by deep residual u-net. IEEE Geoscience and Remote Sensing Letters, 15(5):749-753, 2018.
[69] Nahian Siddique, Paheding Sidike, Colin Elkin, and Vijay Devabhaktuni. U-net and its variants for medical image segmentation: theory and applications. arXiv preprint arXhv:2011.01110 2020.
[70] Nabil lbtehaz and M Sohel Rahman. Multiresunet: multimodal biomedical image segmentation. Neural networks, 121:74-77, 2020.
Rethinking the u-ner architecture for
[71] Mert Lostar and Islem Rekik. Deep hypergraph u-net for brain graph embedding and classifica- tion. arXiv preprint arXiv:2008. 13118, 2020.
[72] Wenchong He, Zhe Jiang, Chengming Zhang, and Arpan Man Sainju. Curvanet: Geometric deep learning based on directional curvature for 3d shape analysis. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 2214 2224, 2020.
[73] Ingrid Daubechies. Ten lectures on wavelets. SIAM, 1992.
[74] Michele Benzi. Preconditioning techniques for large linear systems: survey, Journal of computational Physics, 182(2):4184477 2002.
[75] A. J. Wathen. Preconditioning. Acta Numerica, 24:329-376, 2015.
[76] Ke Chen. Matrix preconditionii techniques and applications, volume 19, Cambridge Univer- sity Press, 2005.
[77] Eli Turkel. Preconditioning techniques in computational fluid dynamies. Annual Review of Fluid Mechanics, 31(1):385 1999.
[78] Wolfgang Dahmen and Angela Kunoth. Multilevel preconditioning. Numerische Mathematik, 63(1):315-344, 1992.
[79] O Axclsson and Panayot S Vassilevski. Algebraic multilevel preconditioning methods. i. Numerische Mathematik, 56(2-3):157-177, 1989.
[80] Owe Axelsson and Panayot S Vassilevski. Algebraic multilevel preconditioning methods, ii
SIAM Journal on Numerical Analysis, 27(6):1569-1523 1990.
[81] Li Zhang, Hengyuan Ma, Xiatian Zhu, and Jianfeng Feng. Preconditioned score-based genera- tive models. arXiv preprint arXiv:2302.06504, 2023.
[82] Lukas Biewald. Experiment tracking with weights and biases, 2020. Software available from wandb.com.
[83] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, pages 8024-8033 8035, 2019.
[84] Fergal Cotter. Uses of complex wavelets deep convolutional neural networks. PhD thesis, University of Cambridge, 2020.
15

---

[85] Gregory Lee, Ralf Gommers, Filip Waselewski, Kai Wohlfahrt, and Aaron O' O'Leary. Pywavelets: python package for wavelet analysis. Journal of Open Source Software, 4(36):1237, 2019.
[86] William Falcon and The PyTorch Lightning team. PyTorch Lightning, March 2019.
[87] J. D. Hunter. Matplotlib: A 2D graphics environment. Computing in Science & Engineering, 9(3):90-95 900990995,20 2007.
[88] Charles R. Harris, K. Jarrod Millman, Stéfan J. der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, et al. Array programming with NumPy. Nature, 585(7825):357-32,, 2020.
[89]
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowiez, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large- scale machine learning heterogeneous systems, 2015. Software available from tensorflow.org.
[90] pyyaml contributors. pyyaml. https: stttg:///gttuu...om/m/mm//yyaaa 2006.
[91] tqdm contributors. Imageio. https://gggtuuu..c/////gg///////
2022.
[92] Pedregosa, G. Varoquaux, A. Gramfort, Michel, B. Thirion, Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830, 2011.
[93] Guido Van Rossum. The Python Library Reference, release 3.8.2. Python Software Foundation, 2020.
[94] OpenAl. Gpt-4 technical report, 2023.
16

---

Appendix for A Unified Framework for U-Net Design and Analysis
A Theoretical Details and Technical Proofs
A.1 Proofs theoretical results in the main text
Theorem Suppose U and U* are solutions of the L regression problem. Then, L (U;)² < (U*) with equality as + cc. Further, if Q,U* is V--mesuuralle, then U = Q,U minimises
C,
Proof. Let (V,,,, Un) be U-Net in canonical form, that is where each encoder map E; in is set to be the identity map, and further assume W Hilbert space with W being a sequence of subspaces spanned by orthogonal basis vectors for W. For a V-valued random variable v, define the o-algebra
Hj = (P, o(())
Now the filtration
H1 C H2 C...
(5)
increases to 1= = (v). These the o-algebras generated by uy. For data W, may define the losses
(U,))
Σ
I|s, (1)
(6)
S
for Uij Vj + W; where Q.w. Let Falj be the set of measurable functions from Vj to W,, then the conditional expectation E( B(ws|1y) is given by the solution of the regression problem
E us [b;)
arg min (U:g))
Uni
where the conditioning respect to the o-algebra Hj for the response variable w;. For a fixed í < and any j > 0 we have that
(UpD))))) = [S]
Σ
Iw - Ugr Ugr (1,))1
(8)
(w,,,,))E
(1Qiiww (11)(1;))11 +1Q U21((1;))11))
(9)
(wyr.uyes
= -0,,(QUUmb))
Σ
11 (wr - Umb((15))11.
(10)
S
(wp.,))
Any U2i admits the parameterisation Upi Usp + Upl where U is in Fals and U e Fel and vanishes W. Under this parameterisaiton gain
(Um))) = (U.b) + S
Σ (wie,,y)S
11 - (1,)11
(11)
and further,
arg min (Um))) arg min (Us(U EaI
arg min
Σ
|1@/ (w> (.)
[
U
(12)
In particular, when i' + if we define C to be the loss W then
arg min (U1)) arg min Ci, (U11s) arg arg min U (w.R))ES
Σ
11 (w - (w,))).
(13)
UF
UqyEFii
17

---

So for any j. we have that (,,(U,}))) - C,;(
Further, as i - oo, the truncation error
Σ
l1 (w U U
(14)
(1.ny)ES
tends to zero. Now assume that Q,UT is Vi-measurable then QU E Fig and if this did not minimise Lij- then choosing the minimiser of Lijj and constructing U = U + U will have (U) < (U*), contradicting U being optimal.
Proposition 1f a residual U-Net, then U; is ResNet precondtioooe on UP (u:) = U-- U--1(Us-1). where 0 (E))())
Proof. For a ResNet R(v;) = (Pr) + (2;) select RPT((:) = U--1(0;-1) where 0;-1 (B))
Theorem 2 For time t 2 0 and j 2 i, (t) _ X,()). Furthermore if (f( = EE (t) 3 be the decomposition of (() in its Haar wavelet frequencies (see Appendix C). Each component (l) of the vector has variance 2/-1 relative the variance of the base Haar wavelet frequency.
Proof. Let X; E Vú represented the standard basis = {or 1,...,22 giving
giving
(R) Ox.
X
(15)
To transform this into its Haar wavelet representation where average-pooling is conjugate to basis projection, we can use the map V, + V, defined by
T = H;
(16)
where H; the Haar-matrix on resolution i and A; is diagonal scaling matrix with entries
(A.) = 2 if E (2) ,,...,2'}, }
(17)
for j 0 and equal to 2 for the initial case
O. For example, the matrix for a four-pixel image is
0
0
2 2
0
1
0
-1
-1
AgH2.
(18)
0
0
21
0
0
0
0
0
2
-1
Now given the vector of coefficients X; (xx())) be given by
=== the coefficients in the Haar frequencies can
X, (X.))
(19)
We define the vectors
(8)s= )-42)1,,
(20)
for j (1....,,fe 21- Now we may write X; in terms of its various frequency components
X; Σ RG 3,
(21)
Suppose that X; is a time varying random variable on V42 that evolve according to
(t) = /1 x,x( (R) + (k)
(22)
18

---

when represented in its pixel basis where e(x) is standard normally distributed random variable. We may represent X; in its frequencies, with a random and time dependence on the coefficient vectors,
X; (0) 3,.
(23)
To analyse the variance of (I), we may analyse the variance of a standard normally distributed ran- dom vector E; = (s*) under the mapping T; Due to the symmetries in the matrix representation for T; the variance for each element of (t) is the same, and by direct computation
VAR (TiE:) = - for i > 0, k
(24)
and 2 for when i= = 0. To see this note that VAR (T) VAR(A, A
(25)
as A; is diagonal matrix. Now each of the elements of are independent and normally distributed, so if (H)) is the ktth row of H;, then (AR(H);;)) = VAR((H)).: = :;:
(26)
where ((HH)),:llo llo is the amount of non-zero elements of (H+),,: as each entry of H, is 0, or +1. For (Ha)k.: where k (2'-1 ..... 2 there are 2'-j+1 entries. Further, in this position 2-i+j-1, so putting this together finishes the count. Now that for the evolution of ( (t), the random part of Eq. (22) is simply a Gaussian random vector multiplied by /ot. Therefore
we have (A,)s 2
have that VAR VARX; x (t) ar2- To get the relative frequency, for k we divide this variance by the variance of the base frequency, the zeroth resolution. This yields a,2-it5-22 ------111/a0 1a, i = 2
Now we show how the diffusion given resolution i can be embedded into higher resolution of i + 1, which maintains consisteney with the projection map Q. Define the diffusion on + 1 to be
(A) (1) = v1 - or x + (20
(27)
Let x (t) be the Haar coefficients for the inital diffusion defined on resolution and 8 (t) be the coefficients defined on resolution i 1. Both of these random vectors are linear transformations of Gaussian random vectors, so we need only confirm that the means and variances on the first
entries agree to show that these are equal in distribution. If x(()
( x +x(41 that is the initial 2
data on resolution + 1 averages to the initial data on resolution i, then by construction the means of the random vectors X (t) and x+1 (t) agree for the first i resolutions. Note that the variances for the components of the noise for either process are 2 g-itj-1 when the diffusion on + 1 given by Eq. (27). In general, for the process on i + é to embed into our original diffusion on resolution i we require scaling the noise term by a factor of 24/2, Again, this shows that the noising process on a frequency i has the natural extension noising the high-frequency components exponentially faster than the reference frequency i.
We may comment on how this is represented in its natural sequence space. Assume that X; E Vi C ([, i)). There is a natural mapping from V; to fz defined by
Ey = {s]
sk
co)
(28)
k=
through the mapping æ I., 11) + Eg via
x(X.) (8
(29)
For elements X; E V; the image of π is not surjective and has image contained in t0o the set of infinite sequences which eventually are zero. In this case we have a finite dimensional mapping. Theorem 2 simply states that the forward noising process in a diffusion model, under transformation into its Haar basis, has variance
VAR (TX,)s 2 --i--1 for > 0, E
2'),
(30)
implying that the sequence image of the datum has monotonically increasing variance in the Haar wavelet sequence space.
19

---

A.2 Diffusions on the sphere.
Suppose that we have L (XX valued data on the globe and we would like form diffusion model on this domain. Instead of constructing local change of coordinate system over manifold and projecting our diffusion geometry. As an example of this, we can take standard triangulation of the globe, then apply our triangular basis U-Net directly to our data.
square domain 16], we could simply design our U-Net to encode our
For the data on the globe, we could then choose a resolution i to refine our triangular domain to, see Figure 9 for a refinement level of i 2 which could be used tessellate the globe. For each 'triangle pixel' in our refinement, we could define the diffusion process
= ax Vāe!
(31)
where is a string of length i+1 encoding which triangle on the globe the pixel is in, then which "triangular pixel' we are We can construct Haar wavelets on the triangle (see Figure 3.3), and hence create natural projection maps in L over this domain. In this case Wo is the span of constant functions over each large triangle for the tessellation of the globe, and each W2 is Wo along with the span of the wavelets of resolution over each triangle. Like in the standard diffusion case, we choose V; W; and can use the mapping given in Section 5.2 construct the natural U-Net design for this geometry, and for this basis. Similarly, for any space that we can form triangulation for, we can use this construction create the natural Haar wavelet space for this geometry, define a diffusion analogously, and implement the natural U-Net design for learning.
33
34
31
32
42
41
13
44
23
14
43
11
12
21
22
Figure 9: Triangulation of the globe with self-similar triangle which admits Haar wavelet basis. Upon refinement through the self-similarity of the triangle we receive finer and finer approximations of the globe, and of functions over it.
A.3 U-Nets on Finite Elements
PDE surrogate modelling is nascent research direction in machine learning research. We have seen in our PDE experiment (see Section 5) that our Multi-ResNc design outperforms Residual U-Nets on square domain. We would like to mimic this in more general geometries and function spaces. Here, we will show on the unit interval how to design a U-Net over basis functions which enforce boundary constraints and smoothness assumptions, mimicking the design of Finite Elements used for PDE solvers.
Recall our model problem
Au =
n(0) u(1)= =0
(32)
which when multiplied by test function o E ((0, 1)) and integrated over the domain puts the equation into its weak form, and by integration by parts
-0', (()))) = f of.
(33)
we recover the standard bilinear representation of Laplaces equation. Now if our solution HG] 1)) and we have an orthogonal basis (Φt..) } for Hb,(0, 1)), then the coefficients for the basis elements of the solution be solved by finding the solution of linear system with diagonal
20

---

mass matrix. Thus, in this way, the information needed to solve the linear system up to some finite resolution only depends on data up that resolution in this weak form. This is Galerkin truncation of the PDE used for forward solvers on finite resolutions, but also gives us the ideal assumptions we need for the construction of our U-Net (the measurability constraint Theorem 1).
For an explicit example, we may make an orthogonal basis for ( 1]) with the initial basis function [Left] of Figure 4 and its refined children on the [Right] of Figure
general we can refine further and create basis resolution through
(x) = x + j/2*),
x(x) = 2x (0.11 + (2 2n) 11/22.1(x);
which are precisely the integrals of the Haar wavelets up to given resolution. We would again construct W; in our U-Net as the span of the first resolution i basis functions, where we now have additionally encoded the boundary constraints and smoothness requirements of our PDE solution into our U-Net design.
21

---

B Additional experimental details and results
In this section, we provide further details on our experiments, and additional experimental results.
Model and training details. We used slightly varying U-Net architectures on the different tasks. While refer to our code base for details, they generally feature single residual block per resolution for Ej and D;, with 2 convolutional layers, and group normalisation. We in general choose P. as average pooling. We further use loss functions and other architectural components which are standard for the respective tasks. This outlined U-Net architecture often required us to make several changes the original repositories which we used. This may also explain small differences to the results that these repositories reported; however, our results are comparable. When using Algorithm 1 during staged training, we require one 'head' and 'tail' network on each resolution which processes the input and output respectively, and hence increases the number of parameters of the overall model relative single stage
single-stage training.
Hyperparameters and hyperparameter tuning. We performed little to no hyperparameter tuning in our experiments. In particular, we did not perform search (e.g. grid search) over hyperparameters. In general, we used the hyperparameters of the original repositories as stated in Appendix D, and changed them only when necessary, for instance adjust the number of parameters so enable fair comparison. There is hence a lot of potential to improve the performance of our experiments, for instance of the Multi-ResNet. or the U-Net on triangular data. We refer to our code repository for specific hyperparameter choices and further details, in particular the respective hyperparan py files.
Evaluation and metrics. We use the following three performance metrics our experimental evaluation: FID score [44], rollout weean-squrred-eroor (r-MSE) [6], and Sørensen- Dice coefficient (Dice) [45, 46]. As these are standard metrics, we only highlight key points that are worth noting. We compute the FID score on holdout dataset not used during training, and using evaluation model where weights are updated with the training weights using an exponential moving average (as is common practice). The r-MSE is computed as MSE over pieces of the PDE trajectory against its ground-truth. Each piece is predicted in an autoregressive fashion, where the model receives the previous predicted piece and historic observations as input [6]. All evaluation metrics are in general computed on the test set and averaged over three random seeds after the same number of iterations in each table, if not stated otherwise. The results are furthermore not depending on these reported evaluation metrics: in our code base, we compute and log large range of other evaluation metrics, and we typically observe similar trends as those reported.
22

---

B.1 Analysis 1: The role of the encoder in a U-Net
In this section, we provide further experimental results analysing the role of the encoder, and hence Residual U-Nets with Multi ResNess. We realise the Multi-ResNet by replacing the parameterised encoder with a multi-level Discrete Wavelet Transform (DWT). More specifically. the skip connection on resolution, relative to the input resolution is computed by taking the lower- part of the DWT level j (I.L,), and then inverting I.I.j to receive coarse, projected version of the original input image.
Generative modelling with difflusion models. In Table 2, analyse the role of the encoder in generative modelling with diffusion models. Following our analysis on PDE modelling and image segmentation, we compare (Haar wavelet) Residual U-Nets with (Haar wavelet) Multi-ResNets where we add the saved parameters in the encoder back into the decoder. We report performance in terms of an exponential moving average FID score [55] on the test set. In contrast to the other two tasks, we find that diffusion models benefit from a parameterised encoder. In particular, the Nusti-ResNet does not outperform the Residual U-Net with approximately the same number of parameters, as the case in the other two tasks. Referring to our theoretical results in $3.1, the input space of Haar wavelets may be suboptimal choice for diffusion models, requiring an encoder learning a suitable change of basis. A second possibility is the FID score itself, which is a useful, yet flawed metric to quantify the fidelity and diversity of images [56, 57]. To underline this point, Figure 10 we compare samples from two runs with the Residual U-Net and Multi- observe that very difficult to tell "by eye" which model produces higher quality samples, in spite of the differences in FID. A third possibility is that the allocation of the saved parameters is suboptimal. We demonstrate the importance of beneficially allocating the saved parameters in the decoder in the context of PDE modelling below.
reported in Table 2.
Table Quantitative performance of the (Haar wavelet) Hulti-ResNet compared a classical (Haar wavelet) Residual U-Net in generative modelling with diffusion models on CIFAR10. We report FID on the test set.
Dataset Neural architecture
# Params.
FID +
CIFARIO
32
Residual U-Net
35.5 M
7.86 +: 0.25
params. added in dec. (ours)
25.8 M
14.87 + 0,50
Multi-ResNeet, saved params, added in dec. (ours)
32.4
12.44 0.22
Residual U-Net (35.5 Parameters)
Multi-ResNel (32.4 Parameters)
Figure 10: Samples of a DDPM-type diffusion model with Residual U-Net [Left] and Multi- ResNet [Right]. We trained both models for 1.2 M iterations at which we obtained the samples.
23

---

Lastly, we present Figure from the main text at larger resolution in Figure 11.
CIFAR10
32 x 32 (at 1 M iter.)
16 x 16 (at 150 K iter.)
8 (at 100 K iter.)
4 x 4 (at 50 K iter.)
Figure 11: Preconditioning enables multi-resolution training and sampling diffusion models. We train a diffusion model [8] with Residual U-Net architecture using Algorithm I with four training stages and no freezing on CIFAR10. This Figure identical with Figure 7, but larger.
PDE modelling. We further present additional quantitative results comparing the (Haar wavelet) Multi-ResNed with (Haar wavelet) Residual U-Nets on our PDE modelling tasks. These shall highlight the importance of allocating the saved parameters in Multi-ResNett to achieve competitive performance. In Table 3, we compare two versions of allocating the parameters we save in Multi- ResNets, which have no parameters in their encoder, on the Navier-stokes and Shallow water PDE modelling datasets, respectively. In 'vl', also presented in the main text in Table but repeated for case comparison, we add the saved parameters by increasing the number residual blocks per resolution. This in effect corresponds to more discretisation steps of the underlying ODE on each resolution [20]. In 'v2', we add the saved parameters by increasing the number hidden channels in the ResNet.
Our results show that 'vl' performs significantly better than 'v2'. In particular, 'v2' does not outperform the Residual U-Net on the two PDE modelling tasks. This indicates that the design choice of how to allocate parameters D matters for the performance Multi-ResNets. We note that in our experiments, beyond this comparison, we did not explore optimal ways of allocating parameters Nesti.ResNets.. Hence, future work should explore this aspect thoroughly with large- scale experiments. Due to its tuni-directionll, 'asymmetric' structure, the Melti-ResNett may further

---

be combined with direction of transformer- based architectures which aim at outperforming and replacing the U-Net [53, 54]. With reference to our earlier presented, inferior results with diffusion models, it is possible that merely designing better decoder would make Multi-ResNet, superior to classical Residual U-Nets for diffusion models.
Table 3: Quantitative performance of the (Haar wavelet) Multi-ResNet compared to classical (Haar wavelet) Residual U-Net two PDE modelling. This table is augmented version of Table I in the main text, where 'v1' indicates the run from the main text, a Multi-ResNer with saved parameters added in the encoder, and 'v2' indicates an alternative parameter allocation. We report Mean- Mean-Squared- Squared-Erre over a rolled out trajectory in time (r-MSE) on the test set, rounded to four decimal digits.
Dataset
Architecture
Params.
r-MSE +
Residual U-Net
34.5 M
0.0057 + 2 10 -5
Navier-
stokes
128 x28x122 128
Multi-ResNee. params. added in dec. (ours)
15.7 M
0.0107 9 90:9-100 10
Multi-ResNet,, saved params. added in dec. v1 (ours)
34.5 M
0.0040 + 2 10 10
Multi-ResNett, saved params. added in dec. v2 (ours)
34.6 M
0.0093 + 10
Shallow water 192
Residual U-Net
34.5 M
0.1712 0.0005
Multi-ResNet, params. added in dec. (ours)
15.7 M
0.4899 0.0156
Multi- ResNet, saved params, added in dec. v1 (ours)
34.5 M
0.1493 - 0.0070
Multi-ResNet saved params. added in dec. v2 (ours)
34.6 M
0.3811 = 0.0091
In Figure 12 we illustrate the full prediction of our Haar wavelet Multi-ResNet when modelling PDE trajectory simulated from the Navier-stokes and Shallow water equations, respectively, corresponding to its truncated version in Figure 6 [Left]. We unroll the trajectories over five timesteps for which we predict the current state. Note that train the Multi-ResNet by predicting the next time step in the trajectory only. We do not condition on previous timesteps.
ScU
Ground truth
S
Prediction
Error
tredeetto
TDY
Figure 12: PDE modelling and image segmentation with a Wavelet-encode out PDE trajectories (ground-trouth prediction, Te--error between ground-truth and prediction) from the Navier-ttokes [top], and the Shallow-Wate equation [bottom]. This is the complete version of the truncated Figure [Left] showing further timesteps for the trajectory. Figure and code as modified from [6, Figure 1].
Rolled
25

---

We also present results comparing our Multi-ResNet, Residual U-Nets and the Fourier Neural Operator (FNO) [52], another competitive model for PDE modelling, in Table We compare models of similar size in terms of parameter count. We also present the same comparison in PDEArena [6] in Table 5 for reference, noting that our experimental configuration slightly differs. Both tables indicate that U-Nets outperform FNO.
Table 4: Quantitative performance of FNO model compared to U-Nets. We compare the FNO model to a (Haar wavelet) Multi-Rest Ness and classical (Haar wavelet) Residual U-Nets with similar number parameters, on two PDE modelling tasks.
Dataset Neural architecture
# Params.
r-MSE +
Shallow Navier- water
stokes
128
Residual U-Net
34.5 M
0.0057 + 2 10-5
x
Nesti- saved params. added in dec.
34.5 M
0.0040 + 2- 10
128
FNO 128-8 mode8 (new)
33.7 M
0.0253 + 0.0
192
Residual U-Net
34.5 M
0.1712 1 0.0005
Multi-ResNer saved params. added in dec.
34.5 M
0.1493 + 0.0070
96
FNO 128-8 mode8 (new)
33.7 M
1.2333 = 0.0115
Table 5: Experimental results as reported in PDEArena [6]): Quantitative performance of an FNO model, in comparison to U-Nets of similar size. Values reported in Table 8 in [7] (5200 trajectories) for Navier-Stokes, and in [Table 2 in [7] (5600 trajectories)] for Shallow water. Here, U-Net 2015 64 clearly outperforms FNO 128-8 mode8. This result is consistent across different numbers of trajectories (dataset sizes), and different data configurations (e.g. velocity function formulation vs. vorticity stream function formulation Shallow water) in the several tables reported in [7].
Dataset
Neural architecture
# Params.
r-MSE +
Navier-
U-Net 2015 64 FNO 128-8 8 mode8
31 M
0.01386 + 0.00004 0.03836 + 0.00037
stokes
33.7 M
128 128
Shallow
U-Net 2015 64
31 M
0.1026 + 0.0161
water
FNO 128-8 mode8
33.7 M
0.8549 + 0.0124
26

---

Image segmentation. In Figure 13, we present further MRI images with overlayed ground-trute (green) and prediction (blue) masks, augmenting Figure [Right] in the main text. The predictions are obtained from our best-performing Multi-ResNet, Note that some MRI images do not contain any ground-truth lesions.
X
Figure 13: MRI images from WMH with overlayed ground truth masks (green) and predictions (blue) obtained from our best-performing Multi-Res.N model.
27

---

B.2 Analysis 2: Staged training enables multi-resolution training and inference
We begin by providing further experimental details. For each dataset, we train each res- olution for the following number of iterations/epochs: MNIST [5K, 5K, 5K, 5K] iterations, CIFAR10 [50K, 50K, 50K, 450K] iterations, Seovirr-Stokee [5, 35] epochs, Shallow water 2,, 14] epochs.
Generative modelling with diffusion models. In Figs. 14 and 15, we illustrate samples of diffusion model (DDPM) trained on MNIST with Algorithm 1, with and without freezing, respectively. These samples show that in both cases, the model produces reasonable samples. Algorithm I with freezing has the advantage that the final model can produce samples on all resolutions, instead of only the intermediate models the end of each training stage. Preliminary tests showed that as perhaps expected, freezing lower-resolution U-Net weights tends to produce worse performance in terms of quantitative metries measured on the highest resolution. In particular, in Figure 16 we compare samples and FID scores from DDPM model trained on CIFAR10, with and without freezing. As can be clearly seen from both the FID score and the quality of the samples, training with Algorithm but without the freezing option produces significantly better samples. This is why we did not further explore the freezing option.
MNIST
8
5
S
/
3
3
9
7
K
5
C
7
9
7
/
/
á
L
[
6
3
6
6
3
8
7
7"
/
9
O
{
9
32 (at 20 iter.)
16 (at 15 iter.)
5
!
5
F
F
7 4.
8 x 8 (at 10 K iter.)
4 4 (at iter.)
Figure 14: Staged training of diffusion models. We train diffiusion model [8] with Residual U-Net architecture using Algorithm I with four training stages and with freezing on MNIST corresponding to Figure 7, We show samples at the end of each training stage.
enables multi-resolution training and sampling of
PDE modelling. In Table 6, we investigate Residual U-Nets and Multi-ResNets using Algorithm I on Sevirr-Stokes and Shallow water. We here find that staged training makes the runs perform substantially worse. In particular, the standard deviation is higher, and some runs are outliers with particularly poor performance. This is in contrast our experiments on generative modelling with diffusion models where staged training had only an insignificant impact on the resulting performance.
We note that we have not further investigated this question, and believe that there are a multitude of ways which could be attempted to stabilise and improve the performance of staged training, for instance in PDE modelling, which we leave for future work.
28

---

MNIST
3
4
5
4
4
8
2
8
3
0
9
5
7 2
8
7
9
7
2
7
7
R
3
2
3
0
7
2
3
4
9
0
32 x 32 (at 20 K iter.)
16 x 16 (at 15 K iter.)
+
4
/'
5
F
x 8 (at 10 K iter.)
x 4 (at 5 K iter.)
Figure 15: Staged training of Multi-ResNets enables multi-resolution training and sampling of diffusion models. train a diffusion model [8] with a Residual U-Net architecture using Algorithm I with four training stages and freezing on MNIST corresponding to Figure 7. We show samples at the end of each training stage.
Algonthim freezing
Alporthmm freesing
Figure 16: A comparison of samples of DDPM model [8] with Residual U-Net trained with Algorithm with [Right] and without [Left] freezing.
Table Staged training with Algorithm on Navier- and Shallow water. Quantitative performance of the (Haar wavelet) Multi-ResNet compared to a classical (Haar wavelet) Residual U-Net.
Dataset
Neural architecture
Params.
r-MSE 1
Navier-stokes 128 x 128
Residual U-Net
37.8 M
0.0083 + 0.0034 0.0105 + 0.0102
Musti-ResNet, saved params. added in dec. (ours)
37.8 M
Shallow water 96 192
Residual U-Net
37.7 M
0.4640 1 0.4545
Nesti-ResNet, saved params. added in dec. (ours)
37.7 M
0.7715 1 0.8327
29

---

B.3 Analysis 3: U-Nets encoding topologiaal structure
Experimental details. We begin by discussing MNIST- Trin-TTranggaa, the dataset we used this experiment. In nutshell, -NIST-Triangull dimensional squared digit into the lower-left half of a 64 x 64 squared image with similar background colour in MNIST in that lower-left half of triangular shape. In the upper-right half, gray background colour to indicate that the image is supported only on the lower-left of the squared image. We illustrate IWIST-Triangen than four times the number of pixels compared to MNIST, yet we trained our DDPM U-Net with hyperparameters and architecture for MNIST without hyperparameter tuning, explaining their perhaps slightly inferior sample quality in Figure 8.
is constructed by shifting the MNIST 28 x 28
in Figure 17. We note that MNIST- TISSTTTriange
has more
4
2
3
7
9
9
3
6
6
7
3
9
8
5
3
0
7
4
g
O
4
4
6
Figure 17: Example images from the MMIST-Triangular dataset.
To process the data into the subspaces W (W)) given in Section 3.3, need formulate the 'triangular pixels' shown Figure 3.3. To do this we use the self-similarity property of the triangular domain to divide the data into desired resolution. For each image, then sample the function value at the center of our triangular pixel and store this in a lexicographical ordering corresponding to the codespace address [58] (e.g. '21") of that pixel. In Figure 18 illustrate the coding map depth two.
Once we have made this encoding map, able to map our function values into array and perform the projections on by push-forwarding through projection Haar wavlets on the triangle to our array through the lexicographic ordering used. For instance, we show in Figure 19
22
24
encoding map
11
12
21
22
21
23
13
14
23
24
31
32
41
42
42
41
33
34
43
44
12
44
32
14
43
34
11
13
31
33
Figure 18: The coding map from the triangular Haar wavelets to their code-space addresses. Such construction can always be made on self-similar object with certain separation properties, such as the Open Set Condition [59].
30

---

Figure 19:
An example of the encoding map taking data from a triangular domain and mapping it to array of each 'triangular pixel' value under lexicographical ordering.
the array, plotted as image, revealing that this processing is inherently different to diffusions square domain.
This transformation, when defined on the infinite resolution object, is discontinuous on a dense set of A, yet encodes all of the necessary data to perform diffusion model over in our triangular MNIST model seen in Figure 8. This is because the transformation we are using is still continuous almost everywhere, and so the parts of the signal that we are losing accounts for negligible amount of the function approximation in the diffusion process.
31

---

B.4
Ablation studies
In several ablation studies, we further analysed Musti-ResNees and our experimental results. We present these below.
B.4.1 Ablation 1: The importance of skip connections in U-Nets
We begin by presenting a key result. As Multi-ResNNes perform a linear transformation in the encoder, which significantly reduces its expressivity, one might hypothesise that the skip connections in U-Net can be removed entirely. Our results show that skip connections are erucial in U-Nets, particularly in Multi-ResNets.
compare Residual U-Nets and Multi-ResNetss with and without skip connections fo- cussing on PDE Modelling Slavier-stoker and Shallow water). The "without skip connections' case is practically realised by feeding zero tensors along the skip connections. This has the benefit of requiring no change to the architecture enabling a fair comparison, while feeding no information via the skip connections as they had been removed. We find that performance significantly deteriorates when using no skip connections. This effect is particularly strong in Kulti-ResNets cannot compensate the lack of skip connections by learning the encoder in such a way that represen- tations learned on the lowest resolution space Vo account for not having access to information via the skip connections when approximating the output via D; in the output space W; on a higher resolution space. This result shows that the encoder actually compresses the information which it provides to the decoder. The decoder crucially depends on this compressed input.
In Table 7
Multi-ResNees
Table 7: Skip connections in U-Nets crucial. Residual U-Nets and (Haar wavelet) Melti-ResNets, with and without skip connections, focussing on the two PDE modelling datasets (Navier-Stoees and Shallow water).
We compare two
(Haar wavelet)
Neural architecture
# Params.
r-MSE |
Residual U-Net w/ skip con.
34.5 M
0.0057 2- 10-5.
Navier- stokes 128
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
Multi-ResNet w/o skip con.
34.5 M
0.0831 = 0.1080
Residual U-Net w/ skip con.
34.5 M
0.1712 + 0.0005
Shallow
122
water
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 V and 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 a single subspace from the Residual U-Net by replacing P; 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-Stann Stokes for approximately 400 K iterations evaluating on the validation set. This is due to significantly longer running times of the ResNets constructed 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 show that preconditioning via the U-Net's self-similarity structure a key reason for their empirical success. The Residual U-Net multiple subspaces outperforms the ResNet
32

---

on 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 plain ResNet. Quantitative performance of the (Haar Wavelet) Residual U-Net over multiple input and output subspaces compared ResNet on single subspace, both trained on Navier -rStoees and Shallow water.
Dataset
Neural architecture
# Params.
r-MEE
Navier-stokes
Residual U-Net (multiple subspaces)
34,5 M
0.0086 + 9.7 10
128 128
ResNet (one subspace)
34.5 M
0.3192 + 0.0731
Shallow water 96 192
Residual U-Net (multiple subspaces)
34.5 M
0.3454 + 0.02402
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 we analyse the effect of different wavelets for the DWT we compute in P, We selected examples from different wavelet families, all of them being orthogonal. A good overview of different wavelets which can forriglt- used compute DWT can be found on this pybytes. com/ On Navier-Stokes 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, 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.
URL:
analyse the effect of choosing
Table 9: On the importance of choosing different wavelets for the Discrete Wavelet Transforms (DWT) Multi-ResNets. We report quantitative performance on the test set of Navier-Stokes and Shallow water.
Dataset
Neural architecture
# Params.
r-MSE +
Multi-ResNet w/ Haar wavelet DWT
34.5 M
0.0040 + 2 10
Navier-stokes 128 x 128
Multi-ResNet w/ Daubechies DWT
34.5
0.0041 + 10-5
Multi-ResNet w/ Daubechies 10 DWT
34,5 M
0.0068 + 1.1 10
Multi ResNet w/ Haar wavelet DWT
34,5 M
0.1493 + 0.0070
Multi-ResNet w/ Daubechies DWT
34,5 M
0.1081
Shallow water 96 192
w/ Dauhechies 10 DWT
34.5
0.1472
Multi-ResNee Coiflets DWT
34.5
0.1305
w/ Discrete Meyer DWT
34.5
0.1237
33

---

B.5 On the importance of preconditioning in residual learning: Synthetic experiment
We begin by providing second example with w(v) 2 in Figure corresponding to Figure in $1. Note that in comparison to the problem w(u) = 2² in Figure 2, the effect of the two precon- ditioners swaps in Figure 20: Using Ro*((o) approximation of the ground-truth function, while choosing R (e)e(v) = |v| results in deteriorate performance. This is because (v)
as the preconditioner produces an almost perfect
a 'good' initial guess for w(()) v², but a very 'poor' guess for w(x) v², and vice versa for Rere(() |u|- This illustrates that the choice of using good preconditioner is hand.
and can be crucial solve an optimisation problem at
w(v] v² (ground-truth)
wivi v² (ground-truth)
R(v (oporoximation)
Rlvi Capproximationn
4P*[v (preconditioner)
(e*[u) (preconditione))
Figure 20: The importance of preconditioning. We learn the ground-trui truth functions w(e) { using a ResNet (v) = R (v) + R(v) with preconditioners [Left] R = and [Right]
Rere |m]-
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 (v) R (v) R(v) over a grid of values {(,,, w;) where v; E [-,,
w; = (1()),,)) on the experiment scenario indicated by the superscripts (k, L), respectively. We do not worry about generalisation in this example and use the same train and test distribution. This functional relationship easy learn by any neural network R of sufficient size, which would overshadow the effect of preconditioning. To construct ResNet which is 'just expressive enough" to learn (kk), we constrain the expressivity of two ways: First, we choose R to small fully-connected network with 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 , of R such that ||o,| < for all where || || is the Frobenius norm, and repeatedly apply R with D = 100 practice, ||0]|| will be close to 1, and as ||0, intuitively measures
50, and w (v) = v² w(
v3, R. (v) = v, (v) = |u] depending
times via weight-sharing. the 'step size" of neural network, constrain the neural network $100 steps of unit length'. As an alternative, one could limit the maximum number training steps, which would indirectly constrain the model's 's expressiveness. We note that reported results robust across variety of hyperparameter combinations leading to qualitatively the same conclusions.
34

---

C Background
C.1 Related work
U-Nets. The U-Net [1] is a go-to, totte-of-thearrr 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]. U-Net which automatises parts of its design, 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 key improvement for the seminal U-Net [1] is to use residual blocks instead of feed-for wredrro 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 unified definition of U-Nets which goes beyond and encompasses the definition in [20], and highlights the key importance of preconditioning in U-Nets well their self self-similarrty structure. Importantly, this definition is not limited to orthogonal wavelet spaces as choices for V and W, hence enabling the use of U-Nets for 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 U-Net. This 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, novel contribution. Fourth, our experiments with Multi-ResNed 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 wide set of applications across various domains.
Miscellaneous. We further briefly discuss various unrelated works which similar concepts as those in this paper. Preconditioning, initialising 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 diffusion in lower-dimensional latent space and experience performance gains [11, 60). Stable Diffusion in particular found that their model learns dineeiiinnioonal latent space in which nigh-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].
35

---

C.2
Hilbert spaces
Hilbert space is vector space W endowed with an inner- inner-product (..)) that also induces complete norm ||||| via ||w|| = (w, w). Due to the inner- amerrprroduc, we have notion of two vectors W1, W2 E W being orthogonal if (wi, w2) 0.
In Section 3 we have paid special attention to two specific Hilbert spaces: ( 1J) and ( 1))-
1. the space (0, 1]) has as elements square-integrable functions and has an inner-product given by
(.9) U. S[ f(z)g(2)d;; and,
(34)
2. the space ([0, 1I) has as elements once weakly differentiable functions which vanish at both zero and one, with inner-product given by
S(' (((r)g(())d
V (. Jo
(35)
where I' E (2( 1J) are the weak derivatives of f and g respectively,
If we have a sequence {dw] of elements of W which are pairwise orthogonal and span W, then we call this orthogonal basis for W. Examples of orthogonal bases for both [² 1]) and H4([0, 1]) are given in Section C.3.
C.3 Introduction to Wavelets
Wavelets are refinable basis functions for ([0, 1]) which obey a self-similarity property in their construction. There is 'mother' and 'father' wavelet Φ and e) of which the children wavelets are derivative of the mother wavelet Φ. For instance, in the case of Haar wavelets with domain [0, we have that (x) and the mother and children wavelets given by
(x,((x) 6(2hx j/2*),
<xx) (2)) (z)
Under the -nnin product, the family (or. Further, we define the functions
2 1 -= forms an orthogonal basis of L ( 1]).
(x)( 1-- (v)
(36)
then each of these functions is in Hd. 1]), and is further orthogonal with respect the H) Htl-inner-product, bootstrapped from the orthogonality the Haar wavelets.
The discrete encoding of the Haar basis can be given by the kronecker product
H, = ) H
Ho ( -1.
(37)
To move between the pixel basis and the Haar basis natural to average pooling, we use the mapping defined on resolution by T, V, + Vú2 through
T((x)) = A,Hoo
(38)
where H; is the Haar matrix above and A, is a diagonal scaling matrix with entries
(A+)a,==3 = 1- if E (2i-1.....
(39)
To get this, we simply identify how represent 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 image with four pixels we can describe the function the weighted sum of the four basis functions [Right], that take values +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 chosen be the local deviance's 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].
36

---

Figure 21: Modelling ID image with four pixel values the weighted sum of Haar wavelet frequencies. The coefficients are such that the local averages of pixel values give the Haar wavelet at 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 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 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 Eg and D; be operators on functions.
f X
y
.3
R
original (512 512)
gcayed original (512 512)
compact set X Im
C1
C2
C3
C4
C5
C6
C7
Figure 22: Images modelled as functions.
D Code, computational resources, datasets, existing assets used
Code. We provide our code base as well instructions to reproduce the key results in this paper under MIT License at code from four Github repositories which are subfolders. For our diffusion experiments on MNIST and MNIST-THee -risS-Trianullar we directly build on top of the repository https://github httpJ//99//tocchhhee
Our code base uses
For our diffusion experiments on CIFAR10, we directly build on top of the repository :ttps://get thub. com/w86763777 https////ptttccc-hheo ddpm. For our PDE experiments on Havier-Stoken and Shallow water, we use the repository https: //github. com/m crosoft/ pdearena [6]. For our image semgentation experiments WMH, are inspired by the repository
37

---

https://githut hwmp/hongweilibam//mm... ibbmTum [51], but write the majority of code our- selves. We note that the MIT License only pertains to the original code in our repository, and 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, well miscellancous adjustments for conveniently running code, for instance with different hyperparameters from the command line.
Datasets. In experiments, we use the following five datasets: MNIST [47], MNIST-Triangul.s a custom version of MNIST where data is supported on a triangle rather than square, CIFAR10 [48], Navier-stokee Shallow water [49], and the MICCAI 2017 White Matter Hyperinten- sity (WMH) segmentation challenge dataset [50, 51]. These five datasets with the exception of MNIST-Tr-h MMIST-Triangular: have in common that data is presented over square 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 MOIST-Triangulaar our custom implementation and dataset class as part of our code hase, referring to Appendix B.3 on how it constructed. It is also worth noting that favier-Stoked and Shallow Water require considerable storage. On our hardware, the two datasets take approximately 120 GB and 150 GB unzipped, respectively.
provide
Computational resources. This work made use two computational resources. First, we used two local machines with latest CPU hardware and one with 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
shared with large number of other users.
To reproduce single run (without error bars) any of the experiments, we provide approximate run times for each dataset using the GPU resources: On MNIST and MMIST-Triangula a single run takes about 30 mins. On CIFAR10, a single run takes several days. On Navier-Stoker and Shallow water, single run takes approximately 1.5 days and day, respectively.
Existing assets used. Our code base uses the following main existing assets: Weights& [82] (MIT License), PyTorch [83] (custom license), in particular the torchvision package, pytorch_waveles [84] (MIT License). PyWavelets [85] (MIT License), pyytorch_lightninn [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] (MPL.v2.0 MIT License), scikit-leara and sklearn [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 writing aid.
38


Reviews:
[{"review_id": "l5DeYH7XvT", "reviewer": "Reviewer_7MgH", "summary": "The work proposes a unified mathematical framework to analyze and design U-Net. Authors highlight the importance of preconditioning and provides several highlights to design unet. Authors propose a new parameter-free encoder based on wavelet space.", "questions": "See weakness", "limitations": "See weakness", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Sec 2 presents a unified mathematical framework for U-Unets. \n2. The discussion of precondition parts and how to design various precondition for different tasks is insightful, such as sec 3.2, 3.3. \n3. Authors also conduct various experiments to support the their claims and highlight the importance of precondition and network design. Some experiments are quite interesting to me, such as Sec 5.1 and 5.3.", "weaknesses": "1. Though it shows improvement in other tasks, there are minimal improvements in diffusion model tasks. Authors should report performance of Multi-ResNet for diffusion models and compare it with existing U-Nets used widely in diffusion models.\n2. I would like authors to comments on how to design Unet for diffusion model so that it can recover high frequency details in high resolution image generation task. The average pooling will filter high frequency noise, but also filter the high frequency data information.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad to hear the reviewer appreciated our submission. \n\n> “Though it shows improvement in other tasks, there are minimal improvements in diffusion model tasks. Authors should report performance of Multi-ResNet for diffusion models and compare it with existing U-Nets [...]”\n\nIt is important for us to stress that generative modelling with diffusion models is only one of the three tasks (besides PDE surrogate modelling and image segmentation) which we tackle to verify hypotheses and ideas derived from our unified U-Net framework. Diffusion models are hence not the core focus of this work, it is rather one (currently very popular) area of research where the U-Net architecture produces strong results. \n\nOur key findings are irrespective of the choice of neural network one uses to parameterise the NNs $\\mathcal{E}$ and $\\mathcal{D}$ (with the caveat that Multi-ResNets require a ResNet decoder). Comparing different U-Net architecture choices for diffusion models is hence a very interesting direction which should definitely be pursued, particularly in light of this paper, and we thank the reviewer for this idea. We yet believe it is not essential to support the contributions of this work. \n\nBesides, we would like to briefly recap 3 key improvements our unified framework provides *specifically for diffusion models*:\n\nA. Theoretical understanding:  While not an improvement for diffusion models in their design or performance, we believe Sec 4 is a crucial improvement and contribution for understanding the understudied success of U-Nets in diffusion models, especially on images. To the best of our knowledge, Theorem 2 for the first time proves why a U-Net with average pooling as their projection operators $\\mathcal{P}$, and regardless of the choice of operators $\\mathcal{E}$ and $\\mathcal{D}$ is a natural inductive bias for data under a forward noising process of a diffusion model. Specifically, Theorem 2 shows that – possibly in contrast to the intuition prior to this submission – that not all frequencies of an image are affected equally by the forward noising process.\n\nB. Diffusions over complicated geometries: Our framework enables the design of U-Nets over complicated geometries beyond the square, for instance CW-complexes or manifolds. This can for instance be used to design neural network architecture for diffusion models on the manifolds merely by encodes the topology of the data into U-Net itself. As a proof-of-concept with U-Nets over a Haar wavelet basis and on a triangular domain (see Sec 3.3 and Sec 5.3), and discuss how this can be extended via triangulations to design a U-net for a sphere (see Appendix A.2) and other triangulated spaces. \n\nThis may be particularly successful as at present, such work often leverages simple, fully-connected networks as their neural architecture. For instance in [4] (also mentioned by reviewer go4z), an award-winning diffusion models paper from last year’s NeurIPS, which extends diffusion models to Riemannian manifolds, the authors exclusively use fully-connected neural networks in their experiments (see appendix O in [4], “Architecture”). Our work is hence complementary to such work, in the sense that it enables a large, previously unrecognised potential in the neural architecture. \n\nC. Multi-resolution training and sampling: As we experimentally demonstrate in Sec 5.2, the identification of the self-similarity structure of a U-Net in Sec 2 enables us to naturally train and sample from diffusion models on multiple resolutions via Algorithm 1 (see Appendix B.2), by training $\\mathcal{U}$ on resolution $i-1$, then training $U_i$ while preconditioning on $U_{i-1}$. Two advantages arise: First, if data is available on a higher resolution $i$, we can reuse the U-Net on a lower resolution $U_{i-1}$ in a principled way, i.e. by preserving the self-similarity structure of the neural network. Second, we can checkpoint sampling models when prototyping, look at and possibly measure the quality of lower-resolution samples (see Fig. 6) first, which may be indicative for the quality of the higher resolution.\n\n> “I would like authors to comments on how to design Unet for diffusion model so that it can recover high frequency details [...].”\n\nWe find the reviewer's question intriguing. While we are uncertain we believe that the reviewer's inquiry pertains to super-resolution applications. Super-resolution techniques hinge on identifying long-range dependencies within image structures to predict higher resolutions. For the purposes of our response, we will assume that our image super-resolution algorithm is trained using high-resolution training data. \n\nFirst, we require the U-Net model to accommodate images of varying resolutions. The essential design requirement is the U-Net's ability to accept images of different resolutions. The original U-Net lacks this capability. However, our Multi-ResNet inherently addresses this issue, enabling the U-Net to handle input images with varying resolutions seamlessly. This is because the encoder has no learnable parameters, and in essence, makes various projections onto varying resolution spaces of the original image. \n\nSuppose now that we are given data on resolution $V_{i+1}$ (a higher resolution than $V_i$) and train a diffusion model with a Multi-ResNet on this data. Then we are given data of resolution $V_i$ and asked to produce a super resolution version of this image. What is important here is, as the reviewer says, that in this setup, any image of arbitrary resolution can be effortlessly embedded into one of the resolution input spaces, denoted as $V_i$, through interpolation. This elegantly extends the U-Net model to support resolutions of any scale, making it a versatile and straightforward approach applicable even to infinite resolution scenarios. This model may not perform well, but being the simplest U-Net model that is adaptive to this context, it forms a good baseline method to ablate from. \n"}, {"review_id": "VoecdsjgRU", "reviewer": "Reviewer_Q8QJ", "summary": "The paper presents a unified framework for U-net design, by generalizing the overall structure of U-nets into different components. The framework is then investigated from a theoretical point of view. At its core, the paper presents Multi-Resnets which is a novel class of U-nets, which is then tested rigorously in the experimental section, in relation to generative modelling, PDE modelling and image segmentation.", "questions": "Maybe out of scope for the paper, but let me ask anyway:\nThe original reason for resnet being introduced was to solve the problem of vanishing gradients. Do you see a connection between this problem and your work?\n\n", "limitations": "The authors provide a lengthy explanation of what the limitations of the framework is in the appendix of the paper, however I think the authors should have definitely included this in the main paper. \nTo me, while the framework generalizes U-Nets it does not necessarily guarantee that the design process becomes easier, as the authors also touch on L162-166. This indeed does seem to be a limitation of the framework, when the initial goal was to \"..designing and analysing general U-Net architectures\".\n", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "It is always great to see papers that take a step back and tries to unify different work into a framework, to conceptualize what is up and down in a certain field. In that sense this paper success in doing this. In particular, I think sections such as 5.1 are some of the stronger sections of the paper because they conceptualize what the role of the encoder is.\n\nOriginality:\nThe framework presented in the paper seems to be original compared to other work.\n\nQuality:\nThe paper seems to be of a high quality, with both theoretical results to backup the claims of the paper but also empirical evidence to support. This in especially true, when the appendix of the paper is taken into account as it provide much more information, especially on the experimental section of the paper.\n\nClarity:\nThe paper is fairly clear in what it is trying to achieve, however since there are a lot of moving parts to the paper the overall structure sometimes suffer from this and the overall story become a bit \"muddy\".\n\nSignificance:\nWhile the conceptualization of U-nets indeed is refreshing, I doubt this paper will have a huge impact on domains that use U-nets as most established fields seems to build incrementally upon each others instead of redesigning architectures. That said, only time will tell. ", "weaknesses": "I did not in particular find any weaknesses of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback on our work. We are delighted that you appreciate the approach of \"taking a step back\" in our research and our analysis of the role of the encoder in U-Nets (Sec 5.1), as well as the originality and quality of our theory, experiments, and Appendix. Your support is greatly valued.\n\nWe would like to respond to your review below. We appreciate the opportunity to address any points you have raised and further engage in constructive discussions about our research.\n\n> “It is always great to see papers that take a step back and tries to unify different work into a framework […] In that sense this paper success in doing this. Significance: [...] I doubt this paper will have a huge impact on domains that use U-nets as most established fields seems to build incrementally upon each others instead of redesigning architectures. That said, only time will tell.”\n\nOn significance, we concur that many fields evolve through incremental advancements, but the very act of 'taking a step back' has enabled us to unearth several novel insights. For instance, our exploration into the role of the encoder in U-Nets led to the design of Multi-ResNets. Furthermore, our findings on U-Nets' self-similarity structure, high-resolution scaling behaviour, and the capability to integrate function constraints (like PDE boundary conditions) directly into the U-Net architecture are pivotal. We've also ventured into designing U-Nets for complex geometries, such as manifolds and spheres in particular.\n\nOne of our aspirations with this paper is to encourage the exchange of U-Net design choices across different domains. As an example, the recent PDEArena benchmark [7] demonstrated the high applicability of U-Nets popular in computer vision to PDE surrogate modelling. Notably, in [7], U-Nets are competitive with other architectures (such as Fourier Neural Operators (FNO)) in that field. We aim to foster a common language for this go-to, widely-used architecture design and believe that many more advancements can be made through such cross-domain efforts based on a unified framework. Our hope is that this work will serve as a catalyst for the exchange of ideas and foster advancements in U-Net design, ultimately benefiting a wide range of applications and research areas.\n\n> “Maybe out of scope for the paper, but let me ask anyway: The original reason for resnet being introduced was to solve the problem of vanishing gradients. Do you see a connection between this problem and your work?”\n\nWe thank the reviewer for this intriguing question. We have not extensively analysed the connection between our work and vanishing gradients, which have indeed been addressed in practice through the use of residual and skip connections. However, we firmly believe that there are fascinating properties on the vanishing gradient problem to be explored now that a U-Net is defined (see Def 1). Specifically, we are interested in defining a ‘stable U-Net’ (akin to similar work on stable ResNets [8]) by understanding the minimal requirements on the choice of $E_i$ and $D_i$ so that the model is guaranteed to have gradients bounded away from zero and infinity, independent of width and depth of the network. We acknowledge the importance of this area of research and look forward to further investigations into the stability properties of U-Nets.\n\nOn another note, our work was indeed greatly motivated by the seminal ResNet paper [6] through the idea of *preconditioning*, which is the core design principle of ResNets and U-Nets alike (also visualised in Figure 2 of our submission). This ultimately led us to the insight of a conjugacy between ResNets and U-Nets (see Prop 1).\n\n> “The authors provide a lengthy explanation of what the limitations of the framework is in the appendix of the paper, however I think the authors should have definitely included this in the main paper.”\n\nWe agree with the reviewer, the limitations discussed in Appendix A are important and useful to be shown in the main text. Following the reviewer’s suggestion, we will include this discussion in Sec 6 (Conclusion) of the main paper as part of the additional page in a potential camera-ready version of this manuscript.\n\n> “[...] To me, while the framework generalizes U-Nets it does not necessarily guarantee that the design process becomes easier [...]. This indeed does seem to be a limitation of the framework [...]”\n\nOur framework offers a decisive advantage in the design process of U-Nets, allowing successful design choices from one area to be effectively applied in another. One illustrative example is the construction of generative models over triangulated spaces, which has significant implications for modelling functions over complicated geometries.\n\nSpecifically, when faced with the task of designing a U-Net for a diffusion model with a favourable inductive bias on a triangulated complicated geometry, we have demonstrated in Theorem 2 that utilising a Haar wavelet decomposition of the geometry naturally induces a desired inductive bias for generative modelling. By opting for a Haar decomposition of the triangulated domain, the selection of spaces $V_i$ and $W_i$ is immediately defined, with Theorem 2 further specifying $P_i$. In this sense, our framework makes designing U-Nets ‘*easier*’ by providing a theoretical framework in which design choices for each of the components of a U-Net can be made based on theoretical insight and by incorporating problem constraints, rather than based on an empirically-driven ‘trial-and-error process’. Notably, our framework also extends the applicability of Theorem 2 to a broader range of geometries, enhancing the generality of its theoretical contributions.\n\nIn light of the reviewer's input, we recognise the need to make these connections more apparent in a potential camera-ready manuscript. We welcome any further suggestions on how to improve the communication of these ideas in our manuscript.\n"}, {"review_id": "QQR2VlpITZ", "reviewer": "Reviewer_ndAy", "summary": "This paper proposes a framework for designing and analyzing general UNet architectures. Theoretical results are presented that can characterize the role of encoder and decoder in UNet and conjugacy to ResNets is pointed out via preconditioning. Furthermore, this paper proposes Multi-ResNets, UNets with a simplified, wavelet-based encoder without learnable parameters. In addition, this paper presents how to design novel UNet architectures which can encode function constraints, natural bases, or the geometry of the data. Experiments of Multi-ResNets on image segmentation, PDE surrogate modeling, and generative modeling with diffusion model are conducted and demonstrate competitive performance compared to classical UNet.", "questions": "Please refer to the weakness part.", "limitations": "The limitations of this paper are well discussed.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Provide the mathematical definition of UNet, which enable identifying self-similarity structure, high-resolution scaling limits, and conjugacy to ResNets via preconditioning. \n- Based on theoretical analysis, authors propose Multi-ResNets, a novel class of UNet with no learnable parameters in its encoder.\n- Multi-ResNets achieves competitive or superior results compared to a classical U-Net in PDE modeling, image segmentation, and generative modeling with diffusion models.\n", "weaknesses": "- The presentation of the paper is not smooth and clear, which needs significant improvement. For example, the authors claim to propose a unified framework for UNet design, but I don’t see how this model can help researchers design specific UNet for different tasks. It’s more like the authors improve UNet on three chosen tasks and provide some theoretical analysis. Please correct me if I misunderstood anything.\n- This paper did a poor evaluation of the PDE part. In the PDE surrogate modeling, UNet is not the state-of-the-art model (authors made a wrong and misleading claim about the performance of UNet in PDE modeling) and performs much worse than the Fourier Neural Operator (FNO) [1][2][3]. So I am not convinced that it is useful to only demonstrate the effectiveness of Multi-ResNets compared to UNet in PDE modeling. In addition, no other representative baselines are compared within the NS and shallow water experiments, such as FNO[1], MPPDE[4], FFNO[5], etc.\n\n[1] Li, Zongyi, et al. \"Fourier neural operator for parametric partial differential equations.\" arXiv preprint arXiv:2010.08895 (2020).\\\n[2] Takamoto, Makoto, et al. \"PDEBench: An extensive benchmark for scientific machine learning.\" Advances in Neural Information Processing Systems 35 (2022): 1596-1611.\\\n[3] Helwig, Jacob, et al. \"Group Equivariant Fourier Neural Operators for Partial Differential Equations.\" arXiv preprint arXiv:2306.05697 (2023).\\\n[4] Brandstetter, Johannes, Daniel Worrall, and Max Welling. \"Message passing neural PDE solvers.\" arXiv preprint arXiv:2202.03376 (2022).\\\n[5] Tran, Alasdair, et al. \"Factorized fourier neural operators.\" arXiv preprint arXiv:2111.13802 (2021).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review, and for highlighting various components and insights of our unified framework as strengths of our work. Below, we would like to focus on the two weaknesses the reviewer mentions.\n\n> “The presentation of the paper is not smooth and clear, which needs significant improvement.”\n\nYour feedback concerning the paper's clarity is invaluable, and we are committed to making necessary revisions. If you have further specific points or sections that you found unclear, we would greatly appreciate further guidance to target those areas effectively.\n\n> “[...] I don’t see how this model can help researchers design specific UNet for different tasks. It’s more like the authors improve UNet on three chosen tasks [...]. Please correct me if I misunderstood anything.”\n\nWe value your recognition of our efforts in enhancing U-Nets for several tasks. Regarding your concern, our paper aims to unify the U-Net definition and component roles based on literature. The inclusion of Multi-ResNets is not for introducing a new U-Net architecture per se, but rather as a tool to analyse the encoder's role (Sec 3.1) and validate it empirically (Sec 5.1). This approach empowers researchers to make informed U-Net component choices, in contrast to empirical iteration practised.\n\nIn Sec 3, we offer three examples of novel design choices within our framework that incorporate task-specific requirements and constraints into U-Nets. Among them are: \n\n***U-Nets for complicated geometries:***\nOur framework enables the design of U-Nets over complicated geometries beyond the square, for instance CW-complexes or manifolds. See “B. Diffusions over complicated geometries” in our response to reviewer 7MgH for a complete explanation.\n\n***U-Nets with constraints:*** \nSuppose that we wish to build a U-Net which is guaranteed to output functions of a certain smoothness class and boundary condition. This could be for making a U-Net for PDE surrogate modelling, where we know a priori that the PDE has a weakly differentiable solution with nullified boundary conditions. Then for $W_i$ in our U-Net definition, we can choose a basis of triangular basis functions (Sec 3.2). We encode this within our neural network by selecting a layer of the $D_i$ to model the coefficients of this basis to a given dimension through what our framework and definition propose for selecting the approximation space which is desirable to the problem at hand. \n\nWe would appreciate it if you could consider these two examples (alongside the Multi-ResNet)  in response to your comment on designing task-specific U-Nets towards your final evaluation of our work.\n\n> “This paper did a poor evaluation of the PDE part. In the PDE surrogate modeling, UNet is not the state-of-the-art model (authors made a wrong and misleading claim about the performance of UNet in PDE modeling) and performs much worse than the Fourier Neural Operator (FNO) [1][2][3]. [...] In addition, no other representative baselines are compared within the NS and shallow water experiments, such as FNO[1], MPPDE[4], FFNO[5], etc.”\n\nIt is important for us to stress that this paper is not about PDE surrogate modelling. This paper is about U-Nets, their design and analysis. In our experiments, we chose PDE surrogate modelling as one of three tasks, because U-Nets are a competitive model candidate for each of them, and often the go-to choice for practitioners. The goal of Sec 5.1 was to verify our theory on the role of the encoder, and indeed provide empirical evidence, including on PDE modelling, that parameters in the encoder may sometimes not be as useful as one might think, if the encoder spaces were chosen suitably—a non-trivial insight enabled by our theory. We believe it would be interesting to benchmark U-Nets and Multi-ResNets against further models such as the ones the reviewer mentions, possibly even on more datasets, yet believe this is beyond the scope of this paper. \n\nWe respectfully disagree with the statement that we had made a “wrong and misleading claim about the performance of [the U-Net] in PDE modelling”. In our submission, we have neither claimed that U-Nets outperform FNO, nor have we claimed that U-Nets are ‘the best’ model for PDE modelling at present. Analysing this was not a goal of this paper. Our intent was to highlight U-Nets as a competitive architecture for the tasks at hand, not necessarily as the unequivocal best choice. We regret any misunderstanding that the term \"state-of-the-art\" may have caused. While sources like PDEBench [2] also reference U-Nets in this manner (“[...] with results for popular state-of-the-art ML models (FNO, U-Net, PINN) [...]” [2]), we'll revise the term to \"competitive\" in all occurences to prevent any misinterpretations. \n\nThat said, we disagree with the reviewer's assessment of FNO's dominance over U-Nets. We believe references [1-3] do not support their claims of clear FNO superiority for various reasons, and we're happy to discuss these papers at length individually during the author-reviewer discussion period. We would further like to refer to the PDEArena benchmark [7], the perhaps most recent, large-scale PDE benchmark which reports U-Nets as competitive, sometimes clearly superior to FNO. In addition, we have run independent, new experiments beyond our submission which underline this on the scale which we report in our submission: in Table 1 in the author rebuttal PDF, we have conducted a comparison of our reported U-Nets with FNO models of similar size. These results match the results in PDEArena by order of magnitude (see Table 2 in author rebuttal PDF), noting that we use slightly different experimental configs. Both report the result that U-Nets outperform FNO. We hence conclude that describing U-Nets as “competitive” for PDE Modelling is an appropriate claim.\n"}, {"review_id": "JPNLmwyI8B", "reviewer": "Reviewer_go4z", "summary": "This paper proposes a formal definition of U-nets, a crucial building block of modern deep learning pipelines such as diffusion models.\nThanks to this definition it is possible to both get theoretical results explaining some of the u-nets behaviours, and generalize u-nets to settings more exotic than the 2D images.\nExperiments are then conducted on various modalities to validate the theory.", "questions": "- What is meant by self-similarity?\n- In the code, for MNIST experiment, a network called GNet is used rather than a U-Net : what is it and does it fit with the experiments reported in the paper?\n- When deriving Theorem 2, it seems that each subspace W_i is an image subspace. However, in practice the U-Net acts on feature spaces, i.e. spaces with much more feature channels than the image itself, how would you take that into account?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- I agree with the main selling point of the paper : U-Nets have been under studied, especially in the context of diffusion models, and it is more than necessary to put an end to that.\n- I think the formalism effort presented in this work is great, and it was indeed, to the best of my knowledge, lacking. \n- the diversity of the modalities studied showcases the potential impact of better understanding U-Nets", "weaknesses": "- **inductive bias**: theorem 2 is the cornerstone of the theoretical explanation of this work as to why U-Nets are effective in diffusion. Yet I think that it has 2 problems. First, the result is about the Haar decomposition of the diffusion process, and therefore it applies to all models that have a multiscale structure with average pooling, not specifically U-Nets. Second, and in my view much more problematic, diffusion pipelines often use an \"epsilon-type\" prediction framework, in which the role of the U-Net is not to predict the denoised image (or signal in general), but rather the noise.\nThere are multiple examples of this that can be found either in the examples of the diffusers library of HuggingFace or in the seminal papers of Ho and Song:\n    - in the unconditional generation example: https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py#L576, when prediction type is \"epsilon\", the loss is computed as the MSE between the output of the model and the noise, and one can see that this is the default mode (L229), not changed in the way the script is called (https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/README.md)\n    - in the text-to-image generation example, the prediction type is taken by default from the noise scheduler which has it by default has \"epsilon\" (https://github.com/huggingface/diffusers/blob/716286f19ddd9eb417113e064b538706884c8e73/src/diffusers/schedulers/scheduling_ddpm.py#L121), not changed in the way the script is called (https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/README.md)\n    - In [7], algorithm 1 clearly shows that it's the noise that is learned\n    - In the implementation of [34], the noise `z` is used used in the MSE with the model output (https://github.com/yang-song/score_sde/blob/main/losses.py#L115)\nAs a consequence I think the theoretical result does not cover the practice. Indeed, it is said \"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.\" while the network's role is precisely to focus on subspaces which feature the noise in order to return it. While I understand that some pipelines still use a \"sample\" type prediction, I think that the understanding of the U-Net inductive bias is not understood if it doesn't take into account both types of predictions, especially if the explanation of one type is contrarian to the other type.\n- **clarity and space use**: I think that some sections of the paper are very unclear, in part because some essential parts were left out in the appendix. For example, Algorithm 1 for staged training should be in the core paper. Similarly the practical construction of a U-Net that guarantees boundary conditions should be in the core paper. For now, we just get an idea of what the decoder subspaces should look like, but not how the decoder should be designed in order to achieve these subspaces. Right now it is not clear to me how these decoders are designed. A last example (although there are some more) is that the formal definition of the U-Net is hard to grasp at first read. There should be a correspondence between the elements introduced and the typical 2D image segmentation U-Net architecture.\n- **prior work comparison/mention**: I noticed two parts that needed more references to prior work/comparison to it. First, the connection between U-Nets and wavelets has been a topic of many works among which [A, B]. Second, diffusion on complicated geometries including manifolds has been recently covered by works such as [C-F]. For both of these, mentioning as well as highlighting the differences would be super important.\n\n*Minor*:\n- To me preconditioning means finding a matrix close to the Hessian inverse in an optimization procedure in order to better take into account the geometry of the problem when optimizing, not finding a good initialization. Here it seems that preconditioning is more used to describe a good initialization rather than taking into account geometric information.\n- \"It also explains why in practice, encoders are often chosen significantly smaller than the decoder [26]\". I am not sure about this fact, so I think it needs much more evidence (see for example in the diffusers library https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d.py#L198\nthe decoder only has one more layer than the encoder). Also [26] is about VAEs, not U-Nets.\n\n\n\n[A]: Ramzi, Z., Michalewicz, K., Starck, JL. et al. Wavelets in the Deep Learning Era. J Math Imaging Vis 65, 240–251 (2023).\n\n[B]: Liu, P., Zhang, H., Zhang, K., Lin, L., & Zuo, W. (2018). Multi-level wavelet-CNN for image restoration. In Proceedings of the IEEE conference on computer vision and pattern recognition workshops (pp. 773-782).\n\n[C]: De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., & Doucet, A. (2022). Riemannian score-based generative modeling. arXiv preprint arXiv:2202.02763.\n\n[D]: Fishman, N., Klarner, L., De Bortoli, V., Mathieu, E., & Hutchinson, M. (2023). Diffusion Models for Constrained Domains. arXiv preprint arXiv:2304.05364.\n\n[E]: Huang, C. W., Aghajohari, M., Bose, J., Panangaden, P., & Courville, A. C. (2022). Riemannian diffusion models. Advances in Neural Information Processing Systems, 35, 2750-2761.\n\n[F]: Liu, L., Ren, Y., Lin, Z., & Zhao, Z. (2022). Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful and detailed review. We are thankful for the appreciation of the formalism of our theoretical framework which contributes to an understudied area, among other strengths. We respectfully would like to point out that we disagree with the concerns on Theorem 2 which seem to be the reviewer’s main reason for discontent.\n\n> “First, the result is about the Haar decomposition of the diffusion process, and therefore it applies to all models that have a multiscale structure with average pooling, not specifically U-Nets.”\n\nTheorem 2 analyses the effect of average pooling (wavelet projection [4]) on data under a forward diffusion process. It is worth noting that Theorem 2 does not require or mention the concept of a U-Net per se, but is consequential to the model design of diffusion models. We appreciate your insight, which highlights the theorem's broader applicability to various architectures leveraging a multiscale structure with average pooling. We don’t think this is a “problem” of Theorem 2, it rather shows its strength and generality beyond U-Nets.\n\nThe reviewer’s description of \"multiscale structure with average pooling\" captures what many might view as a U-Net's essence. This is just a curious example that the lack of a formal definition in the community necessitates our introduction of Definition 1. This definition serves not just as a concise representation, but also facilitates the deep dive into U-Nets' characteristics, from their recursive patterns (2.1) and scaling laws (2.2) to their relationship with ResNets (2.3). \n\n> “Second, [...] diffusion pipelines often use an \"epsilon-type\" prediction framework, in which the role of the U-Net is not to predict the denoised image (or signal in general), but rather the noise. [...] As a consequence I think the theoretical result does not cover the practice.”\n\nWe respectfully disagree with this main concern on Theorem 2, and would like to find an agreement with the reviewer on this point.\n\nFirst, we would like to note that both the “epsilon-type” (noise as output) and the “sample-type” (image as output) solve *the same task*, namely to separate signal from noise. While in both frameworks the network receives image and noise $Im + \\epsilon$ as input, the “epsilon-type” framework achieves this task by subtracting the signal $Im$, while the “sample-type” framework subtracts the noise $\\epsilon$. \n\nThe fundamental quantity of interest which governs how challenging it is to solve this task of separating signal and noise (in both formulations) is the signal-to-noise ratio between $Im$ and $\\epsilon$. In Theorem 2, we demonstrate that the high-frequency components of an image have their signal-to-noise ratio dominated by noise exponentially faster than the low-frequency components. Importantly, a U-Net architecture with average pooling is passing as input to the decoder the image plus noise at the highest resolution, and the low frequency details that are dominated by signal on lower resolutions. This means that apriori to training the network, we have fed as input the image plus noise, and a first approximation of the image without noise. As the network learns to separate noise from signal, these two inputs yield a useful inductive bias for the signal and noise separation task in light of Theorem 2.\n\nWhile not entirely certain, we believe the confusion might arise from the following argument. During average pooling, the high-resolution Haar wavelet spaces, which are dominated by noise (Theorem 2), get discarded. However, in the  \"epsilon-type\" formulation, it may be unclear how the U-Net can now predict the noise, if ‘we just removed the noise’ via average pooling (more precisely, those subspaces dominated by noise). Should this indeed be the source of confusion, we can hopefully resolve this easily. The responsibility for ‘predicting the noise’ does not solely lie on the encoder. Instead, it's the decoder that uses the encoder's signal through skip connections. The encoder primarily functions as an ‘information compressor’ for the input data. The decoder utilises this condensed input to generate an output, which could be either the 'sample' or the 'epsilon' (noise). In either case, a lower-frequency, low-noise signal serves as valuable compressed information that aids the decoder in generating its output.\n\n > ”I noticed two parts that needed more references to prior work/comparison to it [...] mentioning as well as highlighting the differences would be super important.”\n\nU-Nets and Wavelets: We appreciate the reviewer's references, especially [A], and will include [A,B] in our discussion of Haar subspace U-Nets. While these works highlight wavelet-U-Net connections, our Definition 1 unifies diverse wavelet bases and generalises to other bases. Our empirical study of Multi-ResNets, replacing the U-Net's encoder with wavelet projections, adds novel insight. This demonstrates the broader connection between U-Nets and wavelets [5].\n\nDiffusions over Complex Geometries: U-Nets achieve state-of-the-art results in learning diffusion models for image-based tasks [9]. Yet, their performance can waver when modelling functions on non-Euclidean spaces. Notably, recent research into diffusion models on Riemannian manifolds [C-F] has favoured simple, fully-connected networks over U-Nets or modifications thereof. \n\nOur unified framework bridges this gap by unveiling previously overlooked potential in neural architecture. Our work is hence complementary to [C-F], focussing on the neural architecture. It paves the way for U-Net adaptations tailored to complex geometries, including CW-complexes and manifolds by integrating data topology into the U-Net itself. This lays the foundation for neural architectures that suit diffusion models on manifolds. \n\n- - -\n\nLastly, we refer to our response to all authors for a discussion of other points in this review.\n"}]

Author rebuttal:
We thank the reviewers for their valuable and insightful comments and questions. Motivated by the reviews, we make several changes towards a potential camera-ready version of the submission, using the one additional page, which we list below. 

### Summary of key changes towards a camera-ready version

* A note on how to extend our framework to arbitrary channel dimensions using the Kronecker product in Sec 2.1.
* A paragraph on the limitations of our work in Sec 6 (Conclusion), based on the respective section in Appendix A.
* Algorithm 1 and Appendix A.3 as a condensed version, now in Sections 5.2 and 3.2.

### Response on other points from reviewer go4z: 

> “Algorithm 1 for staged training should be in the core paper [...] a U-Net that guarantees boundary conditions should be in the core paper.”

We will address this clarity issue by incorporating practical approaches for guaranteed boundary conditions with a U-Net, adding Algorithm 1 to Sec 5.2 and a condensed version of Appendix A.3 in Sec 3.2 as per the reviewer's insightful feedback.

> The formal definition of the U-Net is hard to grasp at first read. There should be a correspondence between the elements introduced and the typical 2D image segmentation U-Net architecture.

We agree and will provide a figure further illustrating Definition 1 in a potential camera-ready version. 

> To me preconditioning means finding a matrix close to the Hessian inverse in an optimisation procedure [...] Here it seems that preconditioning is more used to describe a good initialisation [...]

We acknowledge the reviewer's observation that “preconditioning” can refer to finding a matrix close to the Hessian inverse for an optimisation procedure. However, we would like to emphasise that “preconditioning” is a term which implies different meanings across various contexts in mathematics. The seminal ResNet paper uses the word ‘precondition’ to indicate ‘initialisation to the identity: “In real cases, it is unlikely that identity mappings are optimal, but our reformulation may help to precondition the problem [6].” As this U-Net work has strong ties with ResNets (e.g. their conjugacy in Prop 1), we considered using it in the same way seems appropriate. We appreciate the reviewer's feedback, which allows us to enhance the clarity and quality of our work.

> “[...] how the decoder should be designed in order to achieve these subspaces.”

We value the reviewer's input on improving our manuscript quality. The decoder's design corresponds to predicting linear combinations of subspace $W_i$ basis vectors. Haar wavelets use a pixel grid as a valid basis, Fourier layers for instance transform into linear combinations of Fourier modes of certain frequencies. Generally, $w_i \in W_i$ is expressed as $w_i = \sum c_i e_i$ with basis vectors $e_i$. The neural network predicts coefficients $c_i$ for a fixed data basis. We'll explicitly address this in the paper update and appreciate the reviewer's observation which helps clarify this point.

> “It also explains why in practice, encoders are often chosen significantly smaller than the decoder [26]’. I am not sure about this fact [...]”

[26] indeed focuses on hierarchical VAEs, which however employ U-Net-style architectures, extensively discussed in [5] (Appendix B.2). Notably, [26] presents an interesting scenario with a smaller encoder than the decoder, possibly explained by our framework. We'll revise 'often' to 'in some cases' to better reflect varying encoder-decoder choices in diffusion models. Thank you for this remark.

> “What is meant by self-similarity?”

We appreciate the reviewer's input on terminology clarity. "Self-similarity" in our context refers to refining a square image into smaller squares to establish nested approximation spaces ${W_i}$ for our U-Net. This enables the U-Nets' adaptability to recursively refined geometries. While acknowledging the non-standard nature of this term, we'll enhance precision in Sec 2.2 to address this feedback. Thank you for helping refine this aspect of our work.

> “What is it [GNet] and does it fit with the experiments reported in the paper?”

We appreciate the reviewer's notice of a code typo. "GNet" was our code name for the ‘general’ U-Net class. This term was specific to initial diffusion experiments, and we will correct it in the potential camera-ready submission.

> “U-Net acts on feature spaces, i.e. spaces with much more feature channels than the image itself, how would you take that into account?”

Our framework easily extends to multiple channels by utilising the kronecker product to expand W_i to the desired channel count in a U-Net. In a potential camera-ready version, we'll add a note in Sec 2.1 and a dedicated appendix section for the construction details. Thank you for this clarifying comment.

### References:

[1-3]: see Reviewer ndAy

[4] De Bortoli, V., et al. 2022. Riemannian score-based generative modelling. Advances in Neural Information Processing Systems, 35, pp.2406-2422.

[5] Falck, F. et al., 2022. A Multi-Resolution Framework for U-Nets with Applications to Hierarchical VAEs. Advances in Neural Information Processing Systems, 35, pp.15529-15544.

[6] He, K., et al., 2016. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).

[7] Gupta, J.K. and Brandstetter, J., 2022. Towards multi-spatiotemporal-scale generalized pde modeling. arXiv preprint arXiv:2209.15616.

[8] Hayou, S., et al., 2021, March. Stable resnet. In International Conference on Artificial Intelligence and Statistics (pp. 1324-1332). PMLR.

[9 Hoogeboom, E.,  et al., 2023. simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint arXiv:2301.11093.


Meta-review: All reviewers agreed that given the success of U-Nets in applications like diffusion modeling and PDEs, the paper does a compelling analysis of different choices in its design space and how it's domain of usefulness can be expanded to more complex geometries. While one reviewer disputed whether U-Nets are superior to FNOs for PDE modeling, this is tangential to the paper's primary focus. Nonetheless, U-Nets have certainly proven to be an important baseline in PDE modeling contexts too and given interest in efficient architectures from the community, would help others in the exploration of this design space.
Decision: Accept (poster)