Variational Autoencoders with Normalizing Flow Decoders

Recently proposed normalizing flow models such as Glow ( Kingma & Dhariwal , 2018 ) have been shown to be able to generate high quality , high dimensional images with relatively fast sampling speed . Due to their inherently restrictive architecture , however , it is necessary that they are excessively deep in order to train effectively . In this paper we propose to combine Glow model with an underlying variational autoencoder in order to counteract this issue . We demonstrate that such our proposed model is competitive with Glow in terms of image quality while requiring far less time for training . Additionally , our model achieves state-of-the-art FID score on CIFAR-10 for a likelihood-based model . 1 INTRODUCTION . The field of deep generative models is experiencing rapid progress , with new state-of-the-art results being achieved regularly . Currently , state-of-the-art in image generation is dominated by adversarial training methods such as Karras et al . ( 2019 ) and Brock et al . ( 2019 ) . However , Generative Adversarial Networks ( GANs ) notoriously suffer from problems with training stability and mode collapse due to the dynamics of an adversarial objective . Many ideas have been proposed for dealing with these issues , such as Arjovsky et al . ( 2017 ) and Miyato et al . ( 2018 ) , however issues inherently remain . Due to this , there is still signficant attention being given to pure likelihood models despite their poorer qualitative performance , as maximum likelihood grants stable training and is able to achieve good mode coverage on account of its heavy penalization of models that do not allocate probability mass to space containing training data . There are numerous classes of likelihood-based deep generative models in the contemporary literature . Autoregressive models such as PixelRNN ( van den Oord et al. , 2016 ) have been shown to achieve impressive test likelihood scores on various datasets , however such models typically perform poorly when evaluated using quantitative measures of image quality , and in addition take a long time to generate samples . Variational autoencoders ( VAE ) ( Kingma & Welling , 2014 ) , a form of latent variable model and one of the earlier proposed deep generative models , are well known to produce samples of poor quality when implemented in their most basic form and as such require combining with more powerful generative models in order to be competitive . More recently , it was shown in Kingma & Dhariwal ( 2018 ) that a class of likelihood models known as normalizing flows can be used to generate high quality images with much faster sampling times than autoregressive models . Heavy restrictions must be imposed on the model architecture however in order to be able to efficiently calculate the determinant of the model Jacobian , forcing implementations of such models to be extremely deep and wide to make up for deficiencies of a restricted architecture . As a result , training requires an exorbitant amount of time ; as an example 40 GPUs and ∼1 week of training was required in order to produce the 256x256 CelebA-HQ samples in Kingma & Dhariwal ( 2018 ) . Finding ways to cut down on these resource requirements is therefore pertinent in order to make such models accessible to the average practitioner . Towards achieving this goal , in this work we propose combining the standard variational autoencoder model with an overlying Glow layer and demonstrate state-of-the-art FID ( Heusel et al. , 2017 ) results while simultaneously achieving much lower training time . 2 BACKGROUND AND RELATED MODELS . 2.1 VARIATIONAL AUTOENCODERS . Variational Autoencoders ( VAEs ) , first proposed in Kingma & Welling ( 2014 ) are a popular class of deep latent variable model inspired by traditional variational inference . Given a dataset X = { x ( 1 ) , x ( 2 ) , ... , x ( n ) } , the goal is to maximize the log-likelihood lower bound Eq ( z|x ( i ) ) [ log p ( x ( i ) |z ) ] −DKL [ q ( z|x ( i ) ) ||p ( z ) ] ( 1 ) where q ( z|x ) is an encoder distribution and p ( x|z ) is a decoder distribution , both parameterized by deep neural networks . The prior distribution p ( z ) may be fixed or learned . Typically the encoder , decoder and prior distributions are all assumed to be some simple distribution , usually Gaussian . In this case the model enjoys the properties of stable training and fast convergence , as well as fast sampling . VAEs are well known to produce sample images of poor quality , however , owing to a number of previously identified factors . It has been pointed out that the typical Gaussian assumption of the posterior q ( z|x ) leads to poor posterior approximations , as the true posterior is often far from Gaussian in shape . As the optimizer is forced to fit mismatching distributions , this results in the marginal encoder distribution q ( z ) diverging from the prior p ( z ) . In Jimenez Rezende & Mohamed ( 2015 ) and many subsequent works such as Kingma et al . ( 2016 ) , it was proposed that the posterior be made more flexible via the use of normalizing flows , which are discussed in detail in the next section . In this way , the approximate posterior can achieve a better fit to the true posterior . In Dai & Wipf ( 2019 ) the authors pointed out that the failure to optimize the decoder variance in typical implementations of VAEs is a major cause of poor quality samples , and that by optimizing decoder variance it is possible to achieve an optimal VAE cost in the case where the encoder and decoder distributions are Gaussian ( specifically , a diagonal Gaussian encoder and isotropic Gaussian decoder ) . Achieving the optimal VAE cost does not necessarily recover the ground truth distribution , though , and so they additionally proposed learning the distribution of the latent space using a secondstage VAE , and proved that by doing so it is theoretically possible to recover the ground truth distribution . Despite there being a theoretically obtainable optimal cost in the case where the decoder variance is learned , in practice using a decoder that assumes independence between pixels is not ideal , as it is well understood that this leads to blurry reconstructions due to the optimizer favouring solutions that are a weighted mean of data points in pixel space . Additionally , because Gaussian decoders are a poor fit for natural images , they may introduce substantial salt-and-pepper noise into the generated samples , and so it is commonplace to restrict decoder sampling to the mean . This artificially restricts the generative distribution to a manifold of dimensionality equal to that of the latent space , and potentially limits sample diversity . It is precisely these deficits of simple decoder distributions used in standard VAE models that we attempt to address in our proposal . Complex , non-Gaussian decoder distributions have been proposed in several previous works . In Gulrajani et al . ( 2017 ) , the authors propose using autoregressive PixelCNN layers ( Van den Oord et al. , 2016 ) in the decoder , and are able to achieve impressive test likelihood scores . An adversarially trained autoencoder was proposed in Larsen et al . ( 2016 ) , with the decoder objective being defined by a standard GAN loss over the generated image plus a representation loss using L2 distance between discriminator hidden features . However , such a model can not be used to evaluate likelihood , and being an adversarial method it suffers from problems with training stability . Shmelkov et al . ( 2019 ) propose a very similar model to our own , with a Real NVP ( Dinh et al. , 2016 ) normalizing flow model stacked on top of a VAE with Gaussian decoder . They combine likelihood maximization with adversarial training to achieve competitive FID scores on various datasets , however their model achieves poor results when using a purely likelihood based approach . 2.2 NORMALIZING FLOWS . Normalizing flows take advantage of the change of variables formula for probability distributions , given by p ( x ) = p ( z ) |det ( ∂f ( x ) /∂xT ) | ( 2 ) where f defines a bijection between the spaces of x and z . They have found use with application to VAEs in order to achieve flexible priors/posteriors ( Jimenez Rezende & Mohamed , 2015 ; Chen et al. , 2017 ) , as well as being used as standalone generative models directly in data space ( Dinh et al. , 2014 ) in which case exact likelihood evaluation is possible . One of the most common implementations is that of autoregressive flows , whereby the bijection f is made to be auto-regressive resulting in a triangular Jacobian . This is important for efficiency , as it allows the determinant in Eq . 2 to be calculated in O ( n ) rather than O ( n3 ) . Composition of multiple autoregressive functions allows for the construction of high capacity models . 3 METHOD AND MODEL DEFINITION . At a high level , our idea is to combine a simple VAE model with a normalizing flow model in the pixel space . As such , we hope to obtain the benefits of both : the smaller model size , faster training , and compact latent representation of VAEs , as well as the high quality samples of normalizing flows . We use a simple VAE to learn a base distribution , and train a normalizing flow using this base distribution , such that we end up with a conditional normalizing flow . As the VAE should be able to learn a distribution that is already close to the ground truth , our normalizing flow , when conditioned on this base distribution , should not need to do as much work as a full marginal normalizing flow model such as Glow ( Kingma & Dhariwal , 2018 ) . Our normalizing flow therefore need not be as deep as such models , saving training time and model size . Our model is trained using pure likelihood maximization , and is capable of evaluating a lower bound on the likelihood of arbitrary data samples . For our approach we use an underlying Gaussian decoder to learn the base distribution of a normalizing flow . We denote the random variable produced by this underlying decoder as x̄ ∼ px̄ ( x̄|z ) = N ( x̄|gµ ( z ) , gΣ ( z ) ) ( 3 ) where gµ decodes the mean of px̄ ( x̄ ) and gΣ decodes the diagonal covariance . The normalizing flow layer then defines a bijection f : X̂ → X̄ and random variable x̂ ∼ px̂ ( x̂|z ) = px̄ ( f ( x̂ ) |z ) |det ( ∂f ( x̂ ) /∂x̂T ) | ( 4 ) where x̂ ∈ X̂ and x̄ ∈ X̄ . 3.1 UNDERLYING GAUSSIAN DECODER . In Dai & Wipf ( 2019 ) the authors propose to add a learnable parameter γ to the model , with the decoder likelihood given byN ( x|gµ ( z ) , γI ) , and prove that doing so enables the model to always be capable of achieving a better VAE cost by lowering the value of γ . This has an additional benefit in the case of our model , as it facilitates training of the overlying Glow layer . By allowing the VAE decoder to shrink its variance as it gains confidence in its estimates , the Glow layer has to do less work to remove any erroneous noise in the decoder distribution . We extend this idea by allowing the decoder distribution to be Gaussian with diagonal covariance , rather than simply an isotropic Gaussian . By doing so , we allow the decoder to express confidence in its estimates at the pixel level . In the case of MNIST , for example , the decoder is likely to be highly confident in its estimate of the uniformly black background pixels , while the pixels on the border of the digit are likely to be the source of lower confidence . When sampling from the underlying VAE ( while keeping decoder samples restricted to the mean ) this does not result in any significant improvement over the simpler isotropic Gaussian decoder . For our full model , however , we found in our experiments that this small change resulted in significant improvements to FID score .

Summary: This paper proposed a variant of the variational autoencoder, specifically by using a simplified normalizing flow model as the decoder. When compared with Glow, the proposed method is simpler and more efficient to train. The authors applied their algorithm on a few datasets, and showed better or competitive performance, both qualitatively and quantitatively.