INFORMATION MAXIMIZATION AUTO-ENCODING

1 INTRODUCTION . A central tenet for designing and learning a model for data is that the resulting representation should be compact yet informative . Therefore , the goal of learning can be formulated as finding informative representations about the data under proper constraints . Generative latent variable models are a popular approach to this problem , where a model parameterized by θ of the form pθ ( x ) = ∫ pθ ( x|z ) p ( z ) dz is used to represent the relationship between the data x and the low dimensional latent variable z . The model is optimized by fitting the generative data distribution pθ ( x ) to the training data distribution p̂ ( x ) , which involves maximizing the likelihood for θ . Typically , this model is intractable even for moderately complicated functions pθ ( x|z ) with continuous z . To remedy this issue , variational autoencoder ( VAE ) ( Kingma and Welling , 2013 ; Rezende et al. , 2014 ) proposes to maximize the evidence lower bound ( ELBO ) of the marginal likelihood objective . However , as was initially pointed out in ( Hoffman and Johnson , 2016 ) , maximizing ELBO also penalizes the mutual information between data and their representations . This in turn makes the representation learning even harder . Many recent efforts have focused on resolving this problem by revising ELBO . Generally speaking , these works fall into two lines . One of them targets “ disentangled representations ” by encouraging the statistical independence between representation components ( Higgins et al. , 2016 ; Kim and Mnih , 2018 ; Gao et al. , 2018 ; Chen et al. , 2018 ; Esmaeili et al. , 2018 ) , while the other line of work seeks to control or encourage the mutual information between data and their representations ( Mary Phuong , 2018 ; Burgess et al. , 2018 ; Alemi et al. , 2017 ; Dupont , 2018 ; Zhao et al. , 2017 ) . However , these approaches either result in an invalid lower bound for the VAE objective or can not avoid sacrificing the mutual information . Instead of building upon the generative latent variable model , we start with a stochastic encoder pθ ( z|x ) and aim at maximizing the mutual information between the data x and its representations z . In this setting , a reconstruction or generating phase can be obtained as the variational inference of the true posterior pθ ( x|z ) . By explicitly seeking for informative representations , the proposed model yields better decoding quality . Moreover , we show that the information maximization objective naturally induces a balance between the informativeness of each latent factor and the statistical independence between them , which gives a more principled way to learn semantically meaningful representations without invalidating ELBO or removing individual terms from it . Another contribution of this work is proposing a framework for simultaneously learning continuous and discrete representations for categorical data . Categorical data are ubiquitous in real-world tasks , where using a hybrid discrete and continuous representation to capture both categorical information and continuous variation in data is more consistent with the natural generation process . In this work , we focus on categorical data that are similar in nature , i.e. , where different categories still share similar variations ( features ) . We seek to learn semantically meaningful discrete representations while maintaining disentanglement of the continuous representations that capture the variations shared across categories . We show that , compared to the VAE based approaches , our proposed objective gives a more natural yet effective way for learning these hybrid representations . 2 RELATED WORK . Recently , there has been a surge of interest in learning interpretable representations . Among them , β-VAE ( Higgins et al. , 2016 ) is a popular method for learning disentangled representations , which modifies ELBO by increasing the penalty on the KL divergence between the variational posterior and the factorized prior . However , by using large weight for the KL divergence term , β-VAE also penalizes the mutual information between the data and the latent representations more than a standard VAE does , resulting in more severe under utilization of the latent representation space . Several follow up works propose different approaches to address the limitations of β-VAE . ( Dupont , 2018 ; Alemi et al. , 2017 ; Burgess et al. , 2018 ; Mary Phuong , 2018 ) propose to constrain the mutual information between the representations and the data by pushing its upper bound , i.e. , the KL divergence term in ELBO , towards a progressively increased target value . However , specifying and tuning this target value can itself be very challenging , which makes this method less practical . Moreover , this extra constraint results in an invalid lower bound for the VAE objective . Alternatively , ( Zhao et al. , 2017 ) drops the mutual information term in ELBO . By pushing only the aggregated posterior towards a factorial prior , they implicitly encourage independence across the dimensions of latent representations without sacrificing the informativeness of the representations . However , simply removing the mutual information term also violates the lower bound of the VAE objective . Another relevant line of work ( Gao et al. , 2018 ; Kim and Mnih , 2018 ; Chen et al. , 2018 ; Esmaeili et al. , 2018 ) seek to learn disentangled representations by explicitly encouraging statistical independence between latent factors . They all propose to minimize the total correlation term of the latent representations , either augmented as an extra term to ELBO or obtained by reinterpreting or re-weighting the terms in the VAE objective , as a way to encourage statistical independence between the representation components . In contrast , we show that our information maximization objective inherently contains the total correlation term while simultaneously seeking to maximize the informativeness of each representation factor . In this paper , we introduce a different perspective to the growing body of the VAE based approaches for unsupervised representation learning . Starting by seeking informative representations for the data , we follow a more intuitive way to maximize the mutual information between the data and the representations . Moreover , we augment the continuous representation with a discrete one , which allows more flexibilities to model real world data that are generated from different categories . We invoke the information maximization principle ( Linsker , 1988 ; Bell and Sejnowski , 1995 ) with proper constraints implied by the objective itself to avoid degenerate solutions . The proposed objective gives a theoretically elegant yet effective way to learn semantically meaningful representations . 3 INFORMATION MAXIMIZATION REPRESENTATION LEARNING . Given data x ∈ Rd , we consider learning a hybrid continuous-discrete representation , denoted respectively with variables z ∈ RK1 and y ∈ { 1 , . . . , K2 } , using a stochastic encoder parameterized by θ , i.e. , pθ ( y , z|x ) . We seek to learn compact yet semantically meaningful representations in the sense that they should be low dimensional but informative enough about the data . A natural approach is to maximize the mutual information ( Cover and Thomas , 2012 ) Iθ ( x ; y , z ) between the data and its representations under the constraint K1 , K2 d. Here the mutual information between two random variables , e.g. , x and z , is defined as Iθ ( x ; z ) = Hθ ( z ) − Hθ ( z|x ) , where Hθ ( z ) = −Epθ ( z ) [ log pθ ( z ) ] is the entropy of z and Hθ ( z|x ) = −Epθ ( x , z ) [ log pθ ( z|x ) ] is the conditional entropy of z given x . The mutual information can be interpreted as the decrease in uncertainty of one random variable given another random variable . In other words , it quantifies how much information one random variable has about the other . A probabilistic decoder qφ ( x|y , z ) is adopted to approximate the true posterior pθ ( x|y , z ) , which can be hard to estimate or even intractable . The dissimilarity between them is optimized by minimizing the KL divergence DKL ( pθ ( x|y , z ) ||qφ ( x|y , z ) ) . In summary , IMAE considers the following , maximizeθ , φ β0Iθ ( x ; y , z ) −DKL ( pθ ( x|y , z ) ||qφ ( x|y , z ) ) . ( 1 ) Given that H ( x ) is independent of the optimization procedure , we can show that optimizing ( 1 ) is equivalent to optimize the following1 , maximizeθ , φ Iθ ( x ; y , z ) + Epθ ( x , y , z ) [ log qφ ( x|y , z ) ] , β = β0 − 1 > 0 . ( 2 ) We set β > 0 to balance between maximizing the informativeness of latent representations and maintaining the decoding quality . The second term is often interpreted as the “ reconstruction error ” which can be optimized using the reparameterization tricks proposed by ( Kingma and Welling , 2013 ) and ( Jang et al. , 2016 ) for continuous representation z and discrete representation y respectively . Now we introduce proper method to optimize the first term Iθ ( x ; y , z ) in ( 2 ) . 3.1 SIMULTANEOUSLY SEEKING INFORMATIVENESS AND DISENTANGLEMENT . We first show that Iθ ( x ; y , z ) inherently involves two keys terms that quantify the informativeness of each representation factor and the statistical dependence between these factors . Assuming the conditional distribution of the representation ( y , z ) given x is factorial , we also assume the marginal distribution of y and z are independent , i.e. , pθ ( y , z ) = pθ ( y ) pθ ( z ) , then1 Iθ ( x ; y , z ) = Iθ ( x ; y ) + ∑K1 k=1Iθ ( x ; zk ) −DKL ( pθ ( z ) ∣∣∣∣ΠK1k=1pθ ( zk ) ) . ( 3 ) The first two terms of the RHS quantify how much information each latent factor , i.e. , y or zk , carry about the data . The last term is known as the total correlation of z ( Watanabe , 1960 ) , which quantifies the statistical independence between the continuous latent factors and achieves the minimum if and only if they are independent of each other . As is implied by ( 3 ) , maximizing Iθ ( x ; y , z ) can be conducted by maximizing informativeness of each latent factor while simultaneously promoting statistical independence between the continuous factors . Various Monte Carlo based sampling strategies have been proposed to optimize the total correlation term ( Chen et al. , 2018 ; Esmaeili et al. , 2018 ) ; in this work we follow this line ( see Appendix B ) . Next we proceed by constructing tractable approximations for Iθ ( x ; zk ) and Iθ ( x ; y ) respectively . 3.2 INFORMATIVE CONTINUOUS REPRESENTATIONS . Without any constraints , the mutual information Iθ ( x ; zk ) between a continuous latent factor and data can be trivially maximized by severely fragmenting the latent space . To be more precise , consider the following proposition . While similar results have likely been established in the information theory literature , we include this proposition to motivate our objective design . Proposition 1 . Suppose the conditional distribution pθ ( z|x ) is a factorial Gaussian distribution with mean µ ( x ) and covariance Σ ( x ) . Let σ ( x ) ∈ RK1 denote the diagonal entries of Σ ( x ) , then Iθ ( x ; zk ) ≤ 1 2 log [ ( Ex [ σ2k ( x ) ] + Varx [ µk ( x ) ] ) ] − 1 2 Ex [ log σ2k ( x ) ] , k = 1 , . . . , K1 . ( 4 ) The equality in ( 4 ) is attained if and only if zk is Gaussian distributed , given which we have Iθ ( x ; zk ) ≥ 1 2 log ( 1 + Varx [ µk ( x ) ] /Ex [ σ2k ( x ) ] ) , k = 1 , . . . , K1 . ( 5 ) Note here both µk ( x ) and σk ( x ) are random variables . The above result implies that zk is more informative about x if it has less uncertainty given x yet captures more variance in data , i.e. , σk ( x ) is small while µk ( x ) disperses within a large space . However , this can result in discontinuity of zk , where in the extreme case each data sample is associated with a delta distribution in the latent space . 1 Detailed derivation is provided in Appendix A . In light of this , we can make what we described above more precise . A vanishing variance of the conditional distribution p ( zk|x ) leads to a plain autoencoder that maps each data sample to a deterministic latent point , which can fragment the latent space in a way that each data sample corresponds with a delta distribution in the latent space pθ ( zk|x ( i ) ) = δ ( z ( i ) k ) . On the other hand , Proposition 1 also implies that controlling the variance σk ( x ) to be finite , Iθ ( x ; zk ) will be maximized by pushing µk ( x ) towards two extremes ( ±∞ ) . To remedy this issue while achieving the upper bound , a natural resolution is to squeeze zk within the domain of a Gaussian distribution with finite mean and variance . By doing so , we can avoid the degenerate solution while achieving a more reasonable trade-off between enlarging the spread of µk ( x ) and maintaining the continuity of z . Therefore , we consider the following as the surrogate for maximizing Iθ ( x ; zk ) , maximize Lθ ( z ) : = − ∑K1 k=1DKL ( pθ ( zk ) ||r ( zk ) ) . ( 6 ) Here r ( zk ) are i.i.d scaled normal distribution with finite variance . That is , we push each pθ ( zk ) towards a Gaussian distribution r ( zk ) by minimizing the KL divergence between them .

Summary: * This paper proposed a principled framework for auto-encoding through information maximization. A novel contribution of this paper is to introduce a hybrid continuous-discrete representation. The authors also related this approach with other related work such as \beta-VAE and info-VAE, putting their work in context. Empirical results show that the learned representation has better trade-off among interpretability and decoding quality.