SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Christopher Williams *1 Fabian Falck *,1,2 George Deligiannidis 1
Chris Holmes  $^{1,2}$  Arnaud Doucet  $^{1}$  Saifuddin Syed  $^{1}$
<sup>1</sup>University of Oxford <sup>2</sup>The Alan Turing Institute
{williams, fabian.falck, deligian, cholmes, doucet, saifuddin.syed}@stats.ox.ac.uk
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.
U-Nets (see Figure 1) are a 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-based diffusion models [8, 9, 10, 11, 12], U-Nets are a go-to architecture yielding state-of-the-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, we provide such a framework for U-Nets. This allows us to design U-Nets for data beyond a square domain, and enable us to incorporate prior knowledge about a problem, for instance a natural basis, functional constraints, or knowledge about its topology, into the neural architecture.
The importance of preconditioning. We begin by illustrating the importance of the core design principle of U-Nets: preconditioning. Preconditioning informally means that initialising an optimisation problem with a 'good' solution greatly benefits learning [13, 14]. Consider a synthetic example using ResNets [15]
Figure 1: A resolution 2 U-Net (Def. 1). If  $E_{i} = \mathrm{Id}_{V_{i}}$ , this is a Multi-ResNet (see Def. 3).
U-Nets as we will show in §2.3: we are interested in learning a ground-truth mapping  $w: V \mapsto W$  and  $w(v) = v^2$  over  $V = [-1, 1]$  and  $W = \mathbb{R}$  using a ResNet  $R^{\mathrm{res}}(v) = R^{\mathrm{pre}}(v) + R(v)$  where  $R^{\mathrm{pre}}, R: V \mapsto W$ . In Figure 2 [Left] we learn a standard ResNet with  $R^{\mathrm{pre}}(v) = v$  on a grid of values from  $V \times W$ , i.e. with inputs  $v_i \in V$  and regression labels  $w_i = w(v_i) = v_i^2$ . In contrast, we train a ResNet with  $R^{\mathrm{pre}}(v) = |v|$  [Right] with the same number of parameters and iterations. Both networks have been purposely designed to be 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  $R^{\mathrm{pre}}(v) = |v|$  is a 'good' initial guess or preconditioner for  $w(v) = v^2$ , but  $R^{\mathrm{pre}}(v) = v$  is a 'bad' one. This shows the importance of encoding good preconditioning into a neural network architecture and motivates us studying how preconditioning is used in U-Net design.
In this paper, we propose a mathematical framework for designing and analysing general U-Net architectures. We begin with a comprehensive definition of a U-Net which characterises its components and identifies its self-similarity structure which is 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 flexible preconditioning from lower-resolutions inputs.
Figure 2: The importance of preconditioning.
Our U-Net framework paves the way to designing U-Nets which can model distributions over complicated geometries beyond the square, for instance CW-complexes or manifolds, or diffusions on 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 a problem. We also analyse why U-Nets with average pooling are a natural inductive bias in diffusion models.
More specifically, our contributions are as follows: (a) We provide the first rigorous definition of U-Nets, which enables us to identify their 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 a hierarchy of orthogonal wavelet spaces of  $L^2(\mathbb{X})$ , for compact domain  $\mathbb{X}$ , with no learnable parameters in its encoder. In our experiments, Multi-ResNets yield competitive and often superior results when compared to a classical U-Net in PDE modelling, image segmentation, and generative modelling with diffusion models. We further show how to encode problem-specific information into a U-Net. In particular, 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 is dominated by noise exponentially faster than lower-frequency terms. We show how U-Nets with average pooling exploit this observation, explaining their go-to usage.
2 U-Nets: Neural networks via subspace preconditioning
The goal of this section is to develop a 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 a U-Net
Definition 1. U-Net. Let  $V$  and  $W$  be measurable spaces. A  $U$ -Net  $\mathcal{U} = (\mathcal{V}, \mathcal{W}, \mathcal{E}, \mathcal{D}, \mathcal{P}, U_0)$  comprises six components:
1. Encoder subspaces:  $\mathcal{V} = (V_i)_{i=0}^{\infty}$  are nested subsets of  $V$  such that  $\lim_{i\to \infty} V_i = V$ .
2. Decoder subspaces:  $\mathcal{W} = (W_i)_{i=0}^{\infty}$  are nested subsets of  $W$  such that  $\lim_{i\to \infty}W_i = W$ .
3. Encoder operators:  $\mathcal{E} = (E_i)_{i=1}^{\infty}$  where  $E_i : V_i \mapsto V_i$  denoted  $E_i(v_i) = \tilde{v}_i$ .
4. Decoders operators:  $\mathcal{D} = (D_i)_{i=1}^{\infty}$  where  $D_i: W_{i-1} \times V_i \mapsto W_i$  at resolution  $i$  denoted  $D_i(w_{i-1}|v_i)$ . The  $v_i$  component is called the skip connection.
5. Projection operators:  $\mathcal{P} = (P_i)_{i=0}^{\infty}$ , where  $P_i : V \mapsto V_i$ , such that  $P_i(v_i) = v_i$  for  $v_i \in V_i$ .
6. Bottleneck:  $U_{0}$  is the mapping  $U_{0}: V_{0} \mapsto W_{0}$ , enabling a compressed representation of the input.
The  $U$ -Net of resolution  $i$  is the mapping  $U_{i}:V_{i}\mapsto W_{i}$  defined through the recursion (see Figure 3):
$$
U _ {i} \left(v _ {i}\right) = D _ {i} \left(U _ {i - 1} \left(P _ {i - 1} \left(\tilde {v} _ {i}\right)\right) \mid \tilde {v} _ {i}\right), \quad i = 1, 2, \dots . \tag {1}
$$
We illustrate the Definition of a U-Net in Figure 1. Definition 1 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 [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_{i}$  and  $W_{i}$  for a given resolution  $i$  to be product spaces  $V_{i} = \bigotimes_{k=1}^{M} V_{i,k}$  and  $W_{i} = \bigotimes_{k=1}^{M} W_{i,k}$  for  $M$  channels 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 1 not only expands the scope of U-Nets beyond
Figure 3: Recursive structure of a U-Net.
problems confined to squared domains, but also naturally incorporates problem-specific 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 a U-Net.
Encoder and decoder subspaces. We begin with the spaces in  $V$  and  $W$  which a U-Net acts on. 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 to suit our task. In order to perform computations in practice, we must restrict ourselves to subspaces  $V_{i}$  and  $W_{i}$  of the potentially infinite-dimensional spaces  $V$  and  $W$ . For instance, if our U-Net is modelling a mapping between images, which can be viewed as bounded functions over a squared domain  $\mathbb{X} = [0,1]^2$ , it is convenient to choose  $V$  and  $W$  as subspaces of  $L^2 (\mathbb{X})$ , the space of square-integrable functions [19, 20] (see Appendix C). Here, a data point  $w_{i}$  in the decoder space  $W_{i}$  is represented by the coefficients  $c_{i,j}$  where  $w_{i} = \sum_{j}c_{i,j}e_{i,j}$  and  $\{e_{i,j}\}_{j}$  is a (potentially orthogonal) basis for  $W_{i}$ . We will consider precisely this case in §3.1. The projected images on the subspace  $V_{i}, W_{i}$  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 in contrast to common understanding, the roles of the encoder and decoder in a U-Net are fundamentally different from each other. The decoder on resolution  $i$  learns the transition from  $W_{i - 1}$  to  $W_{i}$  while incorporating information from the encoder on  $V_{i}$  via the skip connection. The encoder  $E_{i}$  can be viewed as a change of basis mapping on the input space  $V_{i}$  at resolution  $i$  and is not directly tied to the approximation the decoder makes on  $W_{i}$ . This learned change of basis facilitates the decoder's approximation on  $W_{i}$ . In §3.1, we will further extend our 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 a Hilbert space. Note that there is no embedding operator<sup>3</sup>, the operator facilitating the transition to a higher-resolution space, explicitly defined as it is invariably the natural inclusion of  $W_{i - 1}$  into  $W_{i}$ .
Self-similarity of U-Nets via preconditioning. The key design principle of a U-Net is preconditioning. The U-Net of resolution  $i - 1$ ,  $U_{i - 1}$  makes an approximation on  $W_{i - 1}$  which is input of and preconditions  $U_{i}$ . Preconditioning facilitates the transition from  $W_{i - 1}$  to  $W_{i}$  for the decoder. In the encoder, the output of  $P_{i - 1}E_{i}$  is the input of  $U_{i - 1}$ . When our underlying geometry is refinable (such as a square), we may use a refinable set of basis functions. In the standard case of a 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
2Implementation wise this corresponds to taking the Kronecker product across the spaces  $V_{i,k}$  and  $W_{i,k}$ , 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.
U-Nets reveals a 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  $i - 1$  is a coarse approximation for a U-Net on resolution  $i$ . We formalise this notion of preconditioning in Proposition 1 in §2.3.
2.2 High-resolution scaling behavior of U-Nets
Given equation (1), it is clear that the expressiveness of a U-Net  $\mathcal{U}$  is governed by the expressiveness of its decoder operators  $\mathcal{D}$ . If each  $D_{i}$  is a universal approximator [21], then the corresponding U-Nets  $U_{i}$  likewise have this property. Assuming we can represent any mapping  $U_{i}:V_{i}\mapsto W_{i}$  as a U-Net, our goal now is to comprehend the role of increasing the resolution in the design of  $\mathcal{U}$ , and to discern whether any function  $U:V\to W$  can be represented as a high-resolution limit  $\lim_{i\to \infty}U_i$  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 us to define  $\mathcal{W}$  as an increasing sequence of orthogonal subspaces of  $W$ . Possible candidates for orthogonal bases include certain Fourier frequencies, wavelets (of a given order) or 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 us define  $\mathcal{S}$  as infinite resolution data in  $V \times W$  and  $S_i$  as the finite resolution projection in  $V_i \times W_i$  comprising of  $(v_i, w_i) = (P_i(v), Q_i(w))$  for each  $(v, w) \in \mathcal{S}$ . Here,  $P_i$  is the U-Net projection onto  $V_i$ , and  $Q_i: W \mapsto W_i$  is the orthogonal projection onto  $W_i$ . Assume  $U_i^*$  and  $U^*$  are solutions to the finite and infinite resolution regression problems  $\min_{U_i \in \mathcal{F}_i} \mathcal{L}_i(U_i | S_i)$  and  $\min_{U \in \mathcal{F}} \mathcal{L}(U | S)$  respectively, where  $\mathcal{F}_i$  and  $\mathcal{F}$  represent the sets of measurable functions mapping  $V_i \mapsto W_i$  and  $V \mapsto W$ . Let  $\mathcal{L}_i$  and  $\mathcal{L}$  denote the  $L^2$  losses:
$$
\mathcal {L} _ {i} (U _ {i} | \mathcal {S} _ {i}) ^ {2} := \frac {1}{| \mathcal {S} _ {i} |} \sum_ {(w _ {i}, v _ {i}) \in \mathcal {S} _ {i}} \| w _ {i} - U _ {i} (v _ {i}) \| ^ {2}, \quad \mathcal {L} (U | \mathcal {S}) ^ {2} := \frac {1}{| \mathcal {S} |} \sum_ {(w, v) \in \mathcal {S}} \| w - U (v) \| ^ {2}.
$$
The following result analyses the relationship between  $U_{i}^{*}$  and  $U^{*}$  as  $i\to \infty$  where  $\mathcal{L}_{i|j}$  and  $\mathcal{L}_{|j}$  are the losses above conditioned on resolution  $j$  (see Appendix A).
Theorem 1. Suppose  $U_{i}^{*}$  and  $U^{*}$  are solutions of the  $L^2$  regression problems above. Then,  $\mathcal{L}_{i|j}(U_{i}^{*}) \leq \mathcal{L}_{|j}(U^{*})$  with equality as  $i \to \infty$ . Further, if  $Q_{i}U^{*}$  is  $V_{i}$ -measurable, then  $U_{i}^{*} = Q_{i}U^{*}$  minimises  $\mathcal{L}_{i}$ .
Theorem 1 states that solutions of the finite resolution regression problem converge to solutions of the infinite resolution problem. It also informs us how to choose  $V_{i}$  relative to  $W_{i}$ . If we have a  $\mathcal{V}$  where for the decoders on  $\mathcal{W}$ , the  $W_{i}$  component of  $U^{*}$ ,  $Q_{i}U^{*}$ , relies solely on the input up to resolution  $i$ , then the prediction from the infinite-dimensional U-Net projected to  $W_{i}$  can be made by the U-Net of resolution  $i$ . The optimal choice of  $V_{i}$  must be expressive enough to encode the information necessary to learn the  $W_{i}$  component of  $U^{*}$ . This suggests that if  $V_{i}$  lacks expressiveness, we risk efficiency in learning the optimal value of  $U_{i}^{*}$ . However, if  $V_{i}$  is too expressive, no additional information is gained, and we waste computational effort. Therefore, we should choose  $V_{i}$  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_{i}$ , 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 a conjugacy between U-Nets and ResNets. We begin by formalising ResNets in the context of U-Nets.
Definition 2. ResNet, Residual U-Nets. Given a measurable space  $X$  and a vector space  $Y$ , a mapping  $R: X \to Y$  is defined as a ResNet preconditioned on  $R^{\mathrm{pre}}: X \to Y$  if  $R(x) = R^{\mathrm{pre}}(x) + R^{\mathrm{res}}(x)$ , where  $R^{\mathrm{res}}(x) = R(x) - R^{\mathrm{pre}}(x)$  is the residual of  $R$ . A Residual U-Net is a U-Net  $\mathcal{U}$  where  $\mathcal{W}, \mathcal{V}$  are sequences of vector spaces, and the encoder and decoder operators  $\mathcal{E}, \mathcal{D}$  are ResNets preconditioned on  $E_i^{\mathrm{pre}}(v_i) = v_i$  and  $D_i^{\mathrm{pre}}(w_{i-1}|v_i) = w_{i-1}$ , respectively.
A preconditioner initialises a 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 as encoder and decoder operators. In encoders,
preconditioning on the identity on  $V_{i}$  allows the residual encoder  $E_{i}^{\mathrm{res}}$  to learn a change of basis for  $V_{i}$ , which we will discuss in more detail in §3.1. In decoders, preconditioning on the identity on the lower resolution subspace  $W_{i-1}$  allows the residual decoder  $D_{i}^{\mathrm{res}}$  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  $\mathcal{U}$  is a residual  $U$ -Net, then  $U_i$  is a ResNet preconditioned on  $U_i^{\mathrm{pre}}(v_i) = U_{i-1}(\tilde{v}_{i-1})$ , where  $\tilde{v}_{i-1} = P_{i-1}(E_i(v_i))$ .
Proposition 1 states that a U-Net at resolution  $i$  is a ResNet preconditioned on a U-Net of lower resolution. This suggests that  $U_{i}^{\mathrm{res}}$  learns the information arising from the resolution increase. We will discuss the specific case  $E_{i} = \operatorname{Id}_{V_{i}}$ , and  $U^{\mathrm{pre}}(v_i) = U_{i - 1}(P_{i - 1}(v_i))$  in §3.1. Proposition 1 also enables us to interpret  $\lim_{i\to \infty}U_{i}$  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.
Example 1. Haar Wavelet Residual U-Net. A Haar wavelet residual  $U$ -Net  $\mathcal{U}$  is a residual U-Net where:  $V = W = L^{2}(\mathbb{X})$ ,  $V_{i} = W_{i}$  are multi-resolution Haar wavelet spaces (see Appendix C.3), and  $P_{i} = \mathrm{proj}_{V_{i}}$  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  $L^2(\mathbb{X})$ , the space of square integrable functions and a Hilbert space, and use average pooling as the projection operation  $P_i$  [20]. Haar wavelets will in particular be useful to analyse why U-Nets are a 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 a square. Our framework also enables us to encode problem-specific 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
A closer look at the U-Net encoder. To characterise a U-Net in Definition 1, we must in particular choose the encoder subspaces  $\mathcal{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  $\mathcal{V}$  unfavourably? This is where the encoder comes in. While the encoder subspaces  $\mathcal{V}$  define an initial basis for our problem, the encoder learns a change of basis map to a new, implicit basis  $\tilde{v}_i = E_i(v_i)$  which is more favourable. This immediately follows from Eq. (1) since  $U_i$  acts on  $V_i$  through  $\tilde{v}_i$ . The initial subspaces  $\mathcal{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  $\mathcal{V}$ , the question whether and how much work the encoders  $\mathcal{E}$  have to do depends on how far away our choice is from the optimal choice  $\tilde{\mathcal{V}}$  for our problem. This explains why the encoders  $E_{i}$  are commonly chosen to be ResNets preconditioned on the identity  $E_{i}^{\mathrm{pre}} = \mathrm{Id}_{V_{i}}$ , allowing the residual encoder  $E_{i}^{\mathrm{res}}$  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_{i} = \mathrm{Id}_{V_{i}}$ . It also explains why in practice, encoders are in some cases chosen significantly smaller than the decoder [30], as a ResNet encoder need not do much work given a good initial choice of  $\mathcal{V}$ . It is precisely this intuition which motivates our second example of a U-Net, the Multi-Res(ulation) Res(idual) Network (Multi-ResNet).
Definition 3. Multi-ResNets. A Multi-ResNet is a residual U-Net with encoder  $E_{i} = \mathrm{Id}_{V_{i}}$ .
Example 2. Haar Wavelet Multi-ResNets. A Haar Wavelet Multi-ResNet is a Haar Wavelet Residual U-Net with encoder  $E_{i} = \mathrm{Id}_{V_{i}}$ .