Lagging Inference Networks and Posterior Collapse in Variational Autoencoders

1 INTRODUCTION . Variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ) represent a popular combination of a deep latent variable model ( shown in Figure 1 ( a ) ) and an accompanying variational learning technique . The generative model in VAE defines a marginal distribution on observations , x ∈ X , as : pθ ( x ) = ∫ pθ ( x|z ) p ( z ) dz . ( 1 ) The model ’ s generator defines pθ ( x|z ) and is typically parameterized as a complex neural network . Standard training involves optimizing an evidence lower bound ( ELBO ) on the intractable marginal data likelihood ( Eq.1 ) , where an auxiliary variational distribution qφ ( z|x ) is introduced to approximate the model posterior distribution pθ ( z|x ) . VAEs make this learning procedure highly scalable to large datasets by sharing parameters in the inference network to amortize inferential cost . This amortized approach contrasts with traditional variational techniques that have separate local variational parameters for every data point ( Blei et al. , 2003 ; Hoffman et al. , 2013 ) . While successful on some datasets , prior work has found that VAE training often suffers from “ posterior collapse ” , in which the model ignores the latent variable z ( Bowman et al. , 2016 ; Kingma et al. , 2016 ; Chen et al. , 2017 ) . This phenomenon is more common when the generator pθ ( x|z ) is parametrized with a strong autoregressive neural network , for example , an LSTM ( Hochreiter & Schmidhuber , 1997 ) on text or a PixelCNN ( van den Oord et al. , 2016 ) on images . Posterior collapse is especially evident when modeling discrete data , which hinders the usage of VAEs in 1Code and data are available at https : //github.com/jxhe/vae-lagging-encoder . important applications like natural language processing . Existing work analyzes this problem from a static optimization perspective , noting that the collapsed solution is often a reasonably good local optimum in terms of ELBO ( Chen et al. , 2017 ; Zhao et al. , 2017 ; Alemi et al. , 2018 ) . Thus , many proposed solutions to posterior collapse focus on weakening the generator by replacing it with a non-recurrent alternative ( Yang et al. , 2017 ; Semeniuta et al. , 2017 ) or modifying the training objective ( Zhao et al. , 2017 ; Tolstikhin et al. , 2018 ) . In this paper , we analyze the problem from the perspective of training dynamics and propose a novel training procedure for VAEs that addresses posterior collapse . In contrast with other solutions , our proposed procedure optimizes the standard ELBO objective and does not require modification to the VAE model or its parameterization . Recently , Kim et al . ( 2018 ) proposed a new approach to training VAEs by composing the standard inference network with additional mean-field updates . The resulting semi-amortized approach empirically avoided collapse and obtained better ELBO . However , because of the costly instancespecific local inference steps , the new method is more than 10x slower than basic VAE training in practice . It is also unclear why the basic VAE method fails to find better local optima that make use of latents . We consider two questions in this paper : ( 1 ) Why does basic VAE training often fall into undesirable collapsed local optima ? ( 2 ) Is there a simpler way to change the training trajectory to find a non-trivial local optimum ? To this end , we first study the posterior collapse problem from the perspective of training dynamics . We find , empirically , that the posterior approximation often lags far behind the true model posterior in the initial stages of training ( Section 3 ) . We then demonstrate how such lagging behavior can drive the generative model towards a collapsed local optimum , and propose a novel training procedure for VAEs that aggressively optimizes the inference network with more updates to mitigate lag ( Section 4 ) . Without introducing new modeling components over basic VAEs or additional complexity , our approach is surprisingly simple yet effective in circumventing posterior collapse . As a density estimator , it outperforms neural autoregressive baselines on both text ( Yahoo and Yelp ) and image ( OMNIGLOT ) benchmarks , leading to comparable performance with more complicated previous state-of-the-art methods at a fraction of the training cost ( Section 6 ) . 2 BACKGROUND . 2.1 VARIATIONAL AUTOENCODERS . VAEs learn deep generative models defined by a prior p ( z ) and a conditional distribution pθ ( x|z ) as shown in Figure 1 ( a ) . In most cases the marginal data likelihood is intractable , so VAEs instead optimize a tractable variational lower bound ( ELBO ) of log pθ ( x ) , L ( x ; θ , φ ) = Ez∼qφ ( z|x ) [ log pθ ( x|z ) ] ︸ ︷︷ ︸ Reconstruction Loss −DKL ( qφ ( z|x ) ‖p ( z ) ) ︸ ︷︷ ︸ KL Regularizer , ( 2 ) where qφ ( z|x ) is a variational distribution parameterized by an inference network with parameters φ , and pθ ( x|z ) denotes the generator network with parameters θ. qφ ( z|x ) is optimized to approximate the model posterior pθ ( z|x ) . This lower bound is composed of a reconstruction loss term that encourages the inference network to encode information necessary to generate the data and a KL regularizer to push qφ ( z|x ) towards the prior p ( z ) . Below , we consider p ( z ) : = N ( 0 , I ) unless otherwise specified . A key advantage of using inference networks ( also called amortized inference ) to train deep generative models over traditional locally stochastic variational inference ( Hoffman et al. , 2013 ) is that they share parameters over all data samples , amortizing computational cost and allowing for efficient training . The term VAE is often used both to denote the class of generative models and the amortized inference procedure used in training . In this paper , it is important to distinguish the two and throughout we will refer to the generative model as the VAE model , and the training procedure as VAE training . 2.2 POSTERIOR COLLAPSE . Despite VAE ’ s appeal as a tool to learn unsupervised representations through the use of latent variables , as mentioned in the introduction , VAE models are often found to ignore latent variables when using flexible generators like LSTMs ( Bowman et al. , 2016 ) . This problem of “ posterior collapse ” occurs when the training procedure falls into the trivial local optimum of the ELBO objective in which both the variational posterior and true model posterior collapse to the prior . This is undesirable because an important goal of VAEs is to learn meaningful latent features for inputs . Mathematically , posterior collapse represents a local optimum of VAEs where qφ ( z|x ) = pθ ( z|x ) = p ( z ) for all x . To facilitate our analysis about the causes leading up to collapse , we further define two partial collapse states : model collapse , when pθ ( z|x ) = p ( z ) , and inference collapse , when qφ ( z|x ) = p ( z ) for all x . Note that in this paper we use these two terms to denote the posterior states in the middle of training instead of local optima at the end . These two partial collapse states may not necessarily happen at the same time , which we will discuss later . 2.3 VISUALIZATION OF POSTERIOR DISTRIBUTION . Posterior collapse is closely related to the true model posterior pθ ( z|x ) and the approximate posterior qφ ( z|x ) as it is defined . Thus , in order to observe how posterior collapse happens , we track the state of pθ ( z|x ) and qφ ( z|x ) over the course of training , and analyze the training trajectory in terms of the posterior mean space U = { µ : µ = ( µTx , θ , µTx , φ ) } , where µx , θ and µx , φ are the means of pθ ( z|x ) and qφ ( z|x ) , respectively.2 We can then roughly consider µx , θ = 0 as model collapse and µx , φ = 0 as inference collapse as we defined before . Each x will be projected to a point in this space under the current model and inference network parameters . If z is a scalar we can efficiently compute µx , θ and visualize the posterior mean space as shown in Figure 1 ( b ) . The diagonal line µx , θ = µx , φ represents parameter settings where qφ ( z|x ) is equal to pθ ( z|x ) in terms of mean , indicating a well-trained inference network . The collapsed local optimum is located at the origin,3 while the data points at a more desirable local optima may be distributed along the diagonal . In this paper we will utilize this posterior mean space multiple times to analyze the posterior dynamics . 3 A LAGGING INFERENCE NETWORK PREVENTS USING LATENT CODES . In this section we analyze posterior collapse from a perspective of training dynamics . We will answer the question of why the basic VAE training with strong decoders tends to hit a collapsed local optimum and provide intuition for the simple solution we propose in Section 4 . 2µx , θ can be approximated through discretization of the model posterior , which we show in Appendix A . 3Note that the converse is not true : the setting where all points are located at the origin may not be a local optimum . For example when a model is initialized at the origin as we show in Section 3.2 . Published as a conference paper at ICLR 2019Under review as a conferenc paper at ICLR 2019 Figure 2 : The projections of 500 data samples on the posterior mean space . From left to right we show the change of these points at four different time stamps in the course of training . The top row is from the basic VAE training , the bottom row is from our approach . [ Can you make the font slightly bigger ? –GN ] [ Can you also add labels indicating the four training stages ( “ Stage 1 ” , “ Stage 2 ” , ... , or even better “ X Epochs ” , “ X Epochs ” , ... ) . –GN ] [ Add labels to the top of “ Vanilla VAE ” and “ Our Approach ” ( or whatever alternative name you come up with . ) –GN ] p✓ ( z|x ) is log p✓ ( x ) . However , the dependence ( if any ) may be lost if q ( z|x ) is still close to p ( z ) because p✓ ( z|x ) would be pressured to agree with q ( z|x ) [ This is important , but I think not very clear . Could you be more explicit ? –GN ] . To examine the dynamic relationship between approximate posterior q ( z|x ) , model posterior p✓ ( z|x ) , and prior p ( z ) during basic VAE training , we will next visualize the posterior mean-value space by training a basic VAE with a one-dimensional [ “ scalar ” ? –GN ] latent variable on a relatively simple synthetic dataset . 3.2 OBSERVATIONS ON SYNTHETIC DATA As a synthetic dataset we use discrete sequence data , as posterior collapse has been found the most severe in text modeling tasks . To mimic the diversity of the real data distribution , we sample discrete data from a deep generative model , with a [ [ [ four-mode mixture of Gaussian prior — four-modes of a Gaussian mixture as a prior ? ] ] ] and an LSTM generator . More details about this synthetic dataset and experiment details can be found in Appendix B.1 . We train a basic VAE with a one-dimensional latent variable , LSTM encoder , and LSTM decoder2 on our synthetic dataset . We sample 500 data points and show them on the posterior mean space plots at four different training stages from beginning to convergence in Figure 2 . The mean value of the approximate posterior distribution µx , is from the output of the inference network , and µx , ✓ is approximated by : E z⇠p✓ ( z|x ) [ z ] = X zi2C [ zip ( zi|x ) ] , ( 5 ) where C is a grid set with small stride and sufficiently large coverage . We assume the density value outside this grid set is zero . p✓ ( z|x ) is also approximated on this grid . As illustrated in Figure 2 , all points are located at origin [ [ [ after—upon ] ] ] initialization , which means z and x are both almost independent and zero mean at the beginning of training . In the second stage , the points starts to spread along the µx , ✓ axis . This phenomenon implies that for some data points p✓ ( z|x ) moves far away from the prior p ( z ) , and confirms that log p✓ ( x ) is able to help escape model collapse . However , all of these points are still close to the µx , axis , which suggests that q ( z|x ) fails to catch up to p✓ ( z|x ) and these points are still in an inference collapsed state . 2This LSTM decoder has less capacity than the one used for creating the dataset since in real world model capacity is usually insufficient to exactly model the empirical distribution . 4 iter = 0 iter = 200 iter = 2000 convergence B as ic A gg re ss iv e µx , ✓ < latexit sha1_base64= '' XWKglMNT3flz+HvE5Pzzkxay/dQ= '' > AAACJHicZVBNSwMxEM3Wr1q/qh5FCBbBQylbEdRb0YvHCtYWuqVk02wbmmyWZFYsy578MeJVf4cn8eDF/+A/MLvtwdaBMC9v3gwzz48EN+C6X05haXllda24XtrY3NreKe/u3RsVa8paVAmlOz4xTPCQtYCDYJ1IMyJ9wdr++Dqrtx+YNlyFdzCJWE+SYcgDTglYql8+9GTcTzxJYOQHyWNaxZ4vEw9GDEia9ssVt+bmgf+D+gxU0Cya/fKPN1A0liwEKogx3bobQS8hGjgVLC15sWERoWMyZF0LQyKZ6SX5GSk+tswAB0rbFwLO2b8dCZEmW9Qqs2TmahkDSglTtSoYySxlY/K/mUi/6stqJtImMAuLQHDRS3gYxcBCOt0jiAUGhTPL8IBrRkFMLCBUc3sKpiOiCQVrbMl6VF905D9ondYua+7tWaVxNTOriA7QETpBdXSOGugGNVELUfSEXtArenOenXfnw/mcSgvOrGcfzYXz/Qs8kaaP < /latexit > < latexit sha1_base64= '' XWKglMNT3flz+HvE5Pzzkxay/dQ= '' > AAACJHicZVBNSwMxEM3Wr1q/qh5FCBbBQylbEdRb0YvHCtYWuqVk02wbmmyWZFYsy578MeJVf4cn8eDF/+A/MLvtwdaBMC9v3gwzz48EN+C6X05haXllda24XtrY3NreKe/u3RsVa8paVAmlOz4xTPCQtYCDYJ1IMyJ9wdr++Dqrtx+YNlyFdzCJWE+SYcgDTglYql8+9GTcTzxJYOQHyWNaxZ4vEw9GDEia9ssVt+bmgf+D+gxU0Cya/fKPN1A0liwEKogx3bobQS8hGjgVLC15sWERoWMyZF0LQyKZ6SX5GSk+tswAB0rbFwLO2b8dCZEmW9Qqs2TmahkDSglTtSoYySxlY/K/mUi/6stqJtImMAuLQHDRS3gYxcBCOt0jiAUGhTPL8IBrRkFMLCBUc3sKpiOiCQVrbMl6VF905D9ondYua+7tWaVxNTOriA7QETpBdXSOGugGNVELUfSEXtArenOenXfnw/mcSgvOrGcfzYXz/Qs8kaaP < /latexit > < latexit sha1_base64= '' XWKglMNT3flz+HvE5Pzzkxay/dQ= '' > AAACJHicZVBNSwMxEM3Wr1q/qh5FCBbBQylbEdRb0YvHCtYWuqVk02wbmmyWZFYsy578MeJVf4cn8eDF/+A/MLvtwdaBMC9v3gwzz48EN+C6X05haXllda24XtrY3NreKe/u3RsVa8paVAmlOz4xTPCQtYCDYJ1IMyJ9wdr++Dqrtx+YNlyFdzCJWE+SYcgDTglYql8+9GTcTzxJYOQHyWNaxZ4vEw9GDEia9ssVt+bmgf+D+gxU0Cya/fKPN1A0liwEKogx3bobQS8hGjgVLC15sWERoWMyZF0LQyKZ6SX5GSk+tswAB0rbFwLO2b8dCZEmW9Qqs2TmahkDSglTtSoYySxlY/K/mUi/6stqJtImMAuLQHDRS3gYxcBCOt0jiAUGhTPL8IBrRkFMLCBUc3sKpiOiCQVrbMl6VF905D9ondYua+7tWaVxNTOriA7QETpBdXSOGugGNVELUfSEXtArenOenXfnw/mcSgvOrGcfzYXz/Qs8kaaP < /latexit > µx , < latexit sha1_base64= '' JrR8bpJT9WIDXN5JEdhmmw5BeVU= '' > AAACInicZVBNSwMxEM36WetX1WMvwSJ4KGUrgnorevFYwbWFbinZNNuGJpslmRXLsgd/jHjV3+FJPAn+Cf+B2bYH2w6EeXnzZph5QSy4Adf9dlZW19Y3Ngtbxe2d3b390sHhg1GJpsyjSijdDohhgkfMAw6CtWPNiAwEawWjm7zeemTacBXdwzhmXUkGEQ85JWCpXqnsy6SX+pLAMAjTp6yK/UCmfjzkWdYrVdyaOwm8DOozUEGzaPZKv35f0USyCKggxnTqbgzdlGjgVLCs6CeGxYSOyIB1LIyIZKabTo7I8Ill+jhU2r4I8IT935ESafI1rTJPZq6WM6CUMFWrgqHMUz5m8jdjGVQDWc1F2oRmYREIL7spj+IEWESne4SJwKBwbhjuc80oiLEFhGpuT8F0SDShYG0tWo/qi44sA++sdlVz784rjeuZWQVURsfoFNXRBWqgW9REHqLoGb2iN/TuvDgfzqfzNZWuOLOeIzQXzs8fgA2lpg== < /latexit > < latexit sha1_base64= '' JrR8bpJT9WIDXN5JEdhmmw5BeVU= '' > AAACInicZVBNSwMxEM36WetX1WMvwSJ4KGUrgnorevFYwbWFbinZNNuGJpslmRXLsgd/jHjV3+FJPAn+Cf+B2bYH2w6EeXnzZph5QSy4Adf9dlZW19Y3Ngtbxe2d3b390sHhg1GJpsyjSijdDohhgkfMAw6CtWPNiAwEawWjm7zeemTacBXdwzhmXUkGEQ85JWCpXqnsy6SX+pLAMAjTp6yK/UCmfjzkWdYrVdyaOwm8DOozUEGzaPZKv35f0USyCKggxnTqbgzdlGjgVLCs6CeGxYSOyIB1LIyIZKabTo7I8Ill+jhU2r4I8IT935ESafI1rTJPZq6WM6CUMFWrgqHMUz5m8jdjGVQDWc1F2oRmYREIL7spj+IEWESne4SJwKBwbhjuc80oiLEFhGpuT8F0SDShYG0tWo/qi44sA++sdlVz784rjeuZWQVURsfoFNXRBWqgW9REHqLoGb2iN/TuvDgfzqfzNZWuOLOeIzQXzs8fgA2lpg== < /latexit > < latexit sha1_base64= '' JrR8bpJT9WIDXN5JEdhmmw5BeVU= '' > AAACInicZVBNSwMxEM36WetX1WMvwSJ4KGUrgnorevFYwbWFbinZNNuGJpslmRXLsgd/jHjV3+FJPAn+Cf+B2bYH2w6EeXnzZph5QSy4Adf9dlZW19Y3Ngtbxe2d3b390sHhg1GJpsyjSijdDohhgkfMAw6CtWPNiAwEawWjm7zeemTacBXdwzhmXUkGEQ85JWCpXqnsy6SX+pLAMAjTp6yK/UCmfjzkWdYrVdyaOwm8DOozUEGzaPZKv35f0USyCKggxnTqbgzdlGjgVLCs6CeGxYSOyIB1LIyIZKabTo7I8Ill+jhU2r4I8IT935ESafI1rTJPZq6WM6CUMFWrgqHMUz5m8jdjGVQDWc1F2oRmYREIL7spj+IEWESne4SJwKBwbhjuc80oiLEFhGpuT8F0SDShYG0tWo/qi44sA++sdlVz784rjeuZWQVURsfoFNXRBWqgW9REHqLoGb2iN/TuvDgfzqfzNZWuOLOeIzQXzs8fgA2lpg== < /latexit > Figure 2 : The projections of 500 data samples from a synthetic dataset on the posterior mean space over the course of training . “ iter ” denotes the number of updates of generators . The top row is from the basic VAE training , the bottom row is from our aggressive inference network training . The results show that while the approximate posterior is lagging far behind the true model posterior in basic VAE training , our aggressive training approach successfully moves the points onto the diagonal line and away from inference collapse .

Summary: The paper tackles one of the most important problems of learning VAEs, namely, the posterior collapse. Typically, this problem is attacked by either proposing a new model or modifying the objective. Interestingly, the authors considered a third option, i.e., changing the training procedure only, leaving the model and the objective untouched. Moreover, they show that in fact the modified objective (beta-VAE) could drastically harm training a VAE.