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: The paper claims that high quality of generated samples and SOTA bpds are achievable by VAEs if the model is deep enough (deep in terms of the number of stochastic layers). The authors explain the architecture that resemblances the U-net architecture, and explain its building blocks. Interestingly, they are able to learn VAEs with up to 78 stochastic layers, and achieve SOTA bpds on CIFAR-10, ImageNet-32, ImageNet-64, FFHQ-256 (5-bit), and setting a great result on FFHQ-1024 (8bit).