A Multi-Resolution Framework for U-Nets with Applications to Hierarchical VAEs
Authors: Fabian Falck
Christopher Williams
Dominic Danks
George Deligiannidis
Christopher Yau
Christopher C. Holmes
Arnaud Doucet
Matthew Willetts
Venue: NeurIPS 2022 Accept (2022)
Keywords: U-Net
multi-resolution analysis
hierarchical variational autoencoders

Abstract: U-Net architectures are ubiquitous in state-of-the-art deep learning, however their regularisation properties and relationship to wavelets are understudied. In this paper, we formulate a multi-resolution framework which identifies U-Nets as finite-dimensional truncations of models on an infinite-dimensional function space. We provide theoretical results which prove that average pooling corresponds to projection within the space of square-integrable functions and show that U-Nets with average pooling implicitly learn a Haar wavelet basis representation of the data. We then leverage our framework to identify state-of-the-art hierarchical VAEs (HVAEs), which have a U-Net architecture, as a type of two-step forward Euler discretisation of multi-resolution diffusion processes which flow from a point mass, introducing sampling instabilities. We also demonstrate that HVAEs learn a representation of time which allows for improved parameter efficiency through weight-sharing. We use this observation to achieve state-of-the-art HVAE performance with half the number of parameters of existing models, exploiting the properties of our continuous-time formulation.

