Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images

We present a hierarchical VAE that , for the first time , generates samples quickly and outperforms the PixelCNN in log-likelihood on all natural image benchmarks . We begin by observing that , in theory , VAEs can actually represent autoregressive models , as well as faster , better models if they exist , when made sufficiently deep . Despite this , autoregressive models have historically outperformed VAEs in loglikelihood . We test if insufficient depth explains why by scaling a VAE to greater stochastic depth than previously explored and evaluating it CIFAR-10 , ImageNet , and FFHQ . In comparison to the PixelCNN , these very deep VAEs achieve higher likelihoods , use fewer parameters , generate samples thousands of times faster , and are more easily applied to high-resolution images . Qualitative studies suggest this is because the VAE learns efficient hierarchical visual representations . We release our source code and models at https : //github.com/openai/vdvae . 1 INTRODUCTION . One potential path to increased data-efficiency , generalization , and robustness of machine learning methods is to train generative models . These models can learn useful representations without human supervision by learning to create examples of the data itself . Many types of generative models have flourished in recent years , including likelihood-based generative models , which include autoregressive models ( Uria et al. , 2013 ) , variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) , and invertible flows ( Dinh et al. , 2014 ; 2016 ) . Their objective , the negative log-likelihood , is equivalent to the KL divergence between the data distribution and the model distribution . A wide variety of models can be compared and assessed along this criteria , which corresponds to how well they fit the data in an information-theoretic sense . Starting with the PixelCNN ( Van den Oord et al. , 2016 ) , autoregressive models have long achieved the highest log-likelihoods across many modalities , despite counterintuitive modeling assumptions . For example , although natural images are observations of latent scenes , autoregressive models learn dependencies solely between observed variables . That process can require complex function approximators that integrate long-range dependencies ( Oord et al. , 2016 ; Child et al. , 2019 ) . In contrast , VAEs and invertible flows incorporate latent variables and can thus , in principle , learn a simpler model that mirrors how images are actually generated . Despite this theoretical advantage , on the landmark ImageNet density estimation benchmark , the Gated PixelCNN still achieves higher likelihoods than all flows and VAEs , corresponding to a better fit with the data . Is the autoregressive modeling assumption actually a better inductive bias for images , or can VAEs , sufficiently improved , outperform autoregressive models ? The answer has significant practical stakes , because large , compute-intensive autoregressive models ( Strubell et al. , 2019 ) are increasingly used for a variety of applications ( Oord et al. , 2016 ; Brown et al. , 2020 ; Dhariwal et al. , 2020 ; Chen et al. , 2020 ) . Unlike autoregressive models , latent variable models only need to learn dependencies between latent and observed variables ; such models can not only support faster synthesis and higher-dimensional data , but may also do so using smaller , less powerful architectures . We start this work with a simple but ( to the best of our knowledge ) unstated observation : hierarchical VAEs should be able to at least match autoregressive models , because autoregressive models are equivalent to VAEs with a powerful prior and restricted approximate posterior ( which merely outputs observed variables ) . In the worst case , VAEs should be able to replicate the functionality of autoregressive models ; in the best case , they should be able to learn better latent representations , possibly with much fewer layers , if such representations exist . We formalize this observation in Section 3 , showing it is only true for VAEs with more stochastic layers than previous work has explored . Then we experimentally test it on competitive natural image benchmarks . Our contributions are the following : • We provide theoretical justification for why greater depth ( up to the data dimension D , but also as low as some value K D ) could improve VAE performance ( Section 3 ) • We introduce an architecture capable of scaling past 70 layers , when previous work explored at most 30 ( Section 4 ) • We verify that depth , independent of model capacity , improves log-likelihood , and allows VAEs to outperform the PixelCNN on all benchmarks ( Section 5.1 ) • Compared to the PixelCNN , we show the model also uses fewer parameters , generates samples thousands of times more quickly , and can be scaled to larger images . We show evidence these qualities may emerge from the model learning an efficient hierarchical representation of images ( Section 5.2 ) • We release code and models at https : //github.com/openai/vdvae . 2 PRELIMINARIES . We review prior work and introduce some of the basic terminology used in the field . 2.1 VARIATIONAL AUTOENCODERS . Variational autoencoders ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) consist of a generator pθ ( x|z ) , a prior pθ ( z ) , and an approximate posterior qφ ( z|x ) . Neural networks φ and θ are trained end-to-end with backpropagation and the reparameterization trick in order to maximize the evidence lower bound ( ELBO ) : log pθ ( x ) ≥ Ez∼qφ ( z|x ) log pθ ( x|z ) −DKL [ qφ ( z|x ) ||pθ ( z ) ] ( 1 ) See Kingma & Welling ( 2019 ) for an in-depth introduction . There are many choices for what networks are used for pθ ( x|z ) , qφ ( z|x ) , and whether pθ ( z ) is also learned or set to a simple distribution . We study VAEs with independent pθ ( x|z ) – that is , where each observed xi is output without conditioning on any other xj . This ensures generation time does not increase linearly with the dimensionality of the data , and requires that these VAEs learn to incorporate the complexity of the data into a rich distribution over latent variables z . It is possible to have autoregressive pθ ( x|z ) ( Gulrajani et al. , 2016 ) , but generation is slow for these models . They also sometimes ignore latent variables entirely , becoming equivalent to normal autoregressive models ( Chen et al . ( 2016 ) ) . 2.2 HIERARCHICAL VARIATIONAL AUTOENCODERS . Much of the early work on VAEs incorporate fully-factorized Gaussian qφ ( z|x ) and pθ ( z ) . This can lead to poor outcomes if the latent variables required for good generation take on a more complex distribution , as is common with independent pθ ( x|z ) . One of the simplest methods of gaining greater expressivity in both distributions is to use a hierarchical VAE , which has several stochastic layers of latent variables . These variables are emitted in groups z0 , z1 , ... , zN , which are conditionally dependent upon each other in some way . For images , latent variables are typically output in feature maps of varying resolutions , with z0 corresponding to a small number of latent variables at low resolution at the “ top ” of the network , and zN corresponding to a larger number of latent variables at high resolution at the “ bottom ” . One particularly elegant conditioning structure is the top-down VAE , introduced in Sønderby et al . ( 2016 ) . In this model , both the prior and the approximate posterior generate latent variables in the same order : pθ ( z ) = pθ ( z0 ) pθ ( z1|z0 ) ... pθ ( zN |z < N ) ( 2 ) qφ ( z|x ) = qφ ( z0|x ) qφ ( z1|z0 , x ) ... qφ ( zN |z < N , x ) ( 3 ) A diagram of this process appears in Figure 3 . A typical implementation of this model has φ first perform a deterministic “ bottom-up ” pass on the data to generate features , then processes the groups of latent variables from top to bottom , using feedforward networks to generate features which are shared between the approximate posterior , prior , and reconstruction network pθ ( x|z ) . We adopt this base architecture as it is simple , empirically effective , and has been postulated to resemble biological processes of perception ( Dayan et al. , 1995 ) . 3 WHY DEPTH MATTERS FOR HIERARCHICAL VAES . We find that hierarchical VAEs with sufficient depth can not only learn arbitrary orderings over observed variables , but also learn more effective latent variable distributions , if such distributions exist . We present these results below . Definition ( N -layer VAE ) . A deep hierarchical VAE with N stochastic layers , independent p ( x|z ) , and the top-down factorization of the prior and approximate posterior in Equations 2-3 . Proposition 1 . N-layer VAEs generalize autoregressive models when N is the data dimension Proposition 2 . N -layer VAEs are universal approximators of N -dimensional latent densities Proposition 1 ( proof in Appendix , also visualized in Figure 2 , left ) leads to a possible explanation of why autoregressive models to date have outperformed VAEs : they are deeper , in the sense of statistical dependence . A VAE must be as deep as the data dimension D ( 3072 layers in the case of 32x32 images ) if the images truly require D steps to generate . Luckily , however , Proposition 2 ( proof and further technical requirements in Appendix ) suggests that shorter procedures , if they exist , are also learnable . N = D is an extreme case , where the most effective latent variables z ∈ RD may simply be copies of the observed variables . But if for some K < D there exist latent variables z ∈ RK that the generator can use to more efficiently compress the data , Proposition 2 states a K-layer VAE can learn the posterior and prior distribution over those variables . Such shorter generative paths could emerge in two ways . First , as depicted in Figure 2 ( right ) , if the model discovers that certain variables are conditionally independent given others , the model can generate them in parallel inside a single layer , where qφ ( zN |z < N , x ) = ∏ d qφ ( z ( d ) N |z < N , x ) . We hypothesize these efficient hierarchies should emerge in images , as they contain many spatially independent textures , and study this in Section 5.2 . Second , the model could learn a low-dimensional representation of the data . Dai & Wipf ( 2019 ) recently showed that when a VAE is trained on data distributed on a K-dimensional manifold embedded in RD , a VAE will only activate K dimensions in its latent space , meaning that the VAE will require fewer layers unless the manifold dimension is D , which is unlikely to be the case for images . It is difficult to ascertain the lowest possible value of K for a given dataset , but it may be deeper than most hierarchical VAEs to date . Images have many thousands of observed variables , but early hierarchical VAEs did not exceed 3 layers , until Maaløe et al . ( 2019 ) investigated a Gaussian VAE with 15 layers and found it displayed impressive performance along a variety of measures . Kingma et al . ( 2016 ) and Vahdat & Kautz ( 2020 ) additionally explored networks up to 12 and 30 layers . ( These additionally incorporated additional statistical dependencies in the approximate posterior through the usage of inverse autoregressive flow ( Kingma et al. , 2016 ) , an alternative approach which we contrast with our approach in Section A.4 ) . Nevertheless , given these results we hypothesize that greater depth may improve the performance of VAEs . In the next section , we introduce an architecture capable of scaling to a greater number of stochastic layers . In Section 5.1 we show depth indeed improves performance .

Summary: This paper shows that deep hierarchical VAEs can outperform state-of-the-art autoregressive models on images. The authors first argue that autoregressive models are special cases of hierarchical VAEs and that hierarchical VAEs are universal approximators. They introduce a simple top-down (LVAE) architecture that scales past 70 layers. Furthermore, the model can be trained without using freebits or KL annealing -- although additional tricks are required (gradient skipping and prior warmup). They demonstrate that likelihood performance is correlated with depth and report state-of-the-art performances on multiple image datasets.