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
B).
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 and Fig. A.1 for a detailed discussion).
However, a theoretical understanding of these choicesis lacking.
For instance, it has not been shown why a residual backbone may be beneﬁcial, or what thespeciﬁc cell structures 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 contributions areas follows: (a) We provide a multi-resolution framework for U-Nets.
We formally deﬁne U-Netsas multi-resolution bridges over a hierarchy 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.
Our framework both allows us to understandthe empirical 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 instabilitiesand train HVAEs with the largest stochastic depth ever, achieving state-of-the-art performance withhalf the number of parameters by exploiting our theoretical insights.
We explain these results byuncovering that HVAEs secretly represent time in their state and show that they use this informationduring training.
We ﬁnally provide extensive ablation studies which, for instance, rule out otherpotential factors which correlate with stochastic depth, show the empirical gain of multiple resolutions,and ﬁnd that Fourier features (a key performance driver in discrete-time diffusion models [14]) donot 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.
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 istherefore to have to take as few steps away as possible.
The reader is encouraged to physically conductthis experiment with the examples provided in Appendix B.1. We can formalise the aforementionedintuitions by deﬁning a multi-resolution hierarchy [24] of sub-spaces we may project to:
Deﬁnition 1.
[ Daubechies (1992) [24]] 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 functionspace L 2 ( R m ) if: ( A1) S
[Equation]
[-⊂ V 1 ⊂ V 
 ⊂-]=
[Equation]
[Equation]
R m ) ; ( A2) T [Equation]
= [Equation] ; ( A3) f [Equation]
[Equation] j , f [Equation] [Equation]
[Equation] 0 ;( 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 85 of L 2 (
X ) is { V j } j 2 Z 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 negativej corresponds to ‘zooming in’ on the images.
The original image 2 is a member of L 2 ([0 , 1] 2 ) , thespace of square-integrable functions on the unit square, and its ﬁnite projection to 2 j · 2 j many pixelsis a member of V  j .
The nesting property V  j +1 ⇢ V  j ensures that any image with (2 j  1 ) 2 pixelscan also be represented by (2 j ) 2 pixels, but at a higher resolution.
Assumption ( A1 ) states that withinﬁnitely many pixels, we can describe any inﬁnitely detailed image.
In contrast, ( A2 ) says that withno pixels, we cannot approximate any images.
Assumptions ( A3 ) and ( A4 ) allow us to form a basisfor images in any V j if we know the basis of V 0 .
One basis made by extrapolating from V [-
 inthis way is-]0 in this wayis known as a wavelet basis [24].
Wavelets have proven useful for representing images, for instancein the JPEG standard [25], 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 this at a lower resolution.
Similar to how we did for inﬁnitelydetailed images, we want to project the measure ⌫ 1 to a lower dimensional measure ⌫ j on the ﬁnitedimensional space V  j .
In extension to this, we want the ability to reverse this projection so that wemay sample from the lower dimensional measure and create a generative model for ⌫ 1 .
We wouldlike to again prioritise the presence of large-scale features of the original image within the lowerdimensional samples.
We do this by constructing a multi-resolution bridge from ⌫ 1 to ⌫ j , as deﬁnedbelow.
Deﬁnition 2.
Let X be a compact subset of R m and { V  j } 1 j =0 be a multi-resolution partition of
L 2 ( X ) scaled such that L 2 ( X ) =
[Equation] j 2 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 2 [0 , 1] on L 2 ( X ) is a 107 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
[ t j , t j +1 ) implies supp ( ⌫ t ) ⇢ V  j , i.e ⌫ t 2 D ( V  j ) ; and,
(ii) for s 2 (0 , 1) , the mapping s 7!
⌫ s is continuous (in the weak sense, see Appendix A.1, Eq. (A.5)) if s < t j .
The continuous time dependence in Deﬁnition 2 plays a movie of the measure ⌫ 0 supported on V 0growing 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 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 a 116 multi-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 2 (0 , 1) and in the standard basis discussed in §2.2,which will be useful in when viewing HVAEs as discretisations of continuous-time multi-resolutiondiffusion models 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 diffusiontransition 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
B :=
J
X j = [ t j ,t j +1 ) · [Equation]
j , µ : [Equation]
J [Equation]
[Equation] ) ,induce a multi-resolution bridge of measures from the dynamics for t 2 [0 , T ] and on the standardbasis as dX t = µ t ( X t ) dt +  t ( X t ) dW t (see Appendix A.4 for the details of this integration) for
X t 2 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 be straightforwardly deﬁned as the reversal for all time t 2 [0 , T ] .
Here, F and B should be thought ofas the forward and backward passes of an inﬁnitely deep U-Net architecture [1].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 A.2, Eq. (A.9)).
The standard basis is notconvenient to derive theoretical results.
It is instead preferable to use a basis natural to the multiresolution bridge imposed by a U-Net with corresponding projection operation, which for averagepooling is the Haar (wavelet) basis [26] (see Appendix A.2).
The Haar basis, like a Fourier basis, isan orthonormal basis of L 2 ( X ) which has desirable L 2 -approximation properties.
We formalise thisin Theorem 2 which states that the dimension reduction operation of average pooling in the standardbasis 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 148 Haar basis j .
Let ⇡ j : E j 7!
j be the change of basis map illustrated in Fig.
1, then we have theconjugacy ⇡
[Equation]
[Equation]
[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 hierarchy 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 bounded forward/backward kernel operators 160 with parameters  and ✓ . Deﬁne F [Equation]
by
[Equation] : D
[Equation] D ( [Equation]
[Equation] +1 ) 161 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 ) }
Jj =0 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
[Equation]
where W 2 is the Wasserstein- 2 metric and
F j | J,   2 is the Lipschitz constant of F j | J,  .
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 RHS of Eq.(A.65) to the L 2 -reconstruction loss by noting this quantity is bounded by it.
The LHS of Eq.
(A.65)represents what percentage of our data enters the orthogonal subspaces { U  j } Jj =0 which are discarded 169 when using a U-Net architecture with average pooling.
Theorem 3 thus shows that as we minimisethe reconstruction error during training, we minimise the percentage of our data transported to theorthogonal sub-spaces { U  j } Jj =0 .
Consequently, the U-Net architecture implicitly trains our networkto learn a Haar 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, 27, 28, 29] 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 in particular Theorem 3 pertains to all of them.
This opens the doorto exploring conjugacy operations beyond average pooling induced by other, potentially learned,wavelet bases optimised 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 apply itto 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 [6].
Weprovide similar results for Markovian HVAEs [30, 31] and NVAEs [7] (see § 4) in Appendix A.5.
We start by inspecting VDVAEs.
As we show next, we can tie the computations in VDVAE cellto the forward and backward operators F j,  and B j, ✓ within our framework and identify them as atwo-step forward Euler discretisation of a diffusion sum process.
When used with a U-Net, as is donein VDVAE [6], this creates a multi-resolution diffusion bridge by Theorem 4:
Theorem 4.
Let t J :=
T 2 (0 , 1) , and consider B ✓ , 1 | J : D ( V  J ) 7! D ( V 0 ) and its reversal
F  ,J | 1 : D ( V  J ) 7! D ( V 0 ) given through the coupled multi-resolution Markov sum process in the 190 standard basis:
[Equation]
where Y 0 = 0 , dY t =
Z t dt, k Y t k 2 > k Y s k 2 with 0  s < t  T and for a measure ⌫ J 2 D ( V  J ) 192 we have X
T , Z {+
 ⇠+} F  ,J | 1 ⌫ J .
Then, VDVAEs approximate this process, and its residual cells are a 193 two-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.5, Fig. A.1. Notethat this process is Markov on the sum of the variables
Z i , not the variables themselves, which hastwo immediate 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 statein LSTMs [32].
Further, VDVAEs and NVAEs are even claimed to be high-order chains (see Eqs.(2,3) in [6] and Eq.
(1) in [7]), despite only approximating this with a sum process.
Second , thediscretisation of the sum process has better stability under the assumption that the Y i are increasingin magnitude.
This is because the two-step splitting has a relative forward error which grows withk Z (  ) i +1 k 2 / k
Z i k 2 in the regular diffusion case, whilst in the sum process diffusion, it grows withk Z (  ) i +1 k 2 / k Y
i k 2 .
In the latter, the growth of the process increases the denominator and controls thiserror.
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 ⌫ 1 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.
B.107.
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 215 indicates the joint distribution.
Further, the relative backwards error of step i in the two-step forward
Euler discretisation is " i [-+
-]+1 =
k Z (  ) i +1 k 2 / k Y
i k 2 .
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 .
Fromthis insight, we can draw conclusions on instabilities and on parameter redundancies of this HVAEcell.
There are two major instabilities in this discretisation.
First, the imposed ⌫ 0 is disastrouslyunstable as it enforces a data set, with potentially complicated topology to completely collapse att = 0 , and we observe the resulting sampling instability in our experiments in §3.3.
We note thatsimilar arguments are applicable in settings without a latent hierarchy imposed by a U-Net, see forinstance [33].
Second, the discretisation error " i at each step of B ✓ is not explicitly imposed in VDVAEs or NVAEs, but a band-aid solution is to control this by ensuring Y i t ! 1 ! 1 sufﬁcientlyfast.
The VDVAE architecture does, however, bolster this rate through the Z (  ) i, + term, which is absentin NVAEs [7], in the discretisation steps of the residual cell.
We empirically observe this controlledbackward error in Fig. 3 [Right]. We refer to Fig.
A.1 for a detailed comparison of HVAE cells andtheir corresponding discretisation of the coupled SDEs in Eq.
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 A.7 for a precise formulation)—practically speaking, performing weight-sharing across time timesteps/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 model on ﬁve datasets:
MNIST [34], CIFAR10 [35], two downsampled versions of ImageNet [36, 37], and CelebA [38],splitting each into a training, validation and test set (see Appendix D 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 F, G and H 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 | ~ z ) , present experimentaldetails, and additional experimental results, respectively.
We provide our PyTorch code base at<see-supplementary-material>
(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 modelwith time-dependent coefﬁcients, yet uses much less parameters.
We start demonstrating this ef- fect by weight-sharing ResNet blocks across time on a small scale.
In Fig.
2, we train HVAEson 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 ⇥
[Table]
[Table]
C e l e b A ⇥
[Table]
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 Table 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, wemeasure k Y i k 2 , the L 2 -norm of the residual state at every backward/top-down block with index i, overseveral batches for large-scale models trained on MNIST (see Appendix H.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 theresidual state changes.
This supports our claim in §2 that HVAEs are ‘sum-processes’ Y i = P j =
Z j ,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 sum process representation duringtraining: As we show in Table 2, when normalising the state norm at every forward and backwardblock during training, i.e. forcing a “ﬂat line” in Fig.
3 [Left], learning deteriorates after a short while,resulting in poor NLL results compared to the runs with a regular, non-normalised residual state.
Thisevidence conﬁrms our earlier stated hypothesis that the sum process representation in ResNet-based HVAEs encodes important information which recent HVAEs heavily rely on during learning.
3.3 Sampling instabilities in HVAEs
[Table]
High ﬁdelity unconditional samples of faces, .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 H.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 G.4.
In particular, we ﬁnd that the number of latent variables, which correlates with stochasticdepth, 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, incontrast to state-of-the-art diffusion models, where theyare essential for state-of-the-art performance [14].
Thisis consistent with our framework’s ﬁnding that a U-Netarchitecture with pooling is already forced to learn a Haarwavelet basis representation of the data, hence introducinganother basis does not add value.
We also demonstrate thatresidual cells are crucial for the performance of HVAEs asthey are able to approximate the dynamics of a diffusionsum process, and investigate synchronous vs. asynchronous processing in time between the forwardand backward pass.
We conclude this section by highlighting a key result which underlines the beneﬁtof 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 H.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 [39].
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 forthe down-sampling operation [8, 9, 11, 12, 14].
In this work, we focus on U-Nets as linear operators on measures interleaved by average pooling as the down-sampling operation, and we formallycharacterise U-Nets as multi-resolution bridges in Theorem 3.
Prior to our work, the dimensionalityreducing bottleneck structure of U-Nets was widely acknowledged as being useful, however it wasunclear 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 [40], while achieving stable training dynamics that avoid commonissues such as posterior collapse [30, 41] or exploding gradients.
Markovian HVAEs condition eachlatent variable directly on only the previous one by taking samples forward [30, 31].
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 B 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, and we providedthis in Section §2.3.
Conclusion
In this work, we introduced a multi-resolution framework for U-Nets.
We provided theoreticalresults which uncover the regularisation property of U-Nets with average pooling as implicitlylearning a Haar wavelet representation of the data.
We applied our framework to HVAEs, identifyingthem as multi-resolution diffusion sum processes ﬂowing to a point mass, and observed parameterredundancies and instabilities.
We veriﬁed the latter theoretical insights in both small- and large-scaleexperiments, and in doing so trained the deepest ever HVAEs.
We explained these results by showingthat 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 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’scontributions and scope?
[Yes] We state our theoretical and experimental contributionsin §1.
We develop our multi-resolution framework with corresponding Theorems 1 to5 in §2 (proofs in Appendix A).
We provide our experimental results in §3.
(b) Did you describe the limitations of your work?
[Yes] See
(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]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes We stated theassumptions of all theoretical results in
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.
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 F and D.
(c) Did you report error bars (e g., with respect to the random seed after running exper-
iments 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 [6, 7]).
Furthermore, in their
code base, VDVAE [6], 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
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.
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]