A Multi-Resolution Framework for U-Nets with Applications to Hierarchical VAEs
Fabian Falck 1,3,4 Christopher Williams *,1 Dominic Danks 2,4 George Deligiannidis
Christopher Yau 1,3,4 Chris Holmes 1,3,4 Arnaud Doucet
Matthew Willetts
University of Oxford University of Birmingham
Health Data Research UK "The Alan Turing Institute
babian.falcke williams, deligian, cholmes, doucet) ox ox ....c................... ac..uk. uk, {ddanks, cyau, mwilletts)@turinng.. uk}
Abstract
U-Net architectures are ubiquitous in state-of-the-ad deep learning, however their regularisation properties and relationship to wavelets understudied. In this paper, we formulate a multi-resolution framework which identifies U-Nets as finite-dimensional truncations of models We provide theoretical results which prove that average pooling corresponds to projection within the space of square-integrable functions and show that U-Nets with average pooling implicitly learn Haar wavelet basis representation of the data. We then leverage our framework to identify state-of-the-art hierarchical VAES (HVAEs), which have U-Net architecture, as type of two-step forward Euler discretisation of multi-resolution diffusion processes which flow from a point mass, introducing sampling instabilities. We also demonstrate that HVAEs learn representation of time which allows for improved parameter efficiency through weight-sharing. We use this observation to achieve state-of-thee-rt HVAE performance with half the number of parameters of existing models, exploiting the properties of our continuous-ting formulation.
infiite-dimensional function space.
Introduction
U-Net architectures are extensively utilised in modern deep learning models. First developed for
image segmentation in biomedical applications [1], U-Nets have been widely applied for text-to-
image models [2]. image-to-image translation [3]. image restoration [4, 5], super-resolution [6],
and multiview learning [7], amongst other tasks [8]. They also form core building block as the neural architecture of choice in tf-te-of-the-art generative models, particularly for images, such as HVAEs [9, 10, 11, 12] and diffusion models [2, 13, 14, 15, 16, 17. 18, 19, 20J. In spite of their empirical success, it is poorly understood why U-Nets work well, and what regularisation they impose.
In likelihood- based generative modelling, various model classes are competing for superiority, includ- ing normalizing flows [21, 22]. autoregressive models [23, 24]. diffusion models, and hierarchical variational autoencoders (HVAEs), the latter two of which we focus on in this work. HVAEs form groups of latent variables with conditional dependence structure, use U-Net neural architecture, and are trained with the typical VAE ELBO objective (for a detailed introduction to HVAEs, see Appendix B). HVAES show impressive synthesis results on facial images, and yield competitive like- lihood performance, consistently outperforming the previously of-te-of-theart autoregressive models, VAEs and flow models on computer vision benchmarks [9, 10J. HVAEs have undergone a journey
Equual contribution.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Standard basis elements combined by average pooking
average pooling
Σ
e1
C2
C3
04
Σ
0
C4
projection
x 4image graphed as piacowise constant function
00
09
06
01
Resuttngg 2imago
0.
05
Hear wavelet basis elements by projection
Figure 1: U-Nets with average pooling learn a Haar wavelet basis representation of the data.
of design iterations and architectural improvements recent years, for example the introduction a deterministic backbone [25, 26, 27] and ResNet elements [28, 29] with shared parameters between the inference and generative model parts. There has also been a massive increase in the number of latent variables and overall stochastic depth, well the use of different types of residual cells in the decoder [9, 10] (see $4 and Fig. A.1 for detailed discussion). However, a theoretical understanding of these choices is lacking. For instance, it has been shown why residual backbone may be beneficial, or what the specific cell structures in VDVAE [9] and NVAE [10] correspond to, or how they could be improved.
In this paper we provide a theoretical framework for understanding the latent spaces in U-Nets, and apply this to HVAEs specifically. Doing so allows to relate HVAEs to diffusson processes, and also to motivate a new type of piecewise time- mame-homogeenous model which demonstrates state-of-the-ar performance with approximately half the number of parameters of a VDVAE [9]. More formally, (a) We provide a multi- resolutionn framework for U-Nets. We formally define U-Nets as acting over a multi-resolution hierarchy of L (2 1)3). We prove that average pooling is a conjugate operation to projection in the Haar wavelet basis within ([0, 1 We use this insight to show how U-Nets with average pooling implicitly learn Haar wavelet basis representation of the data (see Fig. 1), helping to characterise the regularisation within U-Nets. (b) We apply this framework to state-of-the-an HVAEs as an example, identifying their residual cell structure of two-step forward Euler discretisation of a multi-resolution diffusion bridge. We uncover that this diffusion process flows from a point mass, which causes instabilities, for instance during sampling, and identify parameter redundancies through our continuous time formulation. Our framework both allows us understand the heuristic choices of existing work in HVAEs and enables future work to optimise their design, for instance their residual cell. (c) In our experiments, we demonstrate these sampling instabilities and train HVAEs with the largest stochastic depth ever, achieving state-of- art performance with half the number of parameters by exploiting our theoretical insights. We explain these results by uncovering that HVAEs secretly represent time in their state and show that they use this information during training. We finally provide extensive ablation studies which, for instance, rule out other potential factors which correlate with stochastic depth, show the empirical gain of multiple resolutions, and find that Fourier features (which discrete-time diffusion models strongly benefit from [19]) do not improve performance in the HVAE setting.
contributions are as follows:
type
2 The Multi-Resolutinn Framework
A grayscale image with infinite resolution be thought of as the graph² of two-dimensional function over the unit square. To store these infinitely-detalled images in computers, we project them to some finite resolution. These projections support over the unit square, but they are piecewise constant on finitely many intervals or 'pixels', e.g. 512" pixels, and we store the function values obtained at these pixels in an array or 'grid'. The relationship between the finite-dimensional version and its infinitely-fine counterpart depends entirely on how we construct this projection preserve the details we wish keep. One approach is to prioritise preserving the large-scale details of our images, so unless closely inspected, the projection
still be thought of the graphs of functions with
For a function f(: (.), its graph is the set U. (z, f(x)).

---

is indistinguishable from the original. This can be achieved with multtiresoluuion projection [30] of the image. In this section we introduce multi-resolution framework for constructing neural network architecttres that utilise such projections, prove what regularisation properties they impose, and show as an example how HVAEs with a U-Net [1] architecture can be interpreted in our framework. Proofs of all theorems in the form of an extended exposition of our framework can be found in Appendix A.
2.1 Multi-Resolution Framework: Definitions and Intuition
What makes a multi-resolution projection good at prioritising large-scale details can be informally explained through the following thought experiment. Imagine we have an image, represented as the graph of function, and its finite-dimensional projection drawn on the wall. We look at the wall, start walking away from it and stop when the image and its projection are indistinguishable by eye. The number of steps we took away from the wall can be considered our measure of 'how far away' the approximation is from the underlying function. The goal of the multi-resolution projection is therefore to have to take as few steps away as possible. The reader is encouraged to physically conduct this experiment with the examples provided in Appendix 1.1. We can formalise the aforementioned intuitions by defining
multi-resolution hierarchy [30] of sub-spaces we may project to:
Definition 1. [Daubechies ( (1992) [30]] Given nested sequence of approximation spaces C {V-j is a multi-resolution hierarchy of the function space (RRR") if: (A1) U jez V-, = (R"):
V C Vo- C V-1 C.
AEEE V-j ={0] (A3) f()) E V-j & f(2i.) Vo; (A4) f() E Vo + f(. n) E Vo for n E Z. For compact set X C R a multi-resolutte hierarchy of (X)) is (V-jbez as defined above, restricting functions in V-j to be supported on X.
): (A2)
In Definition! 1, the index references how many steps we took in our thought experiment, so negative j corresponds to 'zooming in' on the images. The original image' is a member of L L², 11²), space of square-integrable functions on the unit square, and its finite projection to 2 2 many pixels is a member of V-j. Images can be represented as piecewise continuous functions in the subspaces V-j- = {f E L ( 1J) f,
), the
= Ck. k E {O.... 2' 1}, C₁ E R]. The nesting property V-5+1 C V-j ensures that any image with (2) (2 pixels can also be represented by (2))² pixels, but at higher resolution. Assumption (A1) states that with infinitely many pixels. we can describe any infinitely detailed image. In contrast, (A2) says that with no pixels, we cannot approximate any images. Assumptions (A3) and (A4) allow us to form a basis for images V-j if we know the basis of Vo. One basis made by extrapolating from Vo in this way is known as wavelet basis [30]. Wavelets have proven useful for representing images, for instance in the JPEG standard [31], and are constructed to be orthonormal.
any
Now suppose we have a probability measure Vs over infinitely detailed images represented in [²0 1 and wish to represent it at a lower resolution. Similar to how we did for infinitely detailed images, we want to project the measure Voc to a lower dimensional measure v/ on the finite dimensional space V-y. In extension to this, we want the ability to reverse this projection so that we may sample from the lower dimensional measure and create a generative model for Mse. We would like to again prioritise the presence of large-scale features of the original image within the lower dimensional samples. We do this by constructing a multiiresolution bridge from V to Vj, as defined below.
Definition 2. Let X C m be compact, {V-j)se be a multi-resolution hierarchy of scaled so (x( = U,eN, V-, and Vo {0}. If D(L2 is the space of probability measures over L (X). then a family of probability measures {1t (etoe on (x(X) is a multi-resolution bridge if:
ii) there exist increasing times I = (t, her, where to 0, lim to 1, such that s Ity, 11) implics supp(v,) C V i. v, E D(V_)); and, (ii) for E (0, 1), the mapping
v, is continuous for (ty, tty1) for some j.
The continuous time dependence in Definition 2 plays a movie of the measure vo supported on Vo growing to 1/... a measure on images with infinite resolution. At time interval [ty, fz-).. the space V-, which the measure is supported on is fixed. We may therefore define a finite-dimensional model
'We here focus on grayscale, squared images for simplicity, but note that our framework can be seamlessly extended to colour images with a Cartesian product L L'( 11²) x (10 113) x <2((0, 1 and other continuous signals such time series.
3

---

transporting probability measures within V-j, but at f the support flows over to V-5-1- Given multi-resolution hierarchy, we may glue these finite models, each acting on disjoint time interval, together in unified fashion. In Theorem 1 we show this for the example of continuous-ting multi-resolution diffusion process truncated up until some time f. = E (0, 1) and in the standard basis discussed in $2.2, which will be useful when viewing HVAEs as discretisations of diffusion processes on functions in $2.3.
Theorem 1. Let Bj (1111++) (V(--)) (D(=))) be a linear operator (such as a diffusion transition kernel, see Appendix A) for J with coefficients 1 g 1b5 (152))41) V-, H V-j- and define the natural extensions within V-J in bold, i.e. Bj == B, 1 Iv Then the operator B [0 x D(V_J) + D(V_)) and the coefficients μ. a: [0, V-J V-J given by
B Σ 10) Bj, 1 E (15,011) 10) o Σ (0.6111) 000)
j=0
induce a multi-resolution bridge of measures from the dynamics for t [0, T) and the standard basis as dZ, = pr(Zz ((ZZz)dt++ (o,+++ (see Appendix A.4 for details) for Ze E V V E [t,, 1g+1), i.e. a multi-resolution diffusion process.
The concept of multi-resolution bridge will become important in Section 2.2 where we will show that current U-Net bottleneck struc-
V
V
V-I
F
A
tures used for unconditional sampling impose a multi-resolution bridge on the modelled densities. To preface this, we here pro- vide a description of U-Net within our framework, illustrated
V
V
V-j
in Consider B, F).0 (VV=)) + (VV-;) as the forwards and backwards passes of U-Net resolution j. Further, let
AB
A
B
je0
1. Y
Skip
P-j+1 : DVV ((- and F. D(V_++1) - (D()))
V-j
V
connection
be the projection (here: average pooling) and embedding maps (e.g. interpolation), respectively. When using an j.-rrecostrruction er- ror, U-Net [1] architecture implicitly learns a sequence of models
A E-i
E-j
P P-j+1 V
V j+1
V
V
j+1
j+1
B, : (D(V=++1) x D(V + (V--;) due to the orthogonal de- composition Vー,- = V-5+1 1-- where U- = VV3 j+1' The backwards operator for the U-Net has (bottleneck) input from D(V-j+1) and a (skip) input yielding information from D(V A simple bottleneck map U3,0 D(V_)) + D(V-,) (without skip connection) given by
A
VO
Vo
Bend/
Vo
Bottleneck
A U-Net in multi-resolution framework. See Appendix B.2 for details.
Figure 2:
Uz.0 B3, F-j P-j+1 o F2.0,
(1)
and a U-Net bottleneck with skip connection is
Uj = -,, P-j+1 Fj,0, FjØ).
(2)
In HVAEs, the map U, DVV-,) + (VV))) is trained to be the identity by minimising reconstruc- tion error, and further shall approximate Ug0 Uj via KL. divergence. The I.errecontt error for U has an orthogonal partition of the inputs from V-j+1 x V-j, hence the only new subspace added U- As each orthogonal U_j+1 added sequentially in HVAEs, the skip connections induce a multi-resolution structure of this hierarchical neural network structure. What we will investigate in Theorem is the regularisation imposed on this partitioning by enforcing US R U j. as is often enforced for generative models with VAEs.
2.2 The regularisation property imposed by U-Net architectures with average pooling
Having defined U-Net architectures within our multi-resolltion framework, we are now interested in the regularisation they impose. We do so by analysing a U-Net when skip connections are absent, so that we may better understand what information is transferred through each skip connection when they are present. In practice, a pixel representation of images is used when training U-Nets, which we henceforth call the standard busis (see A.2, Eq. (A.9)). The standard basis is not convenient to derive theoretical results. It is instead preferable to use basis natural to the multi-resolution bridge imposed by U-Net with corresponding projection operation, which for average pooling the Haar (wavelet) basis [32] (see Appendix A.2). The Haar basis, like Fourier basis, is orthonormal basis

---

of (X)) which has desirable aapproximatioon properties, We formalise this in Theorem 2 which states that the dimension reduction operation of average pooling in the standard basis is conjugate operation to co-ordinate projection within the Haar basis (details are provided in Appendix A.2).
Theorem 2. Given V-j as in Definition 1, let x E v be represented in the standard basis E, and Haar basis y Let x E + , be the change of basis map illustrated in Fig. 3, then we have the conjugacy T- o pool j--++1 = -roojv-.1 o T..
pool 3--41
Theorem 2 means that if we project an image from V-j to V-j+1 the Haar wavelet basis, we can alternatively view this as changing to the standard basis via 7. per- forming average pooling, and reverting back via ; (see Figure 3). This is important because the Haar basis is orthonormal, which in Theorem 3 allows to precisely quantify what information is lost with average pooling.
(V Ej)
(V-++1. J+1, EG-1)
15-
(V v,)
(V V 3-1)
proiy
Figure 3: The function space V-j re- mains the same, but the basis changes under T..
Theorem 3. Let {V-))} be multi- hierar- chy of V_J where V-j = V-j+1 + U-J+1 and further. let F3, BB. D(V-j) + D(V-;) be such that B BysF I with parameters o and B. Define o == Fn o... F2 by Fs D(V-)) + D(V-5+1) where Fj == projv 512 o Fj and analogously define Bilii. o with B, == B3 Beodeeyyy_g. Then, the sequence (B, j (.o((FFF,,oo (Fij,,85))) forms a discrete multi- resolution bridge between FilsoOSs s and B1J1, FilJ,ons y at times {1,))=v and
Σ Ex,y proiu.... 1=0
(3)
where W2 is the Wasserstein-2 metric and | ll is the Lipschitz constant of Fjj.. Theorem 3 states that the bottleneck component of U-Net pushes the latent data distribution to finite multi-resolution basis, specifically a Haar basis when average pooling is used. To see this, note that the RHS of Eq. (A.65) itself upper-bounded by the /r=reconstrocttinn error. This is because the Wasserstein-2 distance finds the infinimum over all possible couplings between the data and the 'reconstruction" measure, hence any coupling (induced by the learned model) bounds it. Note that models using a U-Net, for instance HVAEs or diffusion models, either directly or indirectly optimise for low reconstruction error in their loss function. The LHS of Eq. (A.65) represents what percentage of our data enters the orthogonal subspaces {U- which are (by Theorem 2) discarded by the bottleneck structure when using U-Net architecture with average pooling. Theorem thus shows that as we minimise the reconstruction error during training, we minimise the percentage of our data transported to the orthogonal sub-spaces {U-}}/-0- Consequently, the bottleneck architecture implicitly decomposes our data into Haar wavelet decomposition. and when the skip connections are absent (like traditional auto-encoder) our network learns to compress the discarded subspaces U-j- This characterises the regularisation imposed by U-Net in the absence of skip connections.
These results suggest that U-Nets with average pooling provide a direct alternative to Fourier features [19, 33, 34, 35] which impose a Fourier basis, an alternative orthogonal basis on L'2(X), as with skip connections the U-Net adds each subspace U-j sequentially. However, unlike Fourier bases, there are in fact a multitude of wavelet bases which are all encompassed by the multi-resolution framework, and in particular, Theorem 3 pertains all of them for the bottleneck structure. This opens the door to exploring conjugacy operations beyond average pooling induced by other wavelet bases optimised for specific data types.
2.3 Example: HVAEs Diffusion Discretisations
To show what practical inferences we can derive from our multi-resolution framework, we apply it to analyse of HVAE architectures (see Appendix B.3 for an introduction), identifying parameter redundancies and instabilities. Here and in our experiments, we focus on VDVAEs [9]. We provide similar results for Markovian HVAEs [36, 37] and NVAES [10] (see $ 4) in Appendix A.5.
We start by inspecting VDVAEs. As we show next, we can tie the computations VDVAE cells to the (forward and backward) operators Fia and B.g within our framework and identify them as type of two-step forward Fuler discretisation of diffusion process. When used with U-Net, as is done in VDVAE [9], this creates multi-resolution diffusion bridge by Theorem 4.
5

---

Theorem 4. Lerty = E 1) and consider (the pa backward pass) Bo.i D(V_)) + (V6) given in multi-resolution Markov process the standard basis:
dZ, (..)) + To (Z))d + (Z(.)d..
(4)
where proju_, Z,, =0 0, |Z > 1z ll2 with 0 S s <ts T and for a measure vJ E (VV_J) we have Xr. Zo- ForjiyJ = 6(0). Then, VDVAEs approximates this process, and its residual cells are a rype of two-step forward Euler discretisation of this Stochastic Differential Equation (SDE).
To better understand Theorem 4, we visualise its residual cell structure of VDVAEs and the cor- responding discretisation steps in Fig. 4, and together those of NVAEs and Markovian HVAEs in Appendix A.5, Fig. A.1. Note that this process is Markov and increasing in the Z4 variables. Similar processes have been empirically observed as efficient first-order approximates to higher-order chains, for example the memory state LSTMs [38]. Further, VDVAEs and NVAEs are even claimed to be high-order chains (see Eqs. (2,3) in [9] and Eq. (1) in [10]), despite only approximating this with a accumulative process.
To show how VDVAES impose the growth of the Zz we prove that the bottleneck component of VDVAE's U-Net enforces Zo = 0. This is done by identifying that the measure Vo, which a VDVAE connects to the data Voc via a multi- resolution bridge, is a point mass on the zero function. Consequently the backward pass must grow from this, and the network learns this in monotonic manner we later confirm in our experiments (see $3.2).
4;
Zi 2, 2,
7
Pol 7o(2112)
X,
9o (Z JZ X.,
V choose Po 90
2 Z-,- + (1,12))
2 +x
+5(Z,AW,
sample 2,5
Theorem 5. Consider the SDE in Eq. (A.76). trained through the ELBO in Eq. B.101. Let
S (2
ūy denote the data measure and vo dtog be the initial multi-resolution bridge measure im-
- - 21 -0--
2.7 O
posed by VDVAEs. qo.j and po.j are the densi- ties of B 1 Fyjăj and 15, respectively. then a VDVAE optimises the boundary condition ming.s K LLg.,,,1ll900002 dex indicates the joint distribution.
Z,11 - 2
Zi+z
Figure 4: The VDVAE [9] cell is type of two-step forward Euler discretisations of the continuous- time diffusion process in Eq. A.76. See Fig. A.1 for similar schemas on NVAE [10] and Markovian HVAE [36, 37].
where double in-
Theorem 5 states that the VDVAE architecture forms multi-resolution bridge with the dynamics of Eq. (A.76), and connects our data distribution to the trivial measure on Vj: a Dirac mass at 0 as the pooling here cascades completely to Vo- From this insight, we can draw conclusions on instabilities and on parameter redundancies of this HVAE cell. There are two major instabilities in this discretisation. First, the imposed ио is disastrously unstable as enforces a data set, with potentially complicated topology to derive from paint-mass in U-; at each ( Lg. and we observe the resulting sampling instability in our experiments in $3.3. We note that similar arguments are applicable in settings without latent hierarchy imposed by a U-Net, see for instance [39]. The
VDVAE architecture does, however, bolster this rate through the z term, which is absent in NVAES [10], in the discretisation steps of the residual cell. We empirically observe this controlled backward error in Fig. 6 [Right]. We refer to Fig. A. I for a detailed comparison of HVAE cells and their corresponding discretisation of the coupled SDE in Eq. (A.76).
Moreover, the current form of VDVAEs is paramttreeeee and not informed by this continuous- time formulation. The continuous time analogue of VDVAEs [9] in Theorem 4 has time dependent coefficients Y1 v.1, " 1,2, T t. We hypothesise that the increasing diffusion process in Z; implicitly encodes time. Hence, explicitly representing this in the model, for instance via ResNet blocks with independent parameterisations at every time step, redundant, and homogenooss model (see Appendix A.6 for a precise formulation) practically speaking, performing weight-sharing across time time steps/layers-
the same expressivity, but requires far fewer parameters than the state-of-the-art VDVAE. It is worth noting that such time-homogeneous model would make the parameterisation of HVAES more similar to the recently popular -csee- diffusion models [40, 41] which perform weight-sharing across all time steps.

---

KOWA 00
WOWEE
IN3 VENAR 17
WS- 100AE
325
220-
434
2022 600% 020c
E0O%
1000k
203e
40CK
5008
DOOK
1052e 32008 3400k
Figure 5: A small-scale study parameter efficiency of HVAEs. We compare models with with 1,2,3 and parameterised blocks per resolution (fx1, x:2, 2:3, x44) against models with a single parameterised block per resolution weight-shared {2, 5, 10, 20} times ((r2, r3, r5, r5,, We report NLL (4) measured on the validation set of MNIST [left] and CIFARIO [right]. NLL performance increases with more weight-sharing repetitions and surpasses models without weight- sharing but with more parameters.
Experiments
In the following we probe the the-
oretical understanding of HVAEs gained through our framework, demonstrating its utility in four experimental analyses: (a) Im- proving parameter efficiency HVAEs, (b) Time representation in HVAEs and how they make use of it. (c) Sampling instabilities in HVAEs, and (d) Ablation studies.
Table 1: A large-scale study of parameter efficiency in HVAEs. We compare our runs of VDVAE with original hyperparameters [9] (VDVAE") against While WS-VDVAEs have improved parameter efficiency by fac- tor of 2, they reach similar NLL VDVAE* with the simple modification inspired by our framework (weight sharing). We note that parameter count cannot be provided for VDM [19] as the code is not public and the manuscript does not specify it.
weight-shared VDVAE (WS-VDVAE).
We train HVAEs using VD- VAE [9] as the basis model on five datasets: MNIST [42], CI-
Dataset
Method
Type
#Params
NLL 1
MNIST
WS-VDVAE (ours)
VAE
232k
79.98
28x28
VDVAE" (ours)
VAE
339k
S 80.14
FARIO [43], two downsampled versions of ImageNet [44, 45].
NVAE [10]
VAE
33m
S 78.01
and CelebA [46), splitting each
WS-VDVAE (ours)
VAE
25m
2.88
into a training, validation and
CIFAR10
WS-VDVAE (ours)
VAE
39m
< 2.83
test set (see Appendix D for de-
32 x 32
VDVAE" (ours)
VAE
39m
S 2.87
tails). In general, reported nu-
NVAE [10]
VAE
131m
S 2.91
meric values refer to Negative Log-Likelihood (NLL) in nats (MNIST) or bits per dim (all
VDVAE [9]
VAE
39m
2.87
VDM [19]
Diff
2.65
ImageNet
WS-VDVAE (ours)
VAE
55m
3.68
other datasets) on the test set
WS-VDVAE (ours)
VAE
85m
< 3.65
32 x 32
at model convergence, if not
VDVAE" (ours)
VAE
119m
3.67
stated otherwise. We note that
NVAE [10]
VAE
268m
3.92
performance on the validation
VDVAE [9]
VAE
119m
< 3.80
and test set have similar trends
VDM [19]
Diff
3.72
in general. An optional gra- dient checkpointing implementa- tion to trade in GPU memory for compute is discussed in Ap-
CelebA
64 x 64
WS-VDVAE (ours)
VAE
75m
2.02
VDVAE" (ours)
VAE
125m
< 2.02
NVAE [10]
VAE
153m
< 2.03
pendix F. Appendices F and G
define the HVAE models we train,
i.e. po(z,,) Pa [2:) (2)2))), (x), qs (2 22,,, (2 x) and po(x and present additional experimental details and results. We provide our PyTorc code base http://iwww...ooaaaaaaaaaaaaeeee (see Appendix C for details).

---

200
[A4l wht supput wwt..ttt -ww.ttmmmmm.mm 1pppppp0000 bock
1003
250
Dk inpu blokk
200
2000
-003
200
1000
50-
1.0
20
25
Do0 bocck
Top-doem brock
Figure 6: HVAEs secretly represent a notion of time: We measure the L2-norm of the residual state for the [Left] forward/bottom-up up pass and the [Right] backward/top-dow pass over 10 batches with 100 data points each. In both plots, the thick, central line refers to the average and the thin, outer lines refer to =2 standard deviations.
3.1 "More from less": Improving parameter efficiency in HVAEs
In $2.3, we hypothesised that a time-homogeneous model has the same expressivity as a model with dime-dependent coefficients, yet uses much less parameters. We start demonstrating this effect by weight-sharing ResNet blocks across time on a small scale. In Fig. 5, we train HVAEs on MNIST and CIFARIO with 2,, 4) ResNet blocks (referred resolution with spatial dimensions {322, 162, 22, 42, }2 (VDVAE"), ) and compare their performance when weight-sharing a single parameterised block per resolution 3, 5, 10, 203 times (referred to as 1r2.r3,55.110,222 J; WS-VDVAE) excluding projection and embedding blocks. As hypothesised by our framework, yet very surprising in HVAEs, NLL after Im iterations measured on the validation set gradually increases the more often blocks are repeated even though all weight-sharing models have an identical parameter count the x1 model (MNIST: 107k, CIFAR10: 8.7m). Furthermore, the weight-sharing models often outperform or reach equal NLLs compared to x2, x3, x4, all of which have more parameters (MNIST: 140k: 173k; 206k. CIFARIO: 13.0m; 17.3m; 21.6m), yet fewer activations, latent variables, and number of timesteps at which the coupled SDE in Eq. (A.76) discretised.
[x1, x2, x3, x4)) in each
We now scale these findings up to large-scale hyperparameter configurations. We train VDVAE closely following the state-of-the-art hyperparameter configurations in [9], specifically with the same number of parameterised blocks and without weight-sharing (VDVAE"), and compare them against models with weight-sharing (WS-VDVAE) and fewer parameters, i.e. fewer parameterised blocks, in Table 1. On all four datasets, the weight-shared models achieve similar NLLs with fewer parameters compared to their counterparts without weight-shariee sharrng:: We use 32%, 36%, 54%, and 40% less parameters on the four datasets reported in Table 1, respectively. For the larger runs, weight-sharing has diminishing returns on NLL as these already have many discretisation steps. To the best of our knowledge, our models achieve a new state-of-the-art performance in terms of NIL compared to any HVAE on CIFAR1O. 10, Ie33eNet32 and CelebA. Furthermore, our WS-VI -S-VDVAE models have stochastic depths of 57, 105, 235, 125, respectively, the highest ever trained. In spite of these results, it is worth noting that current HVAEs, and VDVAE in particular remains notoriously unstable to train, partly due to the instabilities identified in Theorem 5, and finding the right hyperparameters helps, but cannot solve this.
3.2 HVAES secretly represent time and make use of it
In $3.1, we showed how we can exploit insight on HVAEs through our framework make HVAEs more parameter efficient. We now want to explain and understand this behavior further. In Fig. 6, we measure |Z the Ln-morm of the residual state at every backward/top-don over several batches for models trained on MNIST (see Appendix G.2 for the corresponding figure of the forward/bottom-uup pass, and similar results on CIFARIO 10 and ImageNet32). On average, we experience an increase in the state norm across time in every resolution, interleaved by discontinuous 'jumps' at the resolution transitions (projection or embedding) where the dimension of the residual state changes. This supports our claim in $2 that HVAEs discretise multi-resolution diffusion
block with index i,
8

---

4
9
8
8
6
3
CIFAR1O (t 0.9)
ImageNet32 (t = 1)
ImageNet64 (t 0.9)
MNIST (t = 0.9)
CelebA (t 0.5)
Figure 7:
Unconditional samples (not cherry- pickry) of VDVAE". While samples on MNIST and
CelebA demonstrate high fidelity and diversity, samples CIFARIO, ImageNet32 and ImageNet64
are diverse, but are unrecognisable, demonstrating the instabilities identified by Theorem 1. Tempera- tures / tuned for maximum fidelity.
processes which are increasing in the Z, variables, and hence learn to represent a notion of time in their residual state.
is now straightforward to ask how HVAEs benefit from this time representation during training: As we show in Table 2, when normalising the state by its norm at every forward and backward block during training, i.c. forcing "flat line" in Fig. 6 [Left], learning deteriorates after short while, resulting in poor NLL results compared to the runs with a regular, non-normalised residual state. This evidence confirms our earlier stated hypothesis: The time representation in ResNet-based HVAEs encodes information which recent HVAEs heavily rely on during learning.
3.3 Sampling instabilities in HVAES
High fidelity unconditional samples of faces, e.g. from models trained on CelebA, cover the front pages of state-of-the-ar HVAE papers [9, 10]. Here, we question whether face datasets are an appropriate benchmark for HVAEs. In Theorem 5, we identified the aforementioned state-of-the-art HVAEs flow from a point mass, hypothesising instabilities during sampling. And indeed, when sampling from our trained VDVAE" with state-of-the-art configura- tions, we observe high fidelity and diversity samples on MNIST and CclebA, but unrecognisable, yet diverse samples CIFAR1O, ImageNet32 and ImageNet64, in spite of state-of-the-as test set NLLs (see Fig. and Appendix G.3). We argue that MNIST and CelebA, i.e. numbers and faces, have a more uni-modal nature, and in this sense easier to learn for discretised multi-resolution process flowing to point mass, which is uni-modal, than the other "in-the-wild", multi-modal datasets. Trying to approximate the latter with the, in this case unsuitable, HVAE model leads to the sampling instabilities observed.
Table 2: NLL of HVAEs with and without normalisation of the residual state Z4.
Residual state
NLL
MNIST Normalised (X)
< 464.68
Non-normalised
S 81.69
CIFARIO
Normalised (X)
6.80
Non-normalised
2.93
ImageNet
Normalised
6.76
Non- Non-normalised
3..88
3.4 Ablation studies
We conducted several ablation studies which support our experimental results and further probe our multi-resolution framework for HVAFs. In this section we note key findings-a detailed account of all ablations can be found in Appendix G.4. In particular, we find that the number of latent variables. which correlates with stochastic depth, does not explain the performance observed in $3.1, supporting our claims. We further show that Fourier features do not provide performance gain in HVAEs, in contrast to tfate-of-the-ar diffusion models, where they significantly improve performance [19]. This consistent with our framework's 's finding that U-Net architecture with pooling is already forced to learn a Haar wavelet basis representation of the data, hence introducing another basis does not add value. We also demonstrate that residual cells are crucial for the performance of HVAES as they are able to approximate the dynamics of diffusion process and impose an SDE structure into the model, empirically compare rulti-resolution bridge to a single-resolution model, and investigate synchronous asynchronous processing time between the forward and backward pass.

---

4 Related work
U-Nets. A U-Net [1] is an autoencoding architecture with multiple resolutions where skip con- nections enable information to pass between matched layers on opposite sides of the autoencoder's bottleneck. These connections also smooth out the network's 's loss landscape [47]. In the litera- ture, U-Nets tend to be convolutional, and wide range of different approaches have been used for up-sampling and down-sampling between resolutions, with many using average pooling for the down-sampling operation [13, 14, 16, 17, 19]. In this work, we focus on U-Nets operators on measures interleaved by average pooling as the down-sampling operation (and a corresponding inclu- sion operation for up-sampling), and formally characterise U-Nets in Section 2.1 and Appendix B.2. Prior to our work, the dimensionality-reduucii acknowledged as being useful, however it was unclear what regularising properties U-Net imposes. We provided these in $2.
bottleneck structure of U-Nets was widely
The evolution of HVAEs can be seen as a quest for parameterisation with more expres- siveness than single-latent-layer VAEs [48], while achieving stable training dynamics that avoid common issues such as posterior collapse [36, 49] or exploding gradients. Early HVAEs such as LVAE condition each latent variable directly on only the previous one by taking samples forward [36, 37]. Such VAFs suffer from stability issues even for very small stochastic depths. Nouveau VAEs (NVAE) [10] and Very Deep VAFs (VDVAE) [9] combine the improvements of several earlier HVAE models (see Appendix B for details), while scaling up larger stochastic depths. Both use ResNet-based backbones, sharing parameters between the generative and recognition parts of the model. VDVAE is the considerably simpler approach, in particular avoiding common tricks such as warm-up deterministic autoencoder training phase or data-specific initialisation. VDVAE achieves stochastic depth of up to 78, improving performance with more ResNet blocks. Worth noting is that while LVAE and NVAE use convolutions with appropriately chosen stride to jump between resolutions, VDVAE average pooling. In all HVAES to date, a theoretical underpinning which explains architectural choices, for instance the choice of residual cell, is missing, and we provided this in Section $2.3.
HVAEs.
5 Conclusion
In this work, we introduced a multi-resolution framework for U-Nets. We provided theoretical results which uncover the regularisation property the U-Nets bottleneck architecture with average pooling as implicitly learning Haar wavelet representation of the data. We applied our framework to HVAEs, identifying them as multi-resolution diffusion processes flowing to a point mass. We characterised their backward cell as type of two-step forward Euler discretisations, providing an alternative to score-matching to approoximate continuous-timg diffusion process [16, 18), and observed parameter redundancies and instabilities. We verified the latter theoretical insights in both small- and large-scale experiments, and in doing so trained the deepest ever HVABs. We explained these results by showing that HVAEs learn representation of time and performed extensive ablation studies.
An important limitation is that the proven regularisation property of U-Nets limited to using average pooling as the down-sampling operation. Another limitation is that we only applied our framework to HVAEs, though is possible to apply to other model classes. It could also be argued that the lack of exhaustive hyperparameter optimisation performed is limitation of the work as it may be possible to obtain improved results. We demonstrate, however, that simply adding weight-sharing to the hyperparameter settings given in the original VDVAE paper [9] leads to state-of-the-art performance with improved parameter efficiency, and hence view it as a strength of our results.
10

---

Acknowledgments and Disclosure of Funding
Fabian Falck acknowledges the receipt of studentship awards from the Health Data Research UK-The Alan Turing Institute Wellcome PhD Programme in Health Data Science (Grant Ref: 218529/2/19/22, and the Enrichment Scheme of The Alan Turing Institute under the EPSRC Grant EP/N5 10129/1. Chris Williams acknowledges support from the Defence Science and Technology (DST) Group and from ESPRC DTP Studentship. Dominic Danks is supported by Doctoral Studentship from The Alan Turing Institute under the EPSRC Grant EP/N5101299
Christopher Yau is funded by UKRI Turing AI Fellowship (Ref: EP/VO23233/1), Chris Holmes acknowledges support from the Medical Research Council Programme Leaders award MC UP A_ 107. 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/RO136 3616/11. This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative. Arnaud Doucet also acknowledges support from the EPSRC grant EP/R0347100 10/1. Matthew Willetts is grateful for the support of UCL Computer Science and The Alan Turing Institute.
MC_UP_L
The authors report no competing interests.
The three compute clusters used in this work were provided by the Alan Turing Institute, the Oxford Biomedical Research Computing (BMRC) facility, and the Baskerville Tier HPC service (ttps://www.baseevville.... www.bbavilllle.cc...k in part through computational resources provided by The Alan Turing Institute under EPSRC grant EP/N510129/ and with the help of generous gift from Microsoft Corporation. Second, used the Oxford BMRC facility, a joint development between the Wellcome Centre for Human Genetics and the Big Data Institute supported by Health Data Research UK and the NIHR Oxford Biomedical Research Centre. The views expressed are those of the author(s) and not necessarily those of the NHS, the NIHR or the Department of Health. Third, Baskerville was funded by the EPSRC and UKRI through the World Class Labs scheme (EP/TO22221/1) and the Digital Research Infrastructure programme (EP/W032244/1) and is operated by Advanced Research Computing at the University of Birmingham.
which we detail in the following. First, this research was supported
We thank Tomas Lazauskas, Jim Madge and Oscar Giles from the Alan Turing Institute's Research Engineering team for their help and support. We thank Adam Huffman, Jonathan Diprose, Geoffrey Ferrari and Colin Freeman from the Biomedical Research Computing team the University of Oxford for their help and support. We thank Haoting Zhang (University of Cambridge) for valuable comments the implementation; Huiyu Wang (Johns Hopkins University) for useful discussion on gradient checkpointing; and Ruining Li and Hanwen Zhu (University of Oxford) for kindly proofreading the manuscript.
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] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint MxXiv:2205.11483 2022.
[3] Phillip Isola, Jun- Yan Zhu, Tinghui Zhou, and Alexei A. Efros. Image-to-image translation with conditional adversarial networks. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), July 2017.
[4] Xiaojiao Mao, Chunhua Shen, and Yu-Bin Yang. Image restoration using very deep convo- lutional encoder-decoder networks with symmetric skip connections. Adyances in Neural Information Processing Systems, volume 29, 2016.

---

[5] Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, Ming- Hsuan Yang, and Ling Shao. Multi-stage progressive image restoration. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2021.
[6] Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Amaud Doucet. Conditional simulation using diffusion schrödinger bridges. arXiv preprint arXiv:2202.13466, 2022.
[7] Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In Computer Vision 2020. Springer International Publishing, 2020.
[8] Zoe Landgraf, Fabian Falck, Michael Bloesch, Stefan Leutenegger, and Andrew J. Davison. Comparing view-based and map-based semantic labelling in real-time slam. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 6884-6890, 2020.
[9] Rewon Child. Very Deep VAES Generalize Autoregressive Models and Can Outperform Them on Images. In International Conference on Learning Representations, 2021.
[10] Arash Vahdat and Jan Kautz. NVAE: deep hierarchical variational autoencoder. In Advances in Neural Information Processing Systems, volume 33, 2020.
[11] Louay Hazami, Rayhane Mama, and Ragavan Thurairatnam. Efficient-vdvac: Less is more.
arXiv preprint arXiv:2203.13733 2022.
[12] Simon AA Kohl, Bernardino Romera-Parede Paredes, Klaus H Maier-Hein, Danilo Jimenez Rezende, SM Eslami, Pushmeet Kohli, Andrew Zisserman, and Olaf Ronneberger. hierarchical probabilistic u-net for modeling multi-scale ambiguities. arXiv preprint arXi::1905.1300 2019.
[13] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, 2020.
[14] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[15] Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv preprint 202.v:200000067. 09672, 2021.
[16] 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.
[17] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems, volume 34, 2021.
[18] Valentin De Bortoli, James Thornton, 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.
[19] Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational Diffusion Models. In Advances in Neural Information Processing Systems, volume 34, 2021.
[20] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
2022.
[21] Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow-based generative models with variational dequantization and architecture design. In International Conference on Machine Learning, 2019.
[22] Durk P Kingma and Prafulla Dhariwal. Glow:
Generative flow with invertible 1xl convolutions.
Advances in Neural Information Processing Systems, 31, 2018.
[23] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint 1rXiv:1904.105000 2019.
12

---

[24] Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al. Con- ditional image generation with pixelenn decoders. Advances in Neural Information Processing Systems, 29, 2016.
[25] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, 2014.
[26] Shengjia Zhao, Jiaming Song, and Stefano Ermon. Learning hierarchical features from deep generative models. In International Conference on Machine Learning, 2017.
[27] Fabian Falck, Haoting Zhang, Matthew Willetts, George Nicholson, Christopher Yau, and Chris C Holmes. Multi-facet clustering variational autoencoders. In Advances in Neural Information Processing Systems, volume 34, 2021.
[28] Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems; 2016.
[29] 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.
[30] Ingrid Daubechies. Ten lectures on wavelets. SIAM, 1992.
[31] David Taubman and Michael Marcellin. JPEG2000 image compression fundamentals, standards and practice: image compression fundamentals, standards and practice, volume 642. Springer Science & Business Media, 2012.
[32] Alfred Haar. Zur theorie der orthogonalen funktionnnsssttttemee.. Georg-August- Lanngg-Auuus---ivessiiat Universitat, Gottingen., 1909.
[33] Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. In Advances Neural Information Processing Systems, volume 33, 2020.
[34] Haohan Wang, Xindi Wu, Zeyi Huang. and Eric P Xing. High- fiehefrryyency component helps explain the generalization of convolutional neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8684 8694, 2020.
[35] Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Advances in neural information processing systems, volume 20, 2007.
[36] Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates,
Inc., 2016
[37] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv preprint arXiv: 1509.00519, 2015.
[38] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 98):1733 1997.
[39] Rob Cornish, Anthony Caterini, George Deligiannidis, and Arnaud Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. In International conference machine learning, pages 2133-2143. PMLR, 2020.
[40] Jascha Sohl--iccsttei, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[41] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint 20X6::00066.1299 2020.
13

---

[42] Yann LeCun, Corinna Cortes, and C. J. Burges. MNIST handwritten digit database. http: //yann. lecun. com/ con/exxb///eee /mnist/i..//////////m//mmm 2010.
[43] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[44] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. lece, 2009.
[45] Patryk Chrabaszez, llya Loshchilov, and Frank Hutter. downsampled variant of imagenet as an alternative to the cifar datasets. arXiv preprint
2017.
[46] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[47] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Advances in Neural Information Processing Systems, volume 31, 2018.
[48] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
[49] Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. arXiv preprint 2015.
[50] D La Torre and F Mendivil. The monge-kantorovich metric multimeasures and self-similar multimeasures. Vel-Valued and Variational Analysis, 23(2):119-311, 2015.
[51] Svetlozar Rachev. Probability metrics and the stability of stochastic models, volume 269.
Wiley, 1991.
[52] Stephane Mallat. Multiresolution upproximations and wavelet orthonormal bases of O). Transactions of the American mathematical society, 315(1):69-87, 1989.
[53] Naftali Tishby, Fernando C. Pereira, and William Bialek. The Information Bottleneck Method. 2000.
[54] Lars Maaløe, Marco Fraccaro, Valentin Liévin, and Ole Winther, BIVA: A very deep hierarchy of latent variables for generative modeling. In Advances in Neural Information Processing Systems, volume 32, 2019.
[55] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1604.600461111 07744 2016.
[56] 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 Detito, 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-8035, 2019.
[57] Charles R. Harris, Jarrod Millman, Stéfan J. van 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):357336. 2020.
[58] Lukas Biewald. Experiment tracking with weights and biases, 2020. Software available from wandb.com.
[59] NVIDIA. Nvidia apex. Ntt.a
com/IVVDDAA/ppx..
[60] Guido Van Rossum, The Python Library Reference, release 3.8.2. Python Software Foundation, 2020.
14

---

[61] J. D. Hunter. Matplotlib: A 2D graphics environment. Computing in Science & Engineering. 9(3)990-9) 95, 2007.
[62] Almar al. Klein, Imageio. http:///eeenooo.oorreerrrorooooe
Yolo_5PMIhg,
[63] Lisandro Dalcin and Yao-Lung L Fang. Mpi4py; Status update after 12 years development. Computing in Science & Engineering, 23(4):47-54, 2021.
[64] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, Prettenhofer, R. Weiss, V. Dubourg. J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Joumnal of Machine Learning Research, 12:2825-2830, 2011.
[65] Umesh. Image processing in python. CSI Communications, 23, 2012.
[66] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. 80 million tiny images. http://groupp csail. mit. edu/visionn wwwooonnnnnnnpneeee
[67] Vinay Uday Prabhu and Abeba Birhane. Large image datasets: A pyrrhic win for computer vision? arXiv preprint arXiv:2006.1668 2020.
[68] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset. 2022.
[69] Tomás Mikolov et Statistical language models based on neural networks. Presentation at Google, Mountain View, 2nd April, 80(26), 2012.
[70] Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European Conference on Computer Vision, pages 646-661. Springer, 2016.
[71] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] We state our theoretical and experimental contributions in $1. We develop our multi-resolution framework with corresponding Theorems to
5 in $2 (proofs in Appendix A). We provide our experimental results in $3.
(b) Did you describe the limitations of your work? [Yes] See $5.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Appendix E.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] We stated the assumptions of all theoretical results in $2 and refer to Appendices A and for further details.
(b) Did you include complete proofs of all theoretical results? [Yes] The complete proofs of Theorems to 5, as well as of all additional experimental results are provided in Appendix A.
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as a URL)? [Yes] We provide our code, and instructions on how to download the data and reproduce the main results in the supplementary material. In particular, we refer to the README, file the code repository for further details which follow the NeurlPS Code Completeness Checklist.
15

---

(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendices F and
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [No] Due to the significant computational cost of training extremely deep HVAEs (multiple Nvidia A100 graphic cards with 40GB of GPU memory each rumning for 3 weeks per run), did not perform the multiple runs per hyperparameter setting required to provide error bars for our runs. We note that this is common practice in large-scale HVAE research (compare [9, 10]). Furthermore, in their code base, VDVAE [9], which we directly base our architecture on, reported highly stable test NIL when varying the random seed, varying only in the second decimal place in bits per dim.
(d) Did you include the total amount of compute and the type of resources used (e.g.. type of GPUs, internal cluster, or cloud provider)? [Yes] We discuss our computational resources used and an estimate of the total amount of compute required to reproduce our results briefly in $3, and in detail in Appendix C.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work existing assets, did you cite the creators? [Yes] We discuss all existing assets used in our code base Appendix C.
(b) Did you mention the license of the assets? [Yes] We mention the license of all assets used in Appendix
(c) Did you include any new assets either the supplemental material or as URL? [Yes] We provide source code with instructions on how to reproduce results in the supplementary material. Further details on our code are provided in Appendix C. (d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [Yes] We address this question in Appendix D.
(e) Did you discuss whether the data you are using/curating contains personally identi- fiable information or offensive content? [Yes] We likewise address this question in Appendix D.
5. If you used crowdsoureing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
16


Reviews:
[{"review_id": "k3Nqgy_jg3m", "reviewer": "Reviewer_sqR1", "summary": "This work tries to make an abstraction of the UNet (through the multi-resolution framework). Based on this abstraction, this work analyzes properties of UNet, revealing a regularisation property of UNet, that is implicitly learning a Haar wavelet basis representation. This work further identifies HVAE (based on UNet) as an Euler discretisation of multi-resolution diffusion sum processes. This paper also proposes a weight sharing of HVAE, which reduces the number of parameters.", "questions": " ", "limitations": " ", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, the analysis framework of multi-resolution is original from my perspective, but I preserve some questions on this framework. The paper builds upon heavy notations with less explanation. Therefore, I think the writting quality and clarity can be improved. Besides, it seems that the UNet used in practice has a gap between the UNet discussed in this work, and this gap is not discussed clearly.\n\nBelow are my concerns and questions in detail:\n\n1. The multi-resolution framework use finite dimensional projection. So I expect a projection is defined on a discrete space, such as $Z^2$. However, the partition is w.r.t. functions with domain $R^2$ in Definition 1.\n\n2. In definition 1, why a negative $j$ corresponds to zooming in? Suppose $f(\\cdot) \\in V_j$, where $j$ is negative, and the input space $[0, 2^j]^2$. Such an input space is scaled to $[0, 1]^2$ for $f(2^j \\cdot)$ to keep the same graph. Therefore, a negative $j$ has a smaller graph (since $2^j < 1$) compared to $j=0$. From my perspective, a smaller graph corresponds to zooming out.\n\n3. In definition 2, the author mentions the multi-resolution partition of $L^2(X)$, where $X$ is a compact set. But the partition on a compact set is not defined. It seems that directly applying definition 1 to compact set will have some problems, since A4 implies $X$ should be unbounded to allow a translation transformation. However, a compact is bounded. Therefore, I think it is necessary to define the multi-resolution partition for a compact set $X$ in some place.\n\n4. This paper does not give an introduction to the UNet. This makes it difficult to link the UNet used in practice to the one discussed in this paper. The author could introduce the UNet used in practice somewhere and discuss how the practical UNet is abstracted in this paper.\n\n5. The practical UNet has different number of channels in different layers. Does this work also support the different number of channels?\n\n6. An important feature of UNet is the skip connection across different layers. It seems that the framework in this paper ignores this feature.\n\n7. This paper relies on heavy notations. These notations should be explained more. For example, the author could explain what the practical counterpart is when introducing a notation.\n\n8. When the author refers to the Appendix, the author could refer to a more detailed position.\n\n9. This paper relies on UNet, HVAE and multi-resolution partition. It would be better to add a background section to introduce these important prior works. This would make this paper more readable.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "_4BVGwWYFIt", "reviewer": "Reviewer_fuoh", "summary": "The paper proposes a multi-resolution framework for U-Net architecture, which provides some insights into the architecture: (1) it can relate UNets as diffusion processes (2) average pooling regularizes it by implicitly learning Haar wavelet basis representation of target data. The authors analyze Hierarchical Gaussian VAE (HVAE) by applying the proposed framework to the model, which can explain the instability of sampling in HVAE and helps reduce parameter redundancies.\n", "questions": "- Sampling instabilities on multi-modal datasets such as CIFAR10 and ImageNet are really due to the instability from the hierarchical property? VAE w/o UNet structure also usually suffers from poor sampling performance on such vision datasets.", "limitations": "Limitations are addressed, no negative societal impact observed.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strong points\nAlthough why UNet has been successful in many applications, its regularization effect is poorly understood ever. The framework can naturally explain the instability issue and parameter redundancies of HVAE. The proposed framework is not only elegant but also insightful. In addition, experimental results are overall convincing and support the theory part.\nThe presentation of the results is correct and easy to follow.\n\n### Weak points\n- My only issue is the connection between the insight into sampling instabilities and empirical observation in Sect. 3.3 is unclear (please see my question).\n- It is not easy to refer to corresponding appendix parts (because appendix. A is long and the numbering for subsection is not used).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "6CsDVe0vNOe", "reviewer": "Reviewer_yjaZ", "summary": "This paper proposes a multi-resolution framework of U-net. Theoretical results show the regularisation property of a U-Net with average pooling as implicitly learning a Haar wavelet representation of the data. This framework is also used to identify parameter redundancies and instabilities of HVAEs and verify them via experiments. ", "questions": "No.", "limitations": "(1) Only test on a single VAE model HVAEs.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n\n(1) This is a paper with theoritical insights that help to understand the regularization properties of U-net. \n\n(2) With the help of the proposed theories, the parameter efficiencies of HVAEs are improved as shown by extensive experiments on 4 benchmarks.\n\nCons:\n\n(1) As admitted in the paper, the author only applies the framework to a single VAE model HVAEs so the theory is not guaranteed to generalize. \n\n(2) Though the theory is insightful and novel, the practical technique used to improve parameter efficiencies is just parameter sharing, which is well studied. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept