SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
X _ {t} ^ {\mathbf {i}} = \sqrt {1 - \widetilde {\alpha}} X _ {0} ^ {\mathbf {i}} + \sqrt {\widetilde {\alpha}} \varepsilon^ {\mathbf {i}}, \tag {31}
$$
where  $\mathbf{i}$  is a string of length  $i + 1$  encoding which triangle on the globe the pixel is in, then which 'triangular pixel' we are in. We can construct Haar wavelets on the triangle (see Figure 3.3), and hence create natural projection maps in  $L^2$  over this domain. In this case  $W_0$  is the span of constant functions over each large triangle for the tessellation of the globe, and each  $W_i$  is  $W_0$  along with the span of the wavelets of resolution  $i$  over each triangle. Like in the standard diffusion case, we choose  $V_i = W_i$  and can use the mapping given in Section 5.2 to construct the natural U-Net design for this geometry, and for this basis. Similarly, for any space that we can form a triangulation for, we can use this construction to create the natural Haar wavelet space for this geometry, define a diffusion analogously, and implement the natural U-Net design for learning.
Figure 9: Triangulation of the globe with a self-similar triangle which admits a 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 a nascent research direction in machine learning research. We have seen in our PDE experiment (see Section 5) that our Multi-ResNet design outperforms Residual U-Nets on a 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
$$
\Delta u = f, \quad u (0) = u (1) = 0, \tag {32}
$$
which when multiplied by a test function  $\phi \in \mathcal{H}_0^1([0,1])$  and integrated over the domain puts the equation into its weak form, and by integration by parts
$$
- \left\langle \phi^ {\prime}, u ^ {\prime} \right\rangle_ {\mathcal {H} _ {0} ^ {1}} = \int \phi f, \tag {33}
$$
we recover the standard bilinear representation of Laplices equation. Now if our solution  $u \in \mathcal{H}_0^1([0,1])$  and we have an orthogonal basis  $\{\phi_{i,k}\}$  for  $\mathcal{H}_0^1([0,1])$ , then the coefficients for the basis elements of the solution can be solved by finding the solution of a linear system with a diagonal
mass matrix. Thus, in this way, the information needed to solve the linear system up to some finite resolution  $i$ , only depends on data up to that resolution in this weak form. This is a 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 in Theorem 1).
For an explicit example, we may make an orthogonal basis for  $\mathcal{H}_0^1 ([0,1])$  with the initial basis function [Left] of Figure 4 and its refined children on the [Right] of Figure 4.
In general we can refine further and create a basis of resolution  $i$  through
$$
\phi_ {k, j} (x) = \phi (2 ^ {k} x + j / 2 ^ {k}), \qquad \phi (x) = 2 x \cdot 1 _ {[ 0, 1 / 2)} (x) + (2 - 2 x) \cdot 1 _ {[ 1 / 2, 1 ]} (x),
$$
which are precisely the integrals of the Haar wavelets up to a given resolution. We would again construct  $\bar{W}_i$  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.
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 we refer to our code base for details, they generally feature a single residual block per resolution for  $E_{i}$  and  $D_{i}$ , with 2 convolutional layers, and group normalisation. We in general choose  $P_{i}$  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 to 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 to single-stage training.
Hyperparameters and hyperparameter tuning. We performed little to no hyperparameter tuning in our experiments. In particular, we did not perform a 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 to adjust the number of parameters so to enable a 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 hyperparam.py files.
Evaluation and metrics. We use the following three performance metrics in our experimental evaluation: FID score [44], rollout mean-squared-error (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 a holdout dataset not used during training, and using an 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 an 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 a large range of other evaluation metrics, and we typically observe similar trends as in those reported.
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-ResNets. 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  $j$  relative to the input resolution is computed by taking the lower-lower part of the DWT at level  $j$  ( $LL_{j}$ ), and then inverting  $LL_{j}$  to receive a coarse, projected version of the original input image.
Generative modelling with diffusion models. In Table 2, we 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 Multi-ResNet does not outperform the Residual U-Net with approximately the same number of parameters, as is the case in the other two tasks. Referring to our theoretical results in §3.1, the input space of Haar wavelets may be a 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, in Figure 10 we compare samples from two runs with the Residual U-Net and Multi-ResNet as reported in Table 2. We observe that it is 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.
Table 2: Quantitative performance of the (Haar wavelet) Multi-ResNet compared to 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 ↓
CIFAR10	Residual U-Net	35.5 M	7.86 ± 0.25
Multi-ResNet, no params. added in dec. (ours)	25.8 M	14.87 ± 0.50
Multi-ResNet, saved params. added in dec. (ours)	32.4 M	12.44 ± 0.22
Residual U-Net (35.5 M Parameters)
Multi-ResNet (32.4 M Parameters)
Figure 10: Samples of a DDPM-type diffusion model with a Residual U-Net [Left] and a Multi-ResNet [Right]. We trained both models for  $1.2\mathrm{M}$  iterations at which we obtained the samples.
Lastly, we present Figure 7 from the main text at a larger resolution in Figure 11.
$32 \times 32$  (at 1 M iter.)
CIFAR10
8 x 8 (at 100 K iter.)
Figure 11: Preconditioning enables multi-resolution training and sampling of diffusion models. We train a diffusion model [8] with a Residual U-Net architecture using Algorithm 1 with four training stages and no freezing on CIFAR10. This Figure is identical with Figure 7, but larger.
16 x 16 (at 150 K iter.)
$4 \times 4$  (at  $50 \mathrm{~K}$  iter.)
PDE modelling. We further present additional quantitative results comparing the (Haar wavelet) Multi-ResNet with (Haar wavelet) Residual U-Nets on our PDE modelling tasks. These shall highlight the importance of allocating the saved parameters in Multi-ResNets 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 'v1', also presented in the main text in Table 1 but repeated for ease of comparison, we add the saved parameters by increasing the number of 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 of hidden channels in the ResNet.
Our results show that 'v1' 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 in  $\mathcal{D}$  matters for the performance of Multi-ResNets. We note that in our experiments, beyond this comparison, we did not explore optimal ways of allocating parameters in Multi-ResNets. Hence, future work should explore this aspect thoroughly with large-scale experiments. Due to its 'uni-directional', 'asymmetric' structure, the Multi-ResNet may further
be combined with a 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 a better decoder would make Multi-ResNets superior to classical Residual U-Nets for diffusion models.
Table 3: Quantitative performance of the (Haar wavelet) Multi-ResNet compared to a classical (Haar wavelet) Residual U-Net on two PDE modelling. This table is an augmented version of Table 1 in the main text, where 'v1' indicates the run from the main text, a Multi-ResNet with saved parameters added in the encoder, and 'v2' indicates an alternative parameter allocation. We report Mean-Squared-Error over a rolled out trajectory in time (r-MSE) on the test set, rounded to four decimal digits.
Dataset	Architecture	# Params.	r-MSE ↓
Navier-stokes128×128	Residual U-Net	34.5 M	0.0057 ± 2·10-5
Multi-ResNet, no params. added in dec. (ours)	15.7 M	0.0107 ± 9·10-5
Multi-ResNet, saved params. added in dec. v1 (ours)	34.5 M	0.0040 ± 2·10-5
Multi-ResNet, saved params. added in dec. v2 (ours)	34.6 M	0.0093 ± 5·10-5
Shallow water96×192	Residual U-Net	34.5 M	0.1712 ± 0.0005
Multi-ResNet, no params. added in dec. (ours)	15.7 M	0.4899 ± 0.0156
Multi-ResNet, saved params. added in dec. 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 a 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 we train the Multi-ResNet by predicting the next time step in the trajectory only. We do not condition on previous timesteps.
Figure 12: PDE modelling and image segmentation with a Wavelet-encoder Multi-ResNet. Rolled out PDE trajectories (ground-truth, prediction,  $L^2$ -error between ground-truth and prediction) from the Navier-Stokes [top], and the Shallow-Water equation [bottom]. This is the complete version of the truncated Figure 6 [Left] showing further timesteps for the trajectory. Figure and code as modified from [6, Figure 1].
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 4. 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 an FNO model compared to U-Nets. We compare the FNO model to a (Haar wavelet) Multi-ResNets and classical (Haar wavelet) Residual U-Nets with similar number of parameters, on two PDE modelling tasks.
Dataset	Neural architecture	# Params.	r-MSE ↓
Navier-stokes 128 × 128	Residual U-Net	34.5 M	0.0057 ± 2 · 10-5
Multi-ResNet, saved params. added in dec.	34.5 M	0.0040 ± 2 · 10-5
FNO 128-8 mode8 (new)	33.7 M	0.0253 ± 0.0
Shallow water 96 × 192	Residual U-Net	34.5 M	0.1712 ± 0.0005
Multi-ResNet, saved params. added in dec.	34.5 M	0.1493 ± 0.0070
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 as 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 on Shallow water) in the several tables reported in [7].
Dataset	Neural architecture	# Params.	r-MSE ↓
Navier-stokes128 × 128	U-Net 2015 64	31 M	0.01386 ± 0.00004
FNO 128-8 mode8	33.7 M	0.03836 ± 0.00037
Shallow water96 × 192	U-Net 2015 64	31 M	0.1026 ± 0.0161
FNO 128-8 mode8	33.7 M	0.8549 ± 0.0124
Image segmentation. In Figure 13, we present further MRI images with overlayed ground-truth (green) and prediction (blue) masks, augmenting Figure 6 [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.
Figure 13: MRI images from WMH with overlayed ground-truth masks (green) and predictions (blue) obtained from our best-performing Multi-ResNet model.
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 resolution for the following number of iterations/epochs: MNIST [5K, 5K, 5K, 5K] iterations, CIFAR10 [50K, 50K, 50K, 450K] iterations, Navier-Stokes [5, 5, 5, 35] epochs, Shallow water [2, 2, 2, 14] epochs.
Generative modelling with diffusion models. In Figs. 14 and 15, we illustrate samples of a 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 1 with freezing has the advantage that the final model can produce samples on all resolutions, instead of only the intermediate models at 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 metrics measured on the highest resolution. In particular, in Figure 16 we compare samples and FID scores from a 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 1 but without the freezing option produces significantly better samples. This is why we did not further explore the freezing option.
$32 \times 32$  (at  $20\mathrm{K}$  liter.)
MNIST
8  $\times$  8 (at 10 K iter.)
16 x 16 (at 15 K iter.)
4x4 (at 5Kiter.)
Figure 14: Staged training of Multi-ResNets enables multi-resolution training and sampling of diffusion models. We train a diffusion model [8] with a Residual U-Net architecture using Algorithm 1 with four training stages and with freezing on MNIST corresponding to Figure 7. We show samples at the end of each training stage.
PDE modelling. In Table 6, we investigate Residual U-Nets and Multi-ResNets using Algorithm 1 on Navier-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 to 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.
8 x 8 (at 10 K iter.)
16 x 16 (at 15 K iter.)
4x4 (at 5 Kiter.)
Figure 15: Staged training of Multi-ResNets enables multi-resolution training and sampling of diffusion models. We train a diffusion model [8] with a Residual U-Net architecture using Algorithm 1 with four training stages and no freezing on MNIST corresponding to Figure 7. We show samples at the end of each training stage.
Algorithm 1 w/o freezing FID:8.2
Figure 16: A comparison of samples of a DDPM model [8] with a Residual U-Net trained with Algorithm 1 with [Right] and without [Left] freezing.
Algorithm 1 w/ freezing FID:59.8
Table 6: Staged training with Algorithm 1 on Navier-Stokes 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 ↓
Navier-stokes	Residual U-Net	37.8 M	0.0083 ± 0.0034
128 × 128	Multi-ResNet, saved params. added in dec. (ours)	37.8 M	0.0105 ± 0.0102
Shallow water	Residual U-Net	37.7 M	0.4640 ± 0.4545
96 × 192	Multi-ResNet, saved params. added in dec. (ours)	37.7 M	0.7715 ± 0.8327
B.3 Analysis 3: U-Nets encoding topological structure
Experimental details. We begin by discussing MNIST-Triangular, the dataset we used in this experiment. In a nutshell, MNIST-Triangular is constructed by shifting the MNIST  $28 \times 28$  dimensional squared digit into the lower-left half of a  $64 \times 64$  squared image with similar background colour as in MNIST in that lower-left half of triangular shape. In the upper-right half, we use a gray background colour to indicate that the image is supported only on the lower-left of the squared image. We illustrate MNIST-Triangular in Figure 17. We note that MNIST-Triangular has more 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.
Figure 17: Example images from the MNIST-Triangular dataset.
To process the data into the subspaces  $\mathcal{W} = (W_i)$  given in Section 3.3, we need to formulate the 'triangular pixels' shown in Figure 3.3. To do this we use the self-similarity property of the triangular domain to divide the data into a desired resolution. For each image, we 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 we illustrate the coding map at depth two.
Once we have made this encoding map, we are able to map our function values into an array and perform the projections on  $V = W$  by push-forwarding through projection of the Haar wavlets on the triangle to our array through the lexicographic ordering used. For instance, we show in Figure 19
Figure 18: The coding map from the triangular Haar wavelets to their code-space addresses. Such a construction can always be made on a self-similar object with certain separation properties, such as the Open Set Condition [59].
Figure 19: An example of the encoding map taking data from a triangular domain and mapping it to an array of each 'triangular pixel' value under lexicographical ordering.
the array, plotted as an image, revealing that this processing is inherently different to diffusions on a square domain.
This transformation, when defined on the infinite resolution object, is discontinuous on a dense set of  $\triangle$ , yet encodes all of the necessary data to perform a 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 a negligible amount of the function approximation in the diffusion process.
B.4 Ablation studies
In several ablation studies, we further analysed Multi-ResNets 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-ResNets perform a linear transformation in the encoder, which significantly reduces its expressivity, one might hypothesise that the skip connections in a U-Net can be removed entirely. Our results show that skip connections are crucial in U-Nets, particularly in Multi-ResNets.
In Table 7 we compare Residual U-Nets and Multi-ResNets with and without skip connections focusing on PDE Modelling (Navier-stokes 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 if they had been removed. We find that performance significantly deteriorates when using no skip connections. This effect is particularly strong in Multi-ResNets. Multi-ResNets cannot compensate the lack of skip connections by learning the encoder in such a way that representations learned on the lowest resolution space  $V_{0}$  account for not having access to higher-frequency information via the skip connections when approximating the output via  $D_{i}$  in the output space  $W_{i}$  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.