A Batch Normalized Inference Network Keeps the KL Vanishing Away (2020)

Abstract: Variational Autoencoder (VAE) is widely used as a generative model to approximate a model’s posterior on latent variables by combining the amortized variational inference and deep neural networks. However, when paired with strong autoregressive decoders, VAE often converges to a degenerated local optimum known as “posterior collapse”. Previous approaches consider the Kullback–Leibler divergence (KL) individual for each datapoint. We propose to let the KL follow a distribution across the whole dataset, and analyze that it is sufficient to prevent posterior collapse by keeping the expectation of the KL’s distribution positive. Then we propose Batch Normalized-VAE (BN-VAE), a simple but effective approach to set a lower bound of the expectation by regularizing the distribution of the approximate posterior’s parameters. Without introducing any new model component or modifying the objective, our approach can avoid the posterior collapse effectively and efficiently. We further show that the proposed BN-VAE can be extended to conditional VAE (CVAE). Empirically, our approach surpasses strong autoregressive baselines on language modeling, text classification and dialogue generation, and rivals more complex approaches while keeping almost the same training time as VAE.

## introduction variational autoencoder (vae) @xcite is one of the most popular generative framework to model complex distributions. different from the autoencoder (ae), vae provides a distribution-based latent representation for the data, which encodes the input x into a probability distribution z and reconstructs the original input using samples from z. when *this work was done when qile zhu was an intern at tencent ai lab. wei bi is the corresponding author. an autoregressive decoder (e.g., a recurrent neural network) is a common choice to model the text data. however, when paired with strong autoregressive decoders such as lstms @xcite and trained under conventional training strategy, vae suffers from a wellknown problem named the posterior collapse or the kl vanishing problem. the decoder in vae learns to reconstruct the data independent of the latent variable z, and the kl vanishes to 0. many convincing solutions have been proposed to prevent posterior collapse. among them, fixing the kl as a positive constant is an important direction @xcite @xcite @xcite @xcite . some change the gaussian prior with other distributions, e.g., a uniform prior @xcite or a von mises-fisher (vmf) distribution @xcite @xcite . however, these approaches force the same constant kl and lose the flexibility to allow various kls for different data points @xcite . without changing the gaussian prior, free-bits @xcite adds a threshold (free-bits) of the kl term in the elbo object and stops the optimization of the kl part when its value is smaller than the threshold. @xcite point out that the objective of free-bits is non-smooth and suffers from the optimization challenges. δ-vae @xcite sets the parameters in a specific range to achieve a positive kl value for every latent dimension, which may limit the model performance. other work analyzes this problem form a view of optimization @xcite @xcite . recently, @xcite observe that the inference network is lagging far behind the decoder during training. they propose to add additional training loops for the inference network only. @xcite further propose to initialize the inference network with an encoder pretrained from an ae objective, then trains the vae with the free-bits. however, these two methods are much slower than the original vae. the limitation of the constant kl and the high cost of additional training motivate us to seek an approach that allows flexible modeling for different data points while keeping as fast as the original vae. in this paper, instead of considering the kl individually for each data point, we let it follow a distribution across the whole dataset. we demonstrate that keeping a positive expectation of the kl's distribution is sufficient to prevent posterior collapse in practice. by regularizing the distribution of the approximate posterior's parameters, a positive lower bound of this expectation could be ensured. then we propose batch normalized-vae (bn-vae), a simple yet effective approach to achieving this goal, and discuss the connections between bn-vae and previous enhanced vae variants. we further extend bn-vae to the conditional vae (cvae). last, experimental results demonstrate the effectiveness of our approach on real applications, including language modeling, text classification and dialogue generation. empirically, our approach surpasses strong autoregressive baselines and is competitive with more sophisticated approaches while keeping extremely higher efficiency. code and data are available at https://github.com/valdersoul/bn-vae . ## background and related work in this section, we first introduce the basic background of vae, then we discuss the lagging problem @xcite . at last, we present more related work. ## batch-normalized vae in this section, we first derive the expectation of the kl's distribution and show that it is enough to avoid posterior collapse by keeping the expectation of the kl's distribution positive. then we propose our regularization method on the parameters of the approximate posterior to ensure a positive lower bound of this expectation. we further discuss the difference between our approach and previous work. ## extension to cvae given an observation x and its output y, cvae @xcite ) models the conditional distribution p(y|x). the variational lower bound of the conditional log-likelihood is: different from vae, the prior p θ (z|x) in cvae is not fixed, which is also parametrized by a neural network. it is possible to apply another bn on the mean of the prior with a different γ so that the expectation of the kl becomes a constant. however, this lower bound is uncontrollable due to the density of µ 1 + µ 2 is the convolution of their densities, which is intractable. 1 to overcome this issue, we propose to constrain the prior with a fixed distribution. we achieve it by adding another kl between the prior and a known gaussian distribution r(z), i.e. kl(p θ (z|x)||r(z)). instead of optimizing the elbo in eq. 10, we optimize a lower bound of the elbo for cvae: the kl term in the new bound is the sum of kl(q φ (z|x, y)||p θ (z|x)) and kl(p θ (z|x)||r(z)), which can be computed as: the kl term vanishes to 0 when and only when q φ and p θ collapse to r(z), which is the normal distribution. as we explained in section 3.2, kl won't be 0 when we apply bn in q φ . we then prove that when q φ collapses to p θ , the kl term is not the minima (details in section 2 of the appendix) so that kl(q φ (z|x, y)||p θ (z|x)) won't be 0. in this way, we can avoid the posterior collapse in cvae. algorithm 2 shows the training details. algorithm 2 bn-cvae training. sample a mini-batch x, y. ## conclusions and future work in this paper, we tackle the posterior collapse problem when vae is paired with autoregressive decoders. instead of considering the kl individually, we make it follow a distribution d kl and show that keeping the expectation of d kl positive is sufficient to prevent posterior collapse. we propose batch normalized vae (bn-vae), a simple but effective approach to set a lower bound of d kl by regularization the approximate posterior's parameters. our approach can also avoid the recently proposed lagging problem efficiently without additional training efforts. we show that our approach can be easily extended to cvae. we test our approach on three real applications, language modeling, text classification and dialogue generation. experiments show that our approach outperforms strong baselines and is competitive with more complex methods which keeping substantially faster. we leverage the gaussian prior as the example to introduce our method in this work. the key to our approach to be applicable is that we can get a formula for the expectation of the kl. however, it is hard to get the same formula for some more strong or sophisticated priors, e.g., the dirichlet prior. for these distributions, we can approximate them by the gaussian distributions (such as in @xcite ). in this way, we can batch normalize the corresponding parameters. further study in this direction may be interesting.