A Multi-Resolution Framework for U-Nets with Applications to Hierarchical VAEs
Anonymous Author(s)
AfﬁliationAddress email
Abstract
U-Net architectures are ubiquitous in state-of-the-art deep learning, however theirregularisation properties are understudied.
In this paper, we formulate a multiresolution framework which identiﬁes U-Nets as ﬁnite-dimensional truncationsof models on an inﬁnite-dimensional function space.
We provide theoreticalresults which prove that average pooling corresponds to projection within thespace of square-integrable functions and show that U-Nets with average poolingimplicitly learn a Haar wavelet basis representation of the data.
We then leverageour framework to identify state-of-the-art hierarchical VAEs (HVAEs), whichhave a U-Net architecture, as forward Euler discretisations of multi-resolutiondiffusion sum processes which ﬂow to a point mass, introducing instabilities.
Wealso demonstrate that HVAEs learn a representation of time which allows forweight-sharing without detriment.
We use this observation to achieve state-of-theart 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 forimage segmentation in biomedical applications [1], U-Nets have been widely applied for image-toimage translation [2], image restoration [3, 4] and multiview learning [5], amongst other tasks.
Theyalso form a core building block as the neural architecture of choice in state-of-the-art generativemodels, particularly for images, such as HVAEs [6, 7] and diffusion models [8, 9, 10, 11, 12, 13, 14].
In spite of their empirical success, the regularisation they impose is poorly understood.
In likelihood-based generative modelling, various model classes are competing for superiority, including normalizing ﬂows [15, 16], autoregressive models [17, 18], diffusion models and hierarchicalvariational autoencoders (HVAEs), the latter of which we focus on in this work.
HVAEs form groupsof latent variables with a conditional dependence structure, use a U-Net neural architecture, and aretrained with the typical VAE ELBO objective (for a detailed introduction to HVAEs, see Appendix
C).
HVAEs show impressive synthesis results on facial images, and yield competitive likelihoodperformance, consistently outperforming the previously state-of-the-art autoregressive models, VAEsand ﬂow models on computer vision benchmarks [6, 7].
HVAEs have undergone a journey of designiterations and architectural improvements in recent years, for example the introduction a deterministicbackbone [19, 20, 21] and ResNet elements [22, 23] with shared parameters between the inferenceand generative model parts.
There has also been a massive increase in the number of latent variablesand overall stochastic depth, as well as the use of different types of residual cells in the decoder [6, 7](see §4 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ﬁccell structure in VDVAEs [6] and NVAEs [7] correspond to, or how they could be improved.
In this paper we provide a theoretical framework for understanding U-Nets and apply it to HVAEs.
Doing so allows us to explain various empirical ﬁndings, to relate HVAEs to diffusion processesand also to motivate a new type of model which demonstrates state-of-the-art performance withapproximately half the number of parameters of a VDVAE [6].
More formally, our contributionsare as follows: (a) We provide a multi-resolution framework for U-Nets.
We formally deﬁne UNets as multi-resolution bridges over a partition of L 2 ([0 , 1] 2 ) .
We prove that average pooling is aconjugate operation to projection in the Haar wavelet basis within L 2 ([0 , 1] 2 ) .
We use this insightto show how U-Nets with average pooling implicitly learn a Haar wavelet basis representation ofthe data, characterising the regularisation within U-Nets.
(b) We apply this framework to state-ofthe-art HVAEs as an example, identifying their residual cell structure as a two-step forward Eulerdiscretisation of a multi-resolution diffusion sum process.
We uncover that this diffusion processﬂows to a point mass, which causes instabilities, for instance during sampling, and identify parameterredundancies through our continuous-time formulation.
(c)
In our experiments, we demonstrate thesesampling instabilities and train HVAEs with the largest stochastic depth ever, achieving state-ofthe-art performance with half the number of parameters by exploiting our theoretical insights.
Weexplain these results by uncovering that HVAEs secretly represent time in their state and show thatthey use this information during training.
We ﬁnally provide extensive ablation studies which, forinstance, rule out other potential factors which correlate with stochastic depth, show the empiricalgain of multiple resolutions, and ﬁnd that Fourier features (a key performance driver in discrete-timediffusion models [14]) 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 1 of a two-dimensionalfunction over the unit square.
To store these inﬁnitely-detailed images in computers, we projectthem to some ﬁnite resolution, e.g. 512 2 pixels.
These projections can still be thought of as thegraphs of functions, but now a discrete one and only deﬁned at a ﬁnite number of grid points.
The relationship between the discretised version and its inﬁnitely-ﬁne counterpart depends entirelyon how we construct this projection to preserve the details we wish to keep.
One approach is toprioritise preserving the large-scale details of our images, so unless closely inspected the projectionis indistinguishable from the original.
This can be achieved with a multi-resolution projection [24] ofthe image.
In this section, we introduce a multi-resolution framework for constructing neural networkarchitectures that utilise such projections, prove what regularisation properties they impose, and showas an example how HVAEs with a U-Net [1] architecture can be interpreted in our framework.
Proofsof all theorems and further mathematical details of the 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 informallyexplained through the following thought experiment.
Imagine we have an image, represented as thegraph of a function, and its ﬁnite-dimensional projection drawn on the wall.
The multi-resolutiongoodness-of-ﬁt metric is to look at the wall, start walking away from it and stop when the image andits projection are indistinguishable by eye.
The number of steps we took away from the wall can beconsidered a measure of ‘how far away’ the approximation is from the underlying function.
The goalof the multi-resolution projection is therefore to have to take as few steps away as possible.
The readeris encouraged to physically conduct this experiment with the examples provided in Appendix B. Wecan formalise the aforementioned intuitions by deﬁning a multi-resolution partition [24] of sub-spaceswe may project to:
Deﬁnition 1.
[ Daubechies (1992) [24]] Given a nested sequence of approximation spaces·
[Equation]
{+⊂ V 1 ⊂ V 
+}
⊂
V − 1 ⊂ · · · , { V j } j ∈ Z is a multi-resolution partition of the functionspace L 2 ( R m ) if: ( A1) (cid:83)
j ∈ Z V j =
L 2 ( R m ) ; ( A2) (cid:84)
j ∈ Z V j = { 0 } ; ( A3) f ( · ) ∈ V j ⇔ f (2 j · ) ∈ V 0 ;( A4) f ( · ) ∈ V 0 ⇔ f ( · − n ) ∈ V 0 .
In Deﬁnition 1, the index j references (the log of) how many steps we took in our thought experiment,so negative j corresponds to zooming in on the image.
The original image 2 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 manypixels is a member of V − j .
The nesting property V − j +1 ⊂ V − j ensures that any image with (2 j − 1 ) 2pixels can also be represented by (2 j ) 2 pixels, but at a higher resolution.
Assumption ( A1 ) states thatwith inﬁnitely many pixels, we can describe any inﬁnitely detailed image.
In contrast, ( A2 ) says thatwith no pixels, we cannot approximate any images.
Assumptions ( A3 ) and ( A4 ) allow us to form abasis for images in any V j if we know the basis of V 0 .
One basis made by extrapolating from V {+
 inthis+} way is known as a wavelet basis [24].
Wavelets have proven useful for representing images, forinstance in the JPEG standard [25], and are constructed to be orthonormal.
Now suppose we have a probability measure ν ∞ over inﬁnitely ﬁne images represented in L 2 ([0 , 1] 2 )and wish to represent this at a lower resolution.
Similar to how we did for inﬁnitely detailed images,we want to project the measure ν ∞ to a lower dimensional measure ν j on the ﬁnite dimensionalspace V − j .
In extension to this, we want the ability to reverse this projection so that we may samplefrom the lower dimensional measure and create a generative model for ν ∞ .
We would like to againprioritise the presence of large-scale features of the original image within the lower dimensionalsamples.
We do this by constructing a multi-resolution bridge from ν ∞ to ν j , as deﬁned below.
Deﬁnition 2.
Let X be a compact subset of R m and { V − j } ∞ j =0 be a multi-resolution partition of
L 2 ( X ) scaled such that L 2 ( X ) =
(cid:83) j ∈ N 0 V − j and V
= { 0 } .
Deﬁne D ( L 2 ( X )) to be the spaceof probability measures over L 2 ( X ) .
A family of probability measures { ν t } t ∈ [0 , 1] on L 2 ( X ) is amulti-resolution bridge if:
(i) there exist increasing times I := { t j } j ∈ N 0 where t 0 = 0 , lim j →∞ t j = 1 , such that s ∈ [ t j , t j +1 ) implies supp ( ν t ) ⊂ V − j , i.e ν t ∈ D ( V − j ) ; and,
(ii) for s ∈ (0 , 1) , the mapping s (cid:55)→ ν
s is continuous (in the weak sense, see Appendix A) if s < t j .
The continuous time dependence in Deﬁnition 2 plays a movie of the measure ν 0 supported on V 0growing to ν ∞ , 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 may therefore deﬁne a ﬁnite-dimensional modeltransporting probability measures within V − j , but at t j +1 the support ﬂows over to V − j − 1 .
Given amulti-resolution partition, 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-timediffusion model up until some time t J = T ∈ (0 , 1) and in the standard basis given in the nextsub-section, which will be useful in §2.3 when viewing HVAEs as discretisations of continuous-timemulti-resolution diffusion models on functions.
Theorem 1.
Let B j : [ t j , t j +1 )
× D ( V − j ) (cid:55)→ D ( V − j ) be a linear operator (such as a diffusiontransition kernel, see Appendix A) for j < J with coefﬁcients µ ( j ) , σ ( j ) :
[ t j , t j +1 )
× V − j
(cid:55)→ 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 ) (cid:55)→ D ( V − J ) and the coefﬁcients µ , σ : [0 , T ] × V − J
(cid:55)→ V − J given by B := J (cid:88)
j = 1 [ t j ,t j +1 ) · B j , µ := J (cid:88) j = 1 [ t j ,t j +1 ) · µ ( j ) , σ := J (cid:88)
j = 1 [ t j ,t j +1 ) · σ ( j ) ,induce a multi-resolution bridge of measures from the dynamics for t ∈ [0 , T ] and on the standardbasis as dX t = µ t ( X t ) dt + σ t ( X t ) dW t (see Appendix A for the details of this integration) for
X t ∈ V − J , i.e. a (backward) multi-resolution diffusion process.
A U-Net [1] reduces the dimension via down-sampling operations towards a bottleneck.
Layers ofdifferent dimensions are not considered to be acting on a common space.
Theorem 1 states that infact these different dimensional representations are not acting on independent spaces, they are allacting on different ﬁnite truncations of the same inﬁnite-dimensional function space.
While Theorem1 only pertains to the backward (i.e. decoder) part of a U-net, a forward (i.e. encoder) part F can bestraightforwardly deﬁned as the reversal for all time t ∈ [0 , T ] .
Here, F and B should be thought ofas the forward and backward passes of an inﬁnitely deep U-Net architecture [1].
We stay within V − Jas, in practice, we do not work in inﬁnite dimensions, and so the diffusion above is truncated to some
T ∈ (0 , 1) such that the obtained resolution space V − J is ﬁnite.
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 inthe regularisation they impose.
In practice, a pixel representation of images is used when training
U-Nets, which we henceforth call the standard basis (see Appendix A).
The standard basis is notconvenient to derive theoretical results.
It is instead preferable to use a wavelet basis natural to themulti-resolution partition imposed by a U-Net projection operation, which for average pooling isthe Haar basis (see Appendix A).
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 idea in Theorem 2 whichstates that the dimension reduction operation of average pooling in the standard basis is a conjugateoperation to co-ordinate projection within the Haar basis (details are provided in Appendix A).
Theorem 2.
Given V − j as in Deﬁnition 1, let x ∈ V − j be represented in the standard basis E j and Haar wavelet basis Ψ j .
Let π j : E j (cid:55)→ Ψ
j be the change of basis map illustrated in Fig.
1, then wehave the conjugacy π j − 1 ◦ [Equation] [Equation]
[Figure]
[Equation] V [Figure]
[Equation]
[Figure] [Equation] [Figure] [Equation]
[Figure]
[Equation]
[Figure] [Equation] [Figure] [Equation]
[Figure]
Theorem 2 means that if we project an image from V jto
V j +1 in the Haar wavelet basis, we can alternativelyview this as changing to the standard basis via π − 1 j , performing average pooling, and reverting back via π j +1 (see
Figure 1).
This is important because the Haar basis isorthonormal, allowing us to precisely quantify what information is lost when average pooling is used to reduce thedimension within a U-Net in Theorem 3.
Theorem 3.
Let { V − j } Jj =0 be a multi-resolution partition of V − J with the orthogonal decompositions V − j = V − j +1 ⊕
U − j +1 .
Let F j,φ , B j,θ : D ( V − j )
(cid:55)→ D ( V − j ) be bounded linear operators with 157 parameters φ
and θ , for instance corresponding to a neural network.
Deﬁne F j,φ : D ( V − j ) (cid:55)→
D ( V − j +1 ) and let its restricted reversal be B j,θ : D ( V − j +1 )
(cid:55)→ D ( V − j ) given through F j,φ :=
proj
V − j +1 ◦ F j,φ and B j,θ := B j,θ ◦ embd V − j 3 with compositions F j 1 | j 2 ,φ := F j 1 ,φ ◦ · · · ◦ F j 2 ,φand analogous B j 1 | j 2 ,θ .
Then the sequence { B 1 | j,θ ( F 1 | J,φ ν J ) }
Jj =0 forms a discrete multi-resolutionbridge between F 1 | J,φ ν J and B 1 | J,θ F 1 | J,φ ν J at times { t j } Jj =1 .
Further, if E : D ( V − J ) (cid:55)→ V − J is a 162 bounded linear operator, then [Equation] [Equation] [Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table] [Equation] [Table] 2 , (1) where X t j := E ( F j,φ F J | j − 1 ,φ δ { x tJ } ) .
Theorem 3 states that U-Nets push the latent data distribution to a ﬁnite multi-resolution basis,speciﬁcally a Haar basis when average pooling is used.
To see this, we relate the right-handside of Eq. 1 to the L 2 reconstruction loss by selecting E to be the expected position operatorν (cid:55)→ (cid:82) xdν ( x ) ∈ L 2 ( X ) .
The left-hand side of Eq.
(1) represents what percentage of our dataenters the orthogonal subspaces { U − j } Jj =0 which are discarded when using a U-Net architecturewith average pooling.
Theorem 3 thus shows that as we minimise the reconstruction error duringtraining, we minimise the percentage of our data transported to the orthogonal sub-spaces { U − j } Jj =0 .
Consequently, the U-Net architecture implicitly trains our network to learn a wavelet basis of L 2 ( X ) .
This characterises the regularisation imposed by a U-Net.
These results also imply that U-Nets with average pooling provide a direct alternative to Fourierfeatures [14, 26, 27, 28] which impose a Fourier basis, an alternative orthogonal basis on L 2 ( X ) .
However, unlike Fourier bases, there are in fact a multitude of wavelet bases which are all encompassed by our framework, and Theorem 3 pertains to all of them.
This opens the door to exploringconjugacy operations beyond average pooling induced by other, potentially learned, wavelet basesoptimised for speciﬁc data types.
2.3 Example: HVAEs are Sum Representation Diffusion Discretisations
To show what practical inferences we can derive from our multi-resolution framework, we applyit to analyse state-of-the-art HVAE architectures (see Appendix C for an introduction), identifying [Figure] increases parameter redundancies and instabilities.
Here and in our experiments, we focus on VDVAEs [6].
Weprovide similar results for Markovian HVAEs [29, 30] and NVAEs [7] (see § 4) in Appendix A.
We start by inspecting VDVAEs.
As we show next, we can tie the computations in VDVAEs to theforward and backward operators F j,φ and B j,θ within our framework and identify them as a two-stepforward Euler discretisation of a diffusion sum process.
When used with a U-Net, as in a VDVAE,this creates a multi-resolution diffusion bridge by Theorem 4:
Theorem 4.
Let t J :=
T ∈ (0 , 1) , and consider B θ, 1 | J : D ( V − J ) (cid:55)→ D ( V 0 ) and its reversal
F φ,J | 1 : D ( V − J ) (cid:55)→ D ( V 0 ) given through the coupled multi-resolution Markov sum process in thestandard basis:
[Equation]
where Y 0 = 0 , dY t =
Z t dt, (cid:107) Y t (cid:107) 2 > (cid:107) Y s (cid:107) 2 with 0 ≤ s < t ≤ T and for a measure ν J ∈ D ( V − J )
we have X T , Z [-
 ⇠-]0 ∼ F φ,J | 1 ν J .
Then, VDVAEs approximate this process, and its residual cells are atwo-step forward Euler discretisation of these coupled Stochastic Differential Equations (SDEs).
To better understand Theorem 4, we visualise its residual cell structure and the correspondingdiscretisation steps of VDVAEs, NVAEs and Markovian HVAEs in Appendix A. Note that thisprocess is Markov on the sum of the variables
Z i , not the variables themselves, which has twoimmediate beneﬁts over a regular diffusion: First , similar sum processes have been empiricallyobserved as efﬁcient ﬁrst-order approximates to higher-order chains, for example the memory state in LSTMs [31].
Further, VDVAEs and NVAEs are even claimed to be high-order chains (Eqs. (2,3) in[6] and Eq.
(1) in [7]), despite only approximating this with a sum process.
Second , the discretisationof the sum process has better stability under the assumption that the Y i are increasing in magnitude.
This is because the two-step splitting has a relative forward error which grows with (cid:107) Z ( σ ) i +1 (cid:107) 2 / (cid:107)
Z i (cid:107) 2in the regular diffusion case, whilst in the sum process diffusion, it grows with (cid:107) Z ( σ ) i +1 (cid:107) 2 / (cid:107)
Y i (cid:107) 2 .
Inthe latter, the growth of the process increases the denominator and controls this error.
To show how VDVAEs impose the growth of the Y t sum process, we prove that the forward passof a VDVAE’s U-Net enforces Y 0 = 0 .
This is done by identifying that the measure ν 0 , which a
VDVAE connects to the data ν ∞ via a multi-resolution bridge, is a point mass on the zero function.
Consequently the backward pass must grow to reverse this, and the network learns this in a monotonicmanner as we later conﬁrm in our experiments (see §3.2).
Theorem 5.
Consider the coupled SDEs in Eq.
(2) , trained through the reconstruction error in Eq.
C.38.
Let ˜ ν
J denote the data measure and ν
= δ { 0 } be the initial multi-resolution bridge measureimposed 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 index 214 indicates the joint distribution.
Further, the relative backwards error of step i in the two-step forward
Euler discretisation is ε i {++
+} =
(cid:107) Z ( σ ) i +1 (cid:107) 2 / (cid:107) Y i (cid:107) 2 . [Figure]
|| y i
|| 2 w i t h i n p u t y i o f t
o p - d o w n b l o c k 1 x 1 4 x 4 8 x16 x
[Figure]
32 x Theorem 5 states that the VDVAE architecture forms multi-resolution bridge with the dynamics of Eq.(2) , and connects our data distribution to the trivial measure on V 0 : a Dirac mass at 0 .
In the following,we now gain two insights into this HVAE cell on instabilities and on parameter redundancies.
Thereare two major instabilities in this discretisation.
First, the imposed ν 0 is disastrously unstable as itenforces a data set, with potentially complicated topology to completely collapse at t = 0 , and weobserve the consequential sampling instability in our experiments in §3.3.
Second, the discretisationerror ε i at each step of B θ is not explicitly imposed in VDVAEs or NVAEs, but a band-aid solutionis to control this by ensuring Y i t → 1 −−−→ ∞ sufﬁciently fast.
The VDVAE architecture does, however,bolster this rate through the Z ( σ ) i, + term, which is absent in NVAEs [7], in the discretisation steps ofthe residual cell (see Appendix A).
We empirically observe this controlled backward error in §3.2.
Moreover, the current form of VDVAEs is over-parameterised and not informed by this continuoustime formulation.
The continuous time analogue of VDVAEs [6] in Theorem 4 has time dependentcoefﬁcients −→ µ t , −→ σ t , ←− µ t , ←− σ t .
We hypothesise that the increasing sum process Y t implicitly encodestime.
Hence, explicitly representing this in the model, for instance via ResNet blocks with independentparameterisations at every time step, is redundant, and a time-homogeneous model (see Appendix Afor a precise formulation)— practically speaking, performing weight-sharing across time steps/layers—has the same expressivity, but requires far fewer parameters than the state-of-the-art VDVAE.
Experiments
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 [6] as the basis architecture on ﬁve datasets:
MNIST [32], CIFAR10 [33], two downsampled versions of ImageNet [34, 35], and CelebA [36],splitting each into a training, validation and test set (see Appendix E for details).
In general, reportednumeric values refer to Negative Log-Likelihood (NLL) in nats (MNIST) or bits per dim (all otherdatasets) on the test set at model convergence, if not stated otherwise.
We note that performanceon the validation and test set is generally very similar.
Appendices G, H and I deﬁne the HVAEmodels 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 | (cid:126) z ) , present experimentaldetails, and additional experimental results, respectively.
We provide our PyTorch code base at<see-supplementary-material>
(see Appendix D 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 modelwith 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.
2, we train HVAEs on MNIST and CIFAR10 with { 1 , 2 , 3 , 4 } ResNet blocks (referred to as { x1 , x2 , x3 , x4 }) ineach 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.
We compare our runs of VDVAE with original hyperparameters [6] ( VDVAE ∗ ) against our weight-shared VDVAEs ( 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 [14] as the code is not public and the manuscript does not specify it.
[Table]
I m ag e N e t ×
WS-VDVAE (ours)
VAE 55m ≤ 3 .
WS-VDVAE (ours)
VAE 85m ≤ 3 . VDVAE ∗ (ours) VAE 119m ≤ 3 . NVAE [7]
VAE 268m ≤ 3 .
VDVAE [6] VAE 119m ≤ 3 . VDM [14] Diff – ≤ 3 .
C e l e b A ×
WS-VDVAE (ours)
VAE 75m ≤ 2 .
VDVAE ∗ (ours) VAE 125m ≤ 2 . NVAE [7] VAE 153m ≤ 2 .
As hypothesised by our framework, yet very surprising in
HVAEs, NLL after 1m iterations measured on the validationset gradually increases the moreoften blocks are repeated eventhough all weight-sharing models have an identical parametercount to the x 1 model (MNIST:107k, CIFAR10: 8.7m).
Furthermore, the weight-sharing modelsoften outperform or reach equal
NLLs compared to x2 , x3 , x4 , allof which have more parameters(MNIST: 140k; 173k; 206k.
CIFAR10: 13.0m; 17.3m; 21.6m),yet fewer
activations, latent variables, and number of timestepsat which the coupled SDEs in Eq.(2) are discretised.
We now scale these ﬁndingsup to large-scale hyperparameter conﬁgurations.
We train VDVAEs closely following the stateof-the-art hyperparameter conﬁgurations in [6], speciﬁcallywith the same number of parameterised blocks and withoutweight-sharing ( VDVAE ∗ ), andcompare them against modelswith weight-sharing ( WS-VDVAE )and fewer parameters, i.e. fewerparameterised blocks, in Table 1.
On all four datasets, the weight-shared models achieve similar NLLs with fewer parameters comparedto their counterparts without weight-sharing: We use 32% , 36% , 54% , and 40% less parameters onthe four datasets reported in 1, respectively.
For the larger runs, weight-sharing has diminishingreturns on NLL as these already have many discretisation steps.
To the best of our knowledge, ourmodels 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 of57, 105, 235, 125, respectively, the highest ever trained.
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 HVAEsmore parameter efﬁcient.
We now want to explain and understand this behavior further.
In Fig.
3,we measure (cid:107) Y i (cid:107) 2 , the L 2 -norm of the residual state at every bottom-up [left] and top-down [right]block with index i, over several batches for large-scale models with weight-sharing trained on MNIST(see Appendix I.2 for similar results on CIFAR10 and ImageNet32).
On average, we experience anincrease in the state norm across time in every resolution, interleaved by discontinuous ‘jumps’ at theresolution transitions (projection or embedding) where the dimension of the residual state changes.
This supports our claim in §2 that HVAEs are ‘sum-processes’ Y i = (cid:80) j =
Z j , and hence learn torepresent a notion of time in their residual state.
It is now straightforward to ask how HVAEs beneﬁt from this sum process representation duringtraining: As we show in Table 2, when normalising the state norm at every forward and backward block during training, i.e. forcing a “ﬂat line” in Fig.
3.[Left], the learning deteriorates after a shortwhile, resulting in very poor NLL results compared to the runs with a regular, non-normalised residualstate above.
This evidence conﬁrms our earlier stated hypothesis that the sum process representationin ResNet-based HVAEs encodes important information which recent HVAEs heavily rely on duringlearning.
3.3 Sampling instabilities in HVAEs
[Table]
High ﬁdelity unconditional samples of faces, e.g. from modelstrained on CelebA, cover the front pages of state-of-the-art HVAEpapers [6, 7].
Here, we question whether face datasets are an appropriate benchmark for HVAEs.
In Theorem 5, we identiﬁed theaforementioned state-of-the-art HVAEs as ﬂows to a point mass,hypothesising instabilities during sampling.
And indeed, whensampling from our trained VDVAE ∗ with state-of-the-art conﬁgurations, we observe high ﬁdelity and diversity samples on MNISTand CelebA, but unrecognisable, yet diverse samples on CIFAR10,
ImageNet32 and ImageNet64, in spite of state-of-the-art test set
NLLs (see Fig.
and Appendix I.3).
We argue that MNIST and
CelebA, i.e. numbers and faces, have a more uni-modal nature, andare in this sense easier to learn for a discretised multi-resolutionprocess ﬂowing to a point mass, which is uni-modal, than the other“in-the-wild”, multi-modal datasets.
Trying to approximate the latterwith the, in this case unsuitable, HVAE model leads to the samplinginstabilities observed.
3.4 Ablation studies
[Table]
We conducted several ablation studies which support ourexperimental results and further probe our multi-resolutionframework for HVAEs.
In this section we note keyﬁndings—a detailed account of all ablations can be foundin Appendix H.4.
In particular, we ﬁnd that the numberof latent variables, which correlates with stochastic depth,does not explain the performance observed in §3.1, supporting our claims.
We show that Fourier features do notprovide a performance gain in HVAEs, in contrast to stateof-the-art diffusion models, where they are essential forstate-of-the-art performance [14].
This is consistent withour framework’s ﬁnding that a U-Net architecture withpooling is already forced to learn a Haar wavelet basisrepresentation of the data, hence introducing another basisdoes not add value.
Furthermore, we demonstrate thatresidual cells are crucial for the performance of HVAEs asthey are able to approximate the dynamics of a diffusionsum process. We also investigate synchronous vs. asynchronous processing in time between theforward and backward pass.
We conclude this section by highlighting a key result which underlinesthe beneﬁt of a multi-resolution bridge compared to a single-resolution model.
On the effect of a multi-resolution bridge.
State-of-the-art HVAEs have a U-Net architecturewith pooling and, hence, are multi-resolution bridges (see Theorem 4).
We investigate the effectof multiple resolutions in HVAEs (here with spatial dimensions { 32 2 , 16 2 , 8 2 , 4 2 , 1 2 } ) against asingle resolution (here with spatial dimension 32 2 ).
We choose the number of blocks for the singleresolution model such that they are distributed in the encoder and decoder proportionally to the multiresolution model and the total number of parameters are equal in both, ensuring a fair comparison.
As we show in Table 3, the multi-resolution models outperform their single-resolution counterparts,while additionally having a representational advantage due to their Haar wavelet basis representation(illustrated in Appendix I.4).
This underlines the importance of the inductive bias a multi-resolutionbridge imposes, as described by our framework, relative to a single-resolution model.
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’sbottleneck.
These connections also smooth out the network’s loss landscape [37].
In the literature, U-Nets tend to be convolutional, and a wide range of different approaches have been usedfor up-sampling and down-sampling between resolutions, with many using average pooling for thedown-sampling operation [8, 9, 11, 12, 14].
In this work, we focus on U-Nets as linear operatorsinterleaved by average pooling as the down-sampling operation, and we formally characterise U-Netsas multi-resolution bridges in Theorem 3.
Prior to our work, the dimensionality-reducing bottleneck structure of U-Nets was widely acknowledged as being useful, however it was unclear whatregularising properties a U-Net imposed.
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 [38], while achieving stable training dynamics that avoid commonissues such as posterior collapse [29, 39] or exploding gradients.
Markovian HVAEs condition eachlatent variable directly on only the previous one by taking samples forward [29, 30].
Markovian
HVAEs suffer from stability issues even for very small stochastic depths.
Nouveau VAEs (NVAEs)
[7]and Very Deep VAEs (VDVAEs) [6] combine the improvements of several earlier HVAE models(see Appendix C for details), while scaling up to larger stochastic depths.
Both use ResNet-basedbackbones, sharing parameters between the generative and recognition parts of the model.
VDVAEsare the considerably simpler approach, in particular avoiding common tricks such as a warm-updeterministic autoencoder training phase or data-speciﬁc initialisation.
VDVAEs achieve a stochasticdepth of up to 78, improving performance with more ResNet blocks.
Worth noting is that while
Markovian HVAEs and NVAEs use convolutions with appropriately chosen stride to jump betweenresolutions, VDVAEs use average pooling.
In all HVAEs to date, a theoretical underpinning whichexplains architectural choices, for instance the choice of residual cell, is missing.
Conclusion
In this work, we introduced a multi-resolution framework for U-Nets.
We provided theoretical resultswhich uncover the regularisation property of a U-Net with average pooling as implicitly learning a Haar wavelet representation of the data.
We then applied our framework to HVAEs as an example,identifying them as multi-resolution diffusion sum processes ﬂowing to a point mass, and observedparameter redundancies and instabilities.
We veriﬁed the latter theoretical insights in both smalland large-scale generative modelling experiments, and in doing so trained the deepest ever HVAEs.
We explained these results by showing that HVAEs learn a representation of time and performedextensive ablation studies.
An important limitation is that the proven regularisation property of U-Nets is limited to using averagepooling 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 lackof exhaustive hyperparameter optimisation performed is a limitation of the work as it may be possibleto obtain improved results.
We demonstrate, however, that simply adding weight-sharing to thehyperparameter settings given in the original VDVAE paper [6] leads to state-of-the-art performancewith improved parameter efﬁciency, and hence view it as a strength of our results.
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 F.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
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 Appendix A 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.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental 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 G and E. [Table]
(c) Did you report error bars (e.g., with respect to the random seed after running exper-
[Table]
iments multiple times)?
[No] Due to the signiﬁcant computational cost of training [Table]
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
[Table]
hyperparameter setting required to provide error bars for our runs.
We note that this is
common practice in large-scale HVAE research (compare [6, 7]).
Furthermore, in their
[Table] code base, VDVAE [6], which we directly base our architecture on, reported highly [Table]
stable test NLL when varying the random seed, varying only in the second decimal
[Table]
place in bits per dim.
[Table] (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
[Table] our results brieﬂy in §3, and in detail in Appendix D. [Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes] We discuss all existingassets used in our code base in Appendix D. [Table]
(b) Did you mention the license of the assets?
[Yes] We mention the license of all assets [Table] used in Appendix D. [Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
[Table]
We provide our source code with instructions on how to reproduce our results in the
[Table]
supplementary material.
Further details on our code are provided in Appendix D.
[Table] (d) Did you discuss whether and how consent was obtained from people whose data you’re [Table]
using/curating?
[Yes] We address this question in Appendix E.
(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 E.
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]