Preventing Posterior Collapse with delta-VAEs

Due to the phenomenon of “ posterior collapse , ” current latent variable generative models pose a challenging design choice that either weakens the capacity of the decoder or requires altering the training objective . We develop an alternative that utilizes the most powerful generative models as decoders , optimize the variational lower bound , and ensures that the latent variables preserve and encode useful information . Our proposed δ-VAEs achieve this by constraining the variational family for the posterior to have a minimum distance to the prior . For sequential latent variable models , our approach resembles the classic representation learning approach of slow feature analysis . We demonstrate our method ’ s efficacy at modeling text on LM1B and modeling images : learning representations , improving sample quality , and achieving state of the art log-likelihood on CIFAR-10 and ImageNet 32× 32 . 1 INTRODUCTION . Deep latent variable models trained with amortized variational inference have led to advances in representation learning on high-dimensional datasets ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . These latent variable models typically have simple decoders , where the mapping from the latent variable to the input space is unimodal , for example using a conditional Gaussian decoder . This typically results in representations that are good at capturing the global structure in the input , but fail at capturing more complex local structure ( e.g . texture ( Larsen et al. , 2016 ) ) . In parallel , advances in autoregressive models have led to drastic improvements in density modeling and sample quality without explicit latent variables ( van den Oord et al. , 2016b ) . While these models are good at capturing local statistics , they often fail to produce globally coherent structures ( Ostrovski et al. , 2018 ) . Combining the power of tractable densities from autoregressive models with the representation learning capabilities of latent variable models could result in higher-quality generative models with useful latent representations . While much prior work has attempted to join these two models , a common problem remains . If the autoregressive decoder is expressive enough to model the data density , then the model can learn to ignore the latent variables , resulting in a trivial posterior that collapses to the prior . This phenomenon has been frequently observed in prior work and has been referred to as optimization challenges of VAEs by Bowman et al . ( 2015 ) , the information preference property by Chen et al . ( 2016 ) , and the posterior collapse problems by several others ( e.g . van den Oord et al. , 2017 ; Kim et al. , 2018 ; Dieng et al. , 2018 ) . Ideally , an approach that mitigates posterior collapse would not alter the evidence lower bound ( ELBO ) training objective , and would allow the practitioner to leverage the most recent advances in powerful autoregressive decoders to improve performance . To the best of our knowledge , no prior work has succeeded at this goal . Most common approaches either change the objective ( Higgins et al. , 2017 ; Alemi et al. , 2017 ; Zhao et al. , 2017 ; Chen et al. , 2016 ; Lucas & Verbeek , 2017 ; Goyal et al. , 2017 ) , or weaken the decoder ( Bowman et al. , 2015 ; Gulrajani et al. , 2016 ) . Additionally , these approaches are often challenging to tune and highly sensitive to hyperparameters ( Alemi et al. , 2017 ; Chen et al. , 2016 ) . In this paper , we propose δ-VAEs , a simple framework for selecting variational families that prevent posterior collapse without altering the ELBO training objective or weakening the decoder . By restricting the parameters or family of the posterior , we ensure that there is a minimum KL divergence , δ , between the posterior and the prior . We demonstrate the effectiveness of this approach at learning latent-variable models with powerful decoders on images ( CIFAR-10 , and ImageNet 32× 32 ) , and text ( LM1B ) . We achieve state of the art log-likelihood results with image models by additionally introducing a sequential latent-variable model with an anti-causal encoder structure . Our experiments demonstrate the utility of δ-VAEs at learning useful representations for downstream tasks without sacrificing performance on density modeling . 2 MITIGATING POSTERIOR COLLAPSE WITH δ-VAES Our proposed δ-VAE builds upon the framework of variational autoencoders ( VAEs ) ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) for training latent-variable models with amortized variational inference . Our goal is to train a generative model p ( x , z ) to maximize the marginal likelihood log p ( x ) on a dataset . As the marginal likelihood requires computing an intractable integral over the unobserved latent variable z , VAEs introduce an encoder network q ( z|x ) and optimize a tractable lower bound ( the ELBO ) : log p ( x ) ≥ Ez∼q ( z|x ) [ log p ( x|z ) ] − DKL ( q ( z|x ) ‖p ( z ) ) . The first term is typically referred to as the reconstruction term , and the second term ( KL ) the rate term , as it measures how many nats on average are required to send the latent variables from the encoder ( q ( z|x ) ) to the decoder ( p ( z|x ) ) using a code designed for the prior ( p ( z ) ) ( Hoffman et al. , 2016 ; Alemi et al. , 2017 ) . The problem of posterior collapse is that the rate term , DKL ( q ( z|x ) ‖p ( z ) ) reduces to 0 . In this case , the approximate posterior q ( z|x ) equals the prior p ( z ) , thus the latent variables do not carry any information about the input x . A necessary condition if we want representations to be meaningful is to have the rate term be positive . In this paper we address the posterior collapse problem with structural constraints so that the KL divergence between the posterior and prior is lower bounded by design . This can be achieved by choosing families of distributions for the prior and approximate posterior , pθ ( z ) and qφ ( z|x ) such that minθ , φDKL ( qφ ( z|x ) ‖pθ ( z ) ) ≥ δ . We refer to δ as the committed rate of the model . Note that a trivial choice for p and q to have a non-zero committed rate is to set them to Gaussian distributions with fixed ( but different ) variance term . We study a variant of this case in the experiments , and provide more details of this setup in Appendix D. In the following section we describe our choices for pθ and qφ , but others should also be explored in future work . 2.1 δ-VAE WITH SEQUENTIAL LATENT VARIABLES Data such as speech , natural images and text exhibit strong spatio-temporal continuity . Our aim is to model variations in such data through latent variables , so that we have control over not just the global characteristics of the generated samples ( e.g. , existence of an object ) , but also can influence their finer , often shifting attributes such as texture and pose in the case of natural images , tone , volume and accent in the case of speech , or style and sentiment in the case of natural language . Sequences of latent variables can be an effective modeling tool for expressing the occurrence and evolution of such features throughout the sequence . To construct a δ-VAE in the sequential setting , we combine a mean field posterior with a correlated prior in time . We model the posterior distribution of each timestep as q ( zt|x ) = N ( zt ; µt ( x ) , σt ( x ) ) . For the prior , we use a first-order linear autoregressive process ( AR ( 1 ) ) , where zt = αzt−1 + t with t zero mean Gaussian noise with constant variance σ2 . The conditional probability for the latent variable can be expressed as p ( zt|z < t ) = N ( zt ; αzt−1 , σ ) . This process is wide-sense stationary ( that is , having constant sufficient statistics through its time evolution ) if |α| < 1 . If so , then zt has zero mean and variance of σ 2 1−α2 . It is thus convenient to choose σ = √ 1− α2 so that the variance is constant over time . The mismatch in the correlation structure of the prior and the posterior results in the following positive lower bound on the KL-divergence between the two distributions ( see Appendix C for derivation ) : DKL ( q ( z|x ) ‖p ( z ) ) ≥ 1 2 d∑ k=1 ( n− 2 ) ln ( 1 + α2k ) − ln ( 1− α2k ) ( 1 ) where n is the length of the sequence and d is the dimension of the latent variable at each timestep . The committed rate between the prior and the posterior is easily controlled by equating the right hand side of the inequality in equation 1 to a given rate δ and solving for α . In Fig . 1 , we show the scaling of the minimum rate as a function of α and the behavior of δ-VAE in 2d . 2.1.1 RELATION TO PROBABILISTIC SLOWNESS PRIOR . The AR ( 1 ) prior over the latent variables specifies the degree of temporal correlation in the latent space . As the correlation α approaches one , the prior trajectories get smoother . On the other hand in the limit of α approaching 0 , the prior becomes the same as the independent standard Gaussian prior where there are no correlations between timesteps . This pairing of independent posterior with a correlated prior is related to the probabilistic counterpart to Slow Feature Analysis ( Wiskott & Sejnowski , 2002 ) in Turner & Sahani ( 2007 ) . SFA has been shown to be an effective method for learning invariant spatio-temporal features ( Wiskott & Sejnowski , 2002 ) . In our models , we infer latent variables with multiple dimensions per timestep , each with a different slowness filter imposed by a different value of α , corresponding to features with different speed of variation . 2.2 ANTI-CAUSAL ENCODER NETWORK . Having a high capacity autoregressive network as the decoder implies that it can accurately estimate p ( xt|x < t ) . Given this premise , what kind of complementary information can latent variables provide ? Encoding information about the past seems wasteful as the autoregressive decoder has full access to past observations already . On the other hand , if we impose conditional independence between observations and latent variables at other timesteps given the current one ( i.e. , p ( xt|z ) = p ( xt|zt ) ) , there will then be at best ( by the data processing inequality ( Cover & Thomas , 2006 ) ) a break-even situation between the KL cost of encoding information in zt and the resulting improvement in the reconstruction loss . There is therefore no advantage for the model to utilize the latent variable even if it would transmit to the decoder the unobserved xt . The situation is different when zt can inform the decoder at multiple timesteps , encoding information about xt and x > t . In this setting , the decoder pays the KL cost for the mutual information once , but is able to leverage the transmitted information multiple times to reduce its entropy about future predictions . To encourage the generative model to leverage the latents for future timesteps , we introduce an anti-causal structure for the encoder where the parameters of the variational posterior for a timestep can not depend on past observations ( Fig . 2 ) . Alternatively , one can consider a non-causal structure that allows latents be inferred from all observations . In this non-causal setup there is no temporal order in either the encoder or the decoder , thus the model resembles a standard non-temporal latentvariable model . While the anti-causal structure is a subgraph of the non-causal structure , we find that the anti-causal structure often performs better , and we compare both approaches in different settings in Appendix F.1 .

Summary: The paper proposes a method to prevent posterior collapse, which refers to the phenomenon that VAEs with powerful autoregressive decoders tend to ignore the latent code, i.e., the decoder models the data distribution independently of the code. Specifically, the encoder, decoder, and prior distribution families are chosen such that the KL-term in the ELBO is bounded away from 0, meaning that the encoder output cannot perfectly match the prior. Assuming temporal data, the authors employ a 1-step autoregressive (across) prior with an encoder whose codes are independent conditionally on the input. Furthermore, they propose to use a causal decoder together with an anti-causal or non-causal encoder, which translates into a PixelSNAIL/PixelCNN style decoder and an anti-causal version thereof as encoder in the case of image data. The proposed approach is evaluated on CIFAR10, Imagenet 32x32, and the LM1B data set (text).