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: the paper proposes a method for unsupervised disentangling of both discrete and continuous factors of variation in image data. It uses an autoencoder learned by optimising an additive loss composed of Mutual Information (MI) I(x;y,z) between the image x and the discrete+cts latents (y,z) and the reconstruction error. The mutual information is shown to decompose into I(x,y), I(x,z) and TC(y;z), and the I(x,z) is treated in a different manner to I(x,y). With Gaussian p(z|x), and it is shown that I(x,z_k) is maximal when p(z_k) is Gaussian. So KL(p(z_k)||N(0,1)) is optimised in lieu of optimising I(x,z), and I(x,y) (and TC(y;z)) is optimised by using mini-batch estimates of marginal distributions of y (and z). The paper claims improved disentangling of discrete and continuous latents compared to methods such as JointVAE and InfoVAE.