A Multi-Resolution Framework for U-Nets with Applications to Hierarchical VAEs
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 identiﬁes U-Nets as ﬁnite-dimensional truncations of models on an inﬁnite-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 ﬂow from a point mass, introducing sampling instabilities.
We also demonstrate that HVAEs learn a representation of time which allows for improved parameter efﬁciency 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.
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-toimage 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 a core building block as the neural architecture of choice in state-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, 20].
In spite of their empirical success, it is poorly understood why U-Nets work so well, and what regularisation they impose.
In likelihood-based generative modelling, various model classes are competing for superiority, including normalizing ﬂows [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 a conditional dependence structure, use a 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 likelihood performance, consistently outperforming the previously state-of-the-art autoregressive models,
VAEs and ﬂow models on computer vision benchmarks [9, 10].
HVAEs have undergone a journey ⇤ Equal contribution. Haar wavelet basis elements nullified by projection
Standard basis elements combined by average pooling image graphed as a piecewise constant function Resulting image average pooling projection of design iterations and architectural improvements in 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, as well as the use of different types of residual cells in the decoder [9, 10] (see §4 and Fig.
A.1 for a detailed discussion).
However, a theoretical understanding of these choices is lacking.
For instance, it has not been shown why a residual backbone may be beneﬁcial, or what the speciﬁc 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 speciﬁcally.
Doing so allows us to relate HVAEs to diffusion processes, and also to motivate a new type of piecewise time-homogenenous model which demonstrates state-of-the-art performance with approximately half the number of parameters of a VDVAE [9].
More formally, our contributions are as follows: (a) We provide a multi-resolution framework for U-Nets.
We formally deﬁne U-Nets as as acting over a multi-resolution hierarchy of L 2 ([0 , 1] 2 ) .
We prove that average pooling is a conjugate operation to projection in the Haar wavelet basis within L 2 ([0 , 1] 2 ) .
We use this insight to show how U-Nets with average pooling implicitly learn a 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-art HVAEs as an example, identifying their residual cell structure as a type of two-step forward Euler discretisation of a multi-resolution diffusion bridge.
We uncover that this diffusion process ﬂows 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 to 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-theart 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 ﬁnally 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 ﬁnd that Fourier features (which discrete-time diffusion models strongly beneﬁt from [19]) do not improve performance in the HVAE setting.
The Multi-Resolution Framework
A grayscale image with inﬁnite resolution can be thought of as the graph 2 of a two-dimensional function over the unit square.
To store these inﬁnitely-detailed images in computers, we project them to some ﬁnite resolution.
These projections can still be thought of as the graphs of functions with support over the unit square, but they are piecewise constant on ﬁnitely many intervals or ‘pixels’, e.g. 512 2 pixels, and we store the function values obtained at these pixels in an array or ‘grid’.
The relationship between the ﬁnite-dimensional version and its inﬁnitely-ﬁne counterpart depends entirely on how we construct this projection to preserve the details we wish to keep.
One approach is to prioritise preserving the large-scale details of our images, so unless closely inspected, the projection is indistinguishable from the original.
This can be achieved with a multi-resolution projection [30] of the image.
In this section we introduce a multi-resolution framework for constructing neural network architectures 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: Deﬁnitions 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 a function, and its ﬁnite-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 B.1.
We can formalise the aforementioned intuitions by deﬁning a multi-resolution hierarchy [30] of sub-spaces we may project to:
Deﬁnition 1.
[ Daubechies (1992) [30]] Given a nested sequence of approximation spaces · · · ⇢ V 1 ⇢ V 0 ⇢ V  1 ⇢ · · · , { V  j } j 2 Z is a multi-resolution hierarchy of the function space
L 2 ( R m ) if: ( A1)
[Equation]
A
[Equation]
( A4) f
( · ) 2 V 0 , f ( ·  n ) 2 V 0 for n 2 Z .
For a compact set X ⇢ R m , a multi-resolution hierarchy of L 2 (
X ) is { V  j } j 2 Z as deﬁned above, restricting functions in V  j to be supported on X .
In Deﬁnition 1, the index j references how many steps we took in our thought experiment, so negative j corresponds to ‘zooming in’ on the images.
The original image 3 is a member of L 2 ([0 , 1] 2 ) , the space of square-integrable functions on the unit square, and its ﬁnite projection to 2 j · 2 j many pixels is a member of V  j .
Images can be represented as piecewise continuous functions in the subspaces
V  j = { f 2 L 2 ([0 , 1]) | f | [2  j · k, 2  j · ( k +1))
=
c k , k 2 { 0 , . . .
, 2 j  1 } , c k 2 R } .
The nesting property V  j +1 ⇢ V  j ensures that any image with (2 j  1 ) 2 pixels can also be represented by (2 j ) 2 pixels, but at a higher resolution.
Assumption ( A1 ) states that with inﬁnitely many pixels, we can describe any inﬁnitely 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 in any V  j if we know the basis of V 0 .
One basis made by extrapolating from V 0 in this way is known as a wavelet basis [30].
Wavelets have proven useful for representing images, for instance in the JPEG standard [31], and are constructed to be orthonormal.
Now suppose we have a probability measure ⌫ 1 over inﬁnitely detailed images represented in
L 2 ([0 , 1] 2 ) and wish to represent it at a lower resolution.
Similar to how we did for inﬁnitely detailed images, we want to project the measure ⌫ 1 to a lower dimensional measure ⌫ j on the ﬁnite dimensional space V  j .
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 ⌫ 1 .
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 multi-resolution bridge from ⌫ 1 to ⌫ j , as deﬁned below.
Deﬁnition 2.
Let X ⇢ R m be compact, { V  j } 1 j =0 be a multi-resolution hierarchy of scaled so
L 2 ( X ) =
S
j 2 N 0 V  j and V
= { 0 } .
If D ( L 2 ( X )) is the space of probability measures over L 2 ( X ) , then a family of probability measures { ⌫ t } t 2 [0 , 1] on L 2 ( X ) is a multi-resolution bridge if:
(i) there exist increasing times I := { t j } j 2 N 0 where t 0 = 0 , lim j !1 t j = 1 , such that s 2 [ t j , t j +1 ) implies supp ( ⌫ s ) ⇢ V  j , i.e ⌫
s 2 D ( V  j ) ; and, (ii) for s 2 (0 , 1) , the mapping s 7!
⌫ s is continuous for s 2 ( t j , t j +1 ) for some j .
The continuous time dependence in Deﬁnition 2 plays a movie of the measure ⌫ 0 supported on Vgrowing to ⌫ 1 , a measure on images with inﬁnite resolution.
At a time interval [ t j , t j +1 ) , the space
V  j which the measure is supported on is ﬁxed.
We therefore deﬁne a ﬁnite-dimensional model transporting probability measures within V  j , but at t j +1 the support ﬂows over to V  j  1 .
Given a multi-resolution hierarchy, we may glue these ﬁnite models, each acting on a disjoint time interval, together in a uniﬁed fashion.
In Theorem 1 we show this for the example of a continuous-time multi-resolution diffusion process truncated up until some time t J = T 2 (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 B j : [ t j , t j +1 ) ⇥
D ( V  j ) 7!
D ( V  j ) be a linear operator (such as a diffusion transition kernel, see Appendix A) for j < J with coefﬁcients µ ( j ) ,  ( j ) :
[ t j , t j +1 ) ⇥
V  j 7!
V  j , and deﬁne the natural extensions within V  J in bold, i.e. B j :=
B j  I V ?
j .
Then the operator
B : [0 , T ] ⇥ D ( V  J ) 7!
D ( V  J ) and the coefﬁcients µ ,  : [0 , T ] ⇥ V
J 7!
V  J given by
[Equation]
induce a multi-resolution bridge of measures from the dynamics for t 2 [0 , T ] and on the standard basis as dZ t = µ t ( Z t ) dt +  t ( Z t ) dW t (see Appendix A.4 for details) for Z t 2 V  j for t 2 [ t j , t j +1 ) , i.e. a multi-resolution diffusion process.
S(cid:78)i(cid:83) c(cid:82)(cid:81)(cid:81)ec(cid:87)i(cid:82)(cid:81)
Be(cid:81)d /
B(cid:82)(cid:87)(cid:87)(cid:79)e(cid:81)ec(cid:78)
See Appendix B.2 for details. The concept of a multi-resolution bridge will become important in Section 2.2 where we will show that current U-Net bottleneck structures used for unconditional sampling impose a multi-resolution bridge on the modelled densities.
To preface this, we here provide a description of a U-Net within our framework, illustrated in 2.
Consider B j, ✓ , F j, ✓ : D ( V  j ) !
D ( V  j ) as the forwards and backwards passes of a U-Net on resolution j .
Further, let P  j +1 : D ( V  j ) !
D ( V  j +1 ) and E  j : D ( V  j +1 ) !
D ( V  j ) be the projection (here: average pooling) and embedding maps (e.g. interpolation), respectively.
When using an L 2 -reconstruction error, a U-Net [1] architecture implicitly learns a sequence of models
B j,  : D ( V  j +1 ) ⇥ D ( V ?
j +1 ) 7!
D ( V  j ) due to the orthogonal decomposition V  j = V
j +1  U
j +1 where U  j
+1 :
= V  j \ V ?
j +1 .
The backwards operator for the U-Net has a (bottleneck) input from D ( V  j +1 ) and a (skip) input yielding information from D ( V ?
j +1 ) .
A simple bottleneck map U j, ✓ : D ( V  j ) !
D ( V  j ) (without skip connection) is given by
[Equation]
and a U-Net bottleneck with skip connection is
[Equation]
In HVAEs, the map U j,  : D ( V  j ) !
D ( V  j ) is trained to be the identity by minimising reconstruction error, and further shall approximate U j, ✓ ⇡ U j,  via a KL divergence.
The L 2 -reconstruction error for U j,  has an orthogonal partition of the inputs from V  j +1 ⇥ V  j , hence the only new subspace added is U  j
+1 .
As each orthogonal U  j +1 is added sequentially in HVAEs, the skip connections induce a multi-resolution structure of this hierarchical neural network structure.
What we will investigate in Theorem 3 is the regularisation imposed on this partitioning by enforcing U j, ✓ ⇡ 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 deﬁned U-Net architectures within our multi-resolution 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 basis (see A.2, Eq.
(A.9)).
The standard basis is not convenient to derive theoretical results.
It is instead preferable to use a basis natural to the multi-resolution bridge imposed by a U-Net with a corresponding projection operation, which for average pooling is the Haar (wavelet) basis [32] (see Appendix A.2).
The Haar basis, like a Fourier basis, is an orthonormal basis of L 2 ( X ) which has desirable L 2 -approximation properties.
We formalise this in Theorem 2 which states that the dimension reduction operation of average pooling in the standard basis is a conjugate operation to co-ordinate projection within the Haar basis (details are provided in Appendix A.2).
Theorem 2.
Given V
j as in Deﬁnition 1, let x 2 V  j be represented in the standard basis E j and Haar basis j .
Let ⇡ j : E j 7!
j be the change of basis map illustrated in Fig.
3, then we have the conjugacy
[Equation]
[Figure]
Theorem 2 means that if we project an image from V  j to V  j +1 in the Haar wavelet basis, we can alternatively view this as changing to the standard basis via ⇡  1 j , performing average pooling, and reverting back via ⇡ j  1 (see
Figure 3).
This is important because the Haar basis is orthonormal, which in Theorem 3 allows us to precisely quantify what information is lost with average pooling.
Theorem 3.
Let { V  j } Jj =0 be a multi-resolution hierar- chy of V  J where V  j = V
j +1  U
j +1 , and further, let F j,  , B j, ✓ : D ( V  j ) 7!
D ( V  j ) be such that B j, ✓ F j,  = I with parameters  and ✓ .
Deﬁne F j 1 | j 2 ,  :=
F j 1 ,   · · ·  F j 2 ,  by F j,  : D ( V  j ) 7!
D ( V  j +1 ) where F j,  :=
proj V
j +1  F j,  , and analogously deﬁne B j 1 | j 2 , ✓ with B j, ✓ := B j, ✓  embd V  j .
Then, the sequence { B 1 | j, ✓ ( F 1 | J,  ⌫ J ) } J j = forms a discrete multi-resolution bridge between F 1 | J,  ⌫ J and B 1 | J, ✓ F 1 | J,  ⌫ J at times { t j } Jj =1 , and J
[Equation]
where W 2 is the Wasserstein- 2 metric and
F
[Equation]
2 is the Lipschitz constant of F j | J,  .
Theorem 3 states that the bottleneck component of a U-Net pushes the latent data distribution to a ﬁnite multi-resolution basis, speciﬁcally a Haar basis when average pooling is used.
To see this, note that the RHS of Eq.
(A.65) is itself upper-bounded by the L 2 -reconstruction error.
This is because the Wasserstein-2 distance ﬁnds the inﬁnimum 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  j } Jj =0 which are (by Theorem 2) discarded by the bottleneck structure when using a U-Net architecture with average pooling.
Theorem 3 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  j } Jj =0 .
Consequently, the bottleneck architecture implicitly decomposes our data into a Haar wavelet decomposition, and when the skip connections are absent (like in a traditional auto-encoder) our network learns to compress the discarded subspaces U  j .
This characterises the regularisation imposed by a 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 to 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 speciﬁc data types.
2.3 Example: HVAEs as Diffusion Discretisations
To show what practical inferences we can derive from our multi-resolution framework, we apply it to analyse state-of-the-art 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 in VDVAE cells to the (forward and backward) operators F j,  and B j, ✓ within our framework and identify them as a type of two-step forward Euler discretisation of a diffusion process.
When used with a U-Net, as is done in VDVAE [9], this creates a multi-resolution diffusion bridge by Theorem 4.
Theorem 4.
Let t J :=
T 2 (0 , 1) and consider (the p ✓ backward pass) B ✓ , 1 | J : D ( V  J ) 7!
D ( V 0 ) given in multi-resolution Markov process in the standard basis:
[Equation]
where proj U  j Z t j = 0 , k Z t k 2 >
k Z s k 2 with 0  s < t  T and for a measure ⌫ J 2 D ( V  J )
we have X T , Z
⇠ F  ,J | 1 ⌫ J =  { 0 } .
Then, VDVAEs approximates this process, and its residual cells are a type 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 corresponding 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 Z i variables. Similar processes have been empirically observed as efﬁcient ﬁrst-order approximates to higher-order chains, for example the memory state in 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.
[Table]
(cid:73)(cid:76)(cid:74)(cid:88)(cid:85)e DO NOT CHANGE!!!!!!!!!! choo(cid:86)e (cid:82)(cid:85) (cid:86)ample
NVAE choo(cid:86)e
(cid:82)(cid:85) (cid:86)ample E(cid:88)(cid:79)(cid:72)(cid:85)-M(cid:68)(cid:85)(cid:88)(cid:92)(cid:68)(cid:80)(cid:68) VAE choo(cid:86)e o(cid:85)
(cid:86)am(cid:83)le
VDVAE choo(cid:86)e
(cid:82)(cid:85) (cid:86)ample NVAE choo(cid:86)e (cid:82)(cid:85) (cid:86)ample
E(cid:88)(cid:79)(cid:72)(cid:85)-M(cid:68)(cid:85)(cid:88)(cid:92)(cid:68)(cid:80)(cid:68) VAE
To show how VDVAEs impose the growth of the Z t , we prove that the bottleneck component of VDVAE’s U-Net enforces Z 0 = 0 .
This is done by identifying that the measure ⌫ 0 , which a VDVAE connects to the data ⌫ 1 via a multiresolution bridge, is a point mass on the zero function.
Consequently the backward pass must grow from this, and the network learns this in a monotonic manner as we later conﬁrm in our experiments (see §3.2).
Theorem 5.
Consider the SDE in Eq.
(A.76) , trained through the ELBO in Eq.
B.101.
Let ˜ ⌫
J denote the data measure and ⌫
=  { 0 } be the initial multi-resolution bridge measure imposed by VDVAEs.
If q  ,j and p ✓ ,j are the densities of B  , 1 | j F J | 1 ˜ ⌫ J and B ✓ , 1 | j ⌫ 0 respectively, then a VDVAE optimises the boundary condition min ✓ ,  KL ( q  , 0 , 1 || q  , 0 p ✓ , 1 ) , where a double in- dex indicates the joint distribution.
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 V 0 : a Dirac mass atas the pooling here cascades completely to V 0 .
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 ⌫ 0 is disastrously unstable as it enforces a data set, with potentially complicated topology to derive from a point-mass in U  j at each t =
t j , and we observe the resulting sampling instability in our experiments in §3.3.
We note that similar arguments are applicable in settings without a latent hierarchy imposed by a U-Net, see for instance [39].
The
VDVAE architecture does, however, bolster this rate through the Z (  ) i, + term, which is absent in NVAEs [10], in the discretisation steps of the residual cell.
We empirically observe this controlled backward error in Fig.
[Right].
We refer to Fig.
A.1 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 over-parameterised and not informed by this continuoustime formulation.
The continuous time analogue of VDVAEs [9] in Theorem 4 has time dependent coefﬁcients  µ t, 1 ,  µ t, 2 ,   t .
We hypothesise that the increasing diffusion process in Z i implicitly encodes time.
Hence, explicitly representing this in the model, for instance via ResNet blocks with independent parameterisations at every time step, is redundant, and a time-homogeneous model (see Appendix A.6 for a precise formulation)—practically speaking, performing weight-sharing across time time steps/layers—has the same expressivity, but requires far fewer parameters than the state-of-the-art VDVAE.
It is worth noting that such a time-homogeneous model would make the parameterisation of HVAEs more similar to the recently popular (score-based) diffusion models [40, 41] which perform weight-sharing across all time steps.
Experiments
We compare our runs of VDVAE with original hyperparameters [9] ( VDVAE ⇤ ) against our weight-shared VDVAE ( WS-VDVAE ).
While WS-VDVAE s have improved parameter efﬁciency by a factor of 2, they reach similar NLL as VDVAE ⇤ with the simple modiﬁcation inspired by our framework (weight sharing).
We note that a parameter count cannot be provided for VDM [19] as the code is not public and the manuscript does not specify it
[Table]
AR
[Table]
Diff –  2 .
I m ag e N e t ⇥
WS-VDVAE (ours)
VAE 55m 
[Table]
WS-VDVAE (ours)
VAE 85m 
[Table]
VDVAE ⇤ [Table] ours)
VAE 119m 
[Table]
67 NVAE [Table]
VAE 268m 
[Table]
92 VDVAE [9]
VAE 119m  3 .
80 VDM [19]
Diff –  3 .
C e l e b A ⇥
WS-VDVAE (ours)
VAE 75m  2 .
VDVAE ⇤ (ours)
VAE 125m  2 .
02 NVAE [10]
VAE 153m  2 .
In the following we probe the theoretical understanding of HVAEs gained through our framework, demonstrating its utility in four experimental analyses: ( a ) Improving parameter efﬁciency in
HVAEs, ( b ) Time representation in HVAEs and how they make use of it, ( c ) Sampling instabilities in HVAEs, and ( d ) Ablation studies. We train HVAEs using VDVAE [9] as the basis model on ﬁve datasets: MNIST [42], CIFAR10 [43], two downsampled versions of ImageNet [44, 45], and CelebA [46], splitting each into a training, validation and test set (see Appendix D for details). In general, reported numeric values refer to Negative
Log-Likelihood (NLL) in nats (MNIST) or bits per dim (all other datasets) on the test set at model convergence, if not stated otherwise.
We note that performance on the validation and test set have similar trends in general.
An optional gradient checkpointing implementation to trade in GPU memory for compute is discussed in Appendix F. Appendices
F and G
deﬁne the HVAE models we train, i.e. p ✓ ( z L ) , p ✓ ( z l | z >l ) , q  ( z L | x ) , q  ( z l | z >l , x ) and p ✓ ( x | ~ z ) , and present additional experimental details and results.
We provide our PyTorch code base at https://github.com/FabianFalck/unet-vdvae (see Appendix C for details).
3.1 “More from less”: Improving parameter efﬁciency in HVAEs
In §2.3, we hypothesised that a time-homogeneous model has the same expressivity as a model with time-dependent coefﬁcients, 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 CIFAR10 with { 1 , 2 , 3 , 4 } ResNet blocks (referred to as { x1 , x2 , x3 , x4 }) in each resolution with spatial dimensions { 32 2 , 16 2 , 8 2 , 4 2 , 1 2 } ( VDVAE ⇤ ), and compare their performance when weight-sharing a single parameterised block per resolution { 2 , 3 , 5 , 10 , 20 } times (referred to as { r2 , r3 , r5 , r10 , r20 }; WS-VDVAE ), excluding projection and embedding blocks.
As hypothesised by our framework, yet very surprising in HVAEs, NLL after 1m 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 to the x 1 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.
CIFAR10: 13.0m; 17.3m; 21.6m), yet fewer activations, latent variables, and number of timesteps at which the coupled SDE in Eq.
(A.76) is discretised.
We now scale these ﬁndings up to large-scale hyperparameter conﬁgurations.
We train VDVAE closely following the state-of-the-art hyperparameter conﬁgurations in [9], speciﬁcally 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-sharing: 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 NLL compared to any
HVAE on CIFAR10, ImageNet32 and CelebA.
Furthermore, our WS-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 identiﬁed in Theorem 5, and ﬁnding 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 to make HVAEs more parameter efﬁcient.
We now want to explain and understand this behavior further.
In Fig.
6, we measure k Z
i k 2 , the L 2 -norm of the residual state at every backward/top-down block with index i, over several batches for models trained on MNIST (see Appendix G.2 for the corresponding ﬁgure of the forward/bottom-up pass, and similar results on CIFAR10 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 processes which are increasing in the Z i variables, and hence learn to represent a notion of time in their residual state.
It is now straightforward to ask how HVAEs beneﬁt 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.e. forcing a “ﬂat line” in Fig.
6 [Left], learning deteriorates after a short while, resulting in poor NLL results compared to the runs with a regular, non-normalised residual state.
This evidence conﬁrms 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
[Table] High ﬁdelity unconditional samples of faces, e.g. from models trained on CelebA, cover the front pages of state-of-the-art HVAE papers [9, 10].
Here, we question whether face datasets are an appropriate benchmark for HVAEs.
In Theorem 5, we identiﬁed the aforementioned state-of-the-art HVAEs as ﬂow from a point mass, hypothesising instabilities during sampling.
And indeed, when sampling from our trained VDVAE ⇤ with state-of-the-art conﬁgurations, we observe high ﬁdelity and diversity samples on MNIST and CelebA, but unrecognisable, yet diverse samples on CIFAR10,
ImageNet32 and ImageNet64, in spite of state-of-the-art test set NLLs (see Fig.
7 and Appendix G.3).
We argue that MNIST and CelebA, i.e. numbers and faces, have a more uni-modal nature, and are in this sense easier to learn for a discretised multi-resolution process ﬂowing to a 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.
3.4 Ablation studies
We conducted several ablation studies which support our experimental results and further probe our multi-resolution framework for HVAEs.
In this section we note key ﬁndings—a detailed account of all ablations can be found in Appendix G.4.
In particular, we ﬁnd 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 a performance gain in HVAEs, in contrast to state-of-the-art diffusion models, where they signiﬁcantly improve performance [19].
This is consistent with our framework’s ﬁnding that a 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 a diffusion process and impose an SDE structure into the model, empirically compare a multi-resolution bridge to a single-resolution model, and investigate synchronous vs. asynchronous processing in time between the forward and backward pass.
Related work
U-Nets.
A U-Net [1] is an autoencoding architecture with multiple resolutions where skip connections enable information to pass between matched layers on opposite sides of the autoencoder’s bottleneck.
These connections also smooth out the network’s loss landscape [47].
In the literature, U-Nets tend to be convolutional, and a 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 as operators on measures interleaved by average pooling as the down-sampling operation (and a corresponding inclusion operation for up-sampling), and we formally characterise U-Nets in Section 2.1 and Appendix B.2.
Prior to our work, the dimensionality-reducing bottleneck structure of U-Nets was widely acknowledged as being useful, however it was unclear what regularising properties a U-Net imposes.
We provided these in §2.
HVAEs.
The evolution of HVAEs can be seen as a quest for a parameterisation with more expressiveness 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 VAEs suffer from stability issues even for very small stochastic depths.
Nouveau
VAEs (NVAE)
[10] and Very Deep VAEs (VDVAE) [9] combine the improvements of several earlier
HVAE models (see Appendix B for details), while scaling up to 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 a warm-up deterministic autoencoder training phase or data-speciﬁc initialisation.
VDVAE achieves a 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 use 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.
Conclusion
In this work, we introduced a multi-resolution framework for U-Nets.
We provided theoretical results which uncover the regularisation property of the U-Nets bottleneck architecture with average pooling as implicitly learning a Haar wavelet representation of the data.
We applied our framework to HVAEs, identifying them as multi-resolution diffusion processes ﬂowing to a point mass.
We characterised their backward cell as a type of two-step forward Euler discretisations, providing an alternative to score-matching to approoximate a continuous-time diffusion process [16, 18], and observed parameter redundancies and instabilities.
We veriﬁed the latter theoretical insights in both small- and large-scale experiments, and in doing so trained the deepest ever HVAEs.
We explained these results by showing that HVAEs learn a representation of time and performed extensive ablation studies.
An important limitation is that the proven regularisation property of U-Nets is limited to using average pooling as the down-sampling operation.
Another limitation is that we only applied our framework to
HVAEs, though it is possible to apply it to other model classes.
It could also be argued that the lack of exhaustive hyperparameter optimisation performed is a 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 efﬁciency, and hence view it as a strength of our results.
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/Z/19/Z), and the Enrichment Scheme of The Alan Turing Institute under the EPSRC Grant EP/N510129/1.
Chris Williams acknowledges support from the Defence Science and Technology (DST) Group and from a ESPRC DTP Studentship.
Dominic Danks is supported by a Doctoral Studentship from
The Alan Turing Institute under the EPSRC Grant EP/N510129/1.
Christopher Yau is funded by a UKRI Turing AI Fellowship (Ref: EP/V023233/1).
Chris Holmes acknowledges support from the Medical Research Council Programme Leaders award MC_UP_A390_1107, The Alan Turing
Institute, Health Data Research, U.K., and the U.K. Engineering and Physical Sciences Research
Council through the Bayes4Health programme grant.
Arnaud Doucet acknowledges support of the
UK Defence Science and Technology Laboratory (Dstl) and EPSRC grant EP/R013616/1.
This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary
University Research Initiative.
Arnaud Doucet also acknowledges support from the EPSRC grant
EP/R034710/1.
Matthew Willetts is grateful for the support of UCL Computer Science and The Alan
Turing Institute.
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 2 HPC service (https://www.baskerville.ac.uk/) which we detail in the following.
First, this research was supported in part through computational resources provided by The Alan Turing Institute under EPSRC grant
EP/N510129/1 and with the help of a generous gift from Microsoft Corporation.
Second, we 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/T022221/1) and the Digital Research Infrastructure programme (EP/W032244/1) and is operated by Advanced Research Computing at the University of
Birmingham.
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 at the University of
Oxford for their help and support.
We thank Haoting Zhang (University of Cambridge) for valuable comments on the implementation; Huiyu Wang (Johns Hopkins University) for a useful discussion on gradient checkpointing; and Ruining Li and Hanwen Zhu (University of Oxford) for kindly proofreading the manuscript.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect 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 1 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 B for further details.
(b) Did you include complete proofs of all theoretical results?
[Yes] The complete proofs of Theorems 1 to 5, as well as of all additional experimental results are provided in Appendix A. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in 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.md ﬁle in the code repository for further details which follow the NeurIPS Code Completeness Checklist.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendices F and D. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] Due to the signiﬁcant computational cost of training extremely deep HVAEs (multiple Nvidia A100 graphic cards with 40GB of GPU memory each running for 3 weeks per run), we 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 NLL 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 brieﬂy 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 uses existing assets, did you cite the creators?
[Yes] We discuss all existing assets used in our code base in Appendix C. (b) Did you mention the license of the assets?
[Yes] We mention the license of all assets used in Appendix C. (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] We provide our source code with instructions on how to reproduce our 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ﬁable information or offensive content?
[Yes] We likewise address this question in Appendix D. 5.
If you used crowdsourcing or 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, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]