MAE: Mutual Posterior-Divergence Regularization for Variational AutoEncoders

1 INTRODUCTION . Representation learning , besides data distributions estimation , is a principle component in generative models . The goal is to identify and disentangle the underlying causal factors , to tease apart the underlying dependencies of the data , so that it becomes easier to understand , to classify , or to perform other tasks ( Bengio et al. , 2013 ) . Among these generative models , VAE ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) gains popularity for its capability of estimating densities of complex distributions , while automatically learning meaningful ( low-dimensional ) representations from raw data . VAE , as a member of latent variable models ( LVMs ) , defines the joint distribution between the observed data ( visible variables ) and a set of latent variables by factorizing it as the product of a prior over the latent variables and a conditional distribution of the visible variables given the latent ones ( detailed in §2 ) . VAEs are usually estimated by maximizing the likelihood of the observed data by marginalizing over the latent variables , typically via optimizing the evidence lower bound ( ELBO ) . By learning a VAE from the data with the appropriate hierarchical structure of latent variables , the hope is to uncover and untangle the causal sources of variations that we are interested in . A notorious problem of VAEs , however , is that the marginal likelihood may not guide the model to learn the intended latent variables . It may instead focus on explaining irrelevant but common correlations in the data ( Ganchev et al. , 2010 ) . Extensive previous studies ( Bowman et al. , 2015 ; Chen et al. , 2017a ; Yang et al. , 2017 ) showed that optimizing the ELBO objective is often completely disconnected from the goal of learning good representations . An extreme case called KL varnishing , happens when using sufficiently expressive decoding distributions such as auto-regressive ones ; the latent variables are often completely ignored and the VAE regresses to a standard auto-regressive model ( Larochelle & Murray , 2011 ; Oord et al. , 2016 ) . This problem has spawned significant interests in analyzing and solving it from both theoretical and practical perspectives . We can only name a few here due to space limits . Some previous work ( Bowman et al. , 2015 ; Sønderby et al. , 2016b ; Serban et al. , 2017 ) attributed the KL varnishing phenomenon to “ optimization challenges ” of VAEs , and proposed training methods including annealing the relative weight of the KL term in ELBO ( Bowman et al. , 2015 ; Sønderby et al. , 2016b ) or adding free bits ( Kingma et al. , 2016 ; Chen et al. , 2017a ) . However , Chen et al . ( 2017a ) pointed out that this phenomenon arises not just due to the optimization challenges , and even if we find the exact solution for the optimization problems , the latent code will still be ignored at optimum . They proposed a solution by limiting the capacity of the decoder and applied PixelCNN ( Oord et al. , 2016 ) with small local receptive fields as the decoder of VAEs to model 2D images , achieving both impressive performance for density estimation and informative latent representations . Yang et al . ( 2017 ) embraced a similar idea and applied VAE to text modeling by using dilated CNN as the decoder . Unfortunately , these approaches require manual and problem-specific design of the decoder ’ s architecture to learn meaningful representations . Other studies attempted to explore alternatives of ELBO . Makhzani et al . ( 2015 ) proposed Adversarial Autoencoders ( AAEs ) by replacing the KL-divergence between the posterior and prior distributions with Jensen-Shannon divergence on the aggregated posterior distribution . InfoVAEs ( Zhao et al. , 2017 ) generalized the Jensen-Shannon divergence in AAEs to a divergence family and linked its objective to the Mutual Information between the data and the latent variables . However , directly optimizing these objectives is intractable , requiring advanced approximate learning methods such as adversarial learning or Maximum-Mean Discrepancy ( Gretton et al. , 2007 ; Dziugaite et al. , 2015 ; Li et al. , 2015 ) . Moreover , these models ’ performance on density estimation significantly falls behind state-of-the-art models ( Salimans et al. , 2017 ; Chen et al. , 2017a ) . In this paper , we propose to tackle the aforementioned representation learning challenges of VAEs by adding a data-dependent regularization to the ELBO objective . Our contributions are three-fold : ( 1 ) Algorithmically , we introduce the mutual posterior-divergence regularization for VAEs , named MAEs ( §3.2 ) , to control the geometry of the latent space during learning by encouraging the learned variational posteriors to be diverse ( i.e . they are favored to be mutually “ different ” from each other ) , to achieve low-redundant , interpretable representation learning . ( 2 ) Theoretically , we establish a close relation between MAE and InfoVAE , by showing that the mutual posterior-devergence regularization maximizes a symmetric version of the KL divergence involved in InfoVAE ’ s mutual information term ( §3.3 ) . ( 3 ) Experimentally , on three benchmark datasets for images , we demonstrate the effectiveness of MAE as a density estimator by state-of-the-art log-likelihood results on MNIST and OMNIGLOT , and comparable result on CIFAR-10 . Moreover , by performing image reconstruction , unsupervised and semi-supervised classification , we show that MAE is also capable of learning meaningful latent representations , even combined with a sufficiently powerful decoder ( §4 ) . 2 VARIATIONAL AUTOENCODERS . 2.1 NOTATIONS . Throughout we use uppercase letters for random variables , and lowercase letters for realizations of the corresponding random variables . Let X ∈ X be the randoms variables of the observed data , e.g. , X is an image or a sentence for image and text generation , respectively . Let P denote the true distribution of the data , i.e. , X ∼ P , and D = { x1 , . . . , xN } be our training sample , where xi , i = 1 , . . . , N , are usually i.i.d . samples of X . Let P = { Pθ : θ ∈ Θ } denote a parametric statistical model indexed by parameter θ ∈ Θ , where Θ is the parameter space . p is used to denote the density of corresponding distribution P . In the literature of deep generative models , deep neural networks are the most widely used parametric models . The goal of generative models is to learn the parameter θ such that Pθ can best approximate the true distribution P . 2.2 VAES . In the framework of VAEs , or general LVMs , a set of latent variables Z ∈ Z are introduced to characterize the hidden patterns of X , and the model distribution Pθ ( X ) is defined as the marginal of the joint distribution between X and Z : pθ ( x ) = ∫ Z pθ ( x , z ) dµ ( z ) = ∫ Z pθ ( x|z ) pθ ( z ) dµ ( z ) , ∀x ∈ X , ( 1 ) where the joint distribution pθ ( x , z ) is factorized as the product of a prior pθ ( z ) over the latent Z , and the “ generative ” distribution pθ ( x|z ) . µ ( z ) is the base measure on the latent space Z . Typically , prior pθ ( z ) is modeled with a simple distribution like multivariate Gaussian , or transforming simple priors to complex ones by normalizing flows and variants ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ; Sønderby et al. , 2016a ) . To learn parameters θ , we wish to minimizes the negative log-likelihood of the parameters : min θ∈Θ 1 N N∑ i=1 − log pθ ( xi ) = min θ∈Θ EP̃ ( X ) [ − log pθ ( X ) ] ( 2 ) where P̃ ( X ) is the empirical distribution derived from training data D. In general , this marginal likelihood is intractable to compute or differentiate directly for high-dimensional latent space Z . Variational Inference ( Wainwright et al. , 2008 ) provides a solution to optimize the evidence lower bound ( ELBO ) an alternative objective by introducing a parametric inference model qφ ( z|x ) : Lelbo ( θ , φ ) = Ep ( X ) [ −Eqφ ( Z|X ) [ log pθ ( X|Z ) ] + KL ( qφ ( Z|X ) ||pθ ( Z ) ) ] = Ep ( X ) [ − log pθ ( X ) + KL ( qφ ( Z|X ) ||pθ ( Z|X ) ) ] ≥ Ep ( X ) [ − log pθ ( X ) ] ( 3 ) where Lelbo could be seen as an autoencoding loss with qφ ( z|x ) being the encoder and pθ ( x|z ) being the decoder , with the first term in the RHS in ( 3 ) as the reconstruction error . 2.3 AUTOENCODING PROBLEM IN VAES . As discussed in Chen et al . ( 2017a ) , without further assumptions , the ELBO objective Lelbo in ( 3 ) may not guide the model towards the intended role for the latent variables Z , or even learn uninformative Z with the observation that the KL term KL ( qφ ( Z|X ) ||pθ ( Z ) ) varnishes to zero . For example , suppose we use an auto-regressive decoder , pθ ( x|z ) = ∏ i pθ ( xi|x < i , z ) , which is sufficiently expressive that it can model the data distribution P ( X ) without the assistance of Z , i.e , pθ ( xi|x < i , z ) = pθ ( xi|x < i ) . In this case , the optimal Z w.r.t . Lelbo is the one independent with X , with the inference model reducing to the prior , i.e. , qφ ( Z|X ) = pθ ( Z ) , ∀X ∈ X . The essential reason of this problem is that , under absolutely unsupervised setting , the marginal likelihood based objectiveLelbo incorporates no ( direct ) supervision on the latent space to characterize the latent variable Z with preferred properties w.r.t . representation learning . The main goal of this work is to explicitly control the geometry of the latent space , in the hope that preferred latent representations would be characterized and selected . 3 MUTUAL POSTERIOR-DIVERGENCE REGULARIZATION . 3.1 GEOMETRIC PROPERTIES OF MEANINGFUL LATENT SPACE . Motivated by the Diversity-Inducing Mutual Angular Regularization ( Xie et al. , 2015 ) which is widely used in LVMs , we propose to regularize the posteriors qφ ( z|x ) of different data x ∈ X , to encourage them to diversely , smoothly , and evenly spread out in the data space X . The intuition is : ( 1 ) To make posteriors mutually diverse from each other , the patterns captured by different posteriors are likely to have less redundancy and hence characterizing and interpreting different data x . ( 2 ) To make posteriors smoothly and evenly distributed in the whole space of X , the shared patterns of similar data points are likely to be captured by their posteriors , to avoid isolating each data point from others . By balancing the diversity and smoothness of the distribution of posteriors , learned representations are encouraged to maintain global structured information , discarding detailed texture of local dependencies in the data .

Summary: This paper presents a new regularization technique for VAEs similar in motivation and form to the work on InfoVAE.  The basic intuition is to encourage different training samples to occupy different parts of z-space, by maximizing the expected KL divergence between pairwise posteriors, which they call Mutual Posterior-Divergence (MPD).  They show that this objective is a symmetric version (sum of the forward and reverse KL) of the Mutual Info regularization used by the InfoVAE.  In practice however, they do not actually use this objective.  They use a different regularization which is based on the MPD loss but they say is more stable because it's always greater than zero, and ensures that all latent dimensions are used.  In addition to the MPD based term, they also add another term which encouraging the pairwise KL-divergences to have a low standard-deviation, to encourge more even spreading over the z-space rather than the clumpy distribution that they observed with only the MPD based term.