Year: 2018

Abstract:
We propose the Information Maximization Autoencoder (IMAE), an information theoretic approach to simultaneously learn continuous and discrete representations in an unsupervised setting. Unlike the Variational Autoencoder framework, IMAE starts from a stochastic encoder that seeks to map each input data to a hybrid discrete and continuous representation with the objective of maximizing the mutual information between the data and their representations. A decoder is included to approximate the posterior distribution of the data given their representations, where a high fidelity approximation can be achieved by leveraging the informative representations. We show that the proposed objective is theoretically valid and provides a principled framework for understanding the tradeoffs regarding informativeness of each representation factor, disentanglement of representations, and decoding quality.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 INFORMATION MAXIMIZATION REPRESENTATION LEARNING
text: 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).

heading: 3.1 SIMULTANEOUSLY SEEKING INFORMATIVENESS AND DISENTANGLEMENT
text: 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.

heading: 3.2 INFORMATIVE CONTINUOUS REPRESENTATIONS
text: 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.

heading: 3.3 INFORMATIVE DISCRETE REPRESENTATIONS
text: Unlike the continuous representation, the mutual information Iθ(x;y) between a discrete representation and data can be well approximated, given the fact that the cardinality of the space of y is typically low. To be more specific, given N i.i.d samples {xn}Nn=1 of the data, the empirical estimation of Iθ(x;y) under the conditional distribution pθ(y|xn) follows as
Îθ(x;y) = Ĥθ(y)− Ĥθ(y|x) = H ( 1
N
∑N n=1pθ(y|xn) ) − 1 N ∑N n=1H (pθ(y|xn)) . (7)
As shown in Proposition 2, with a suitably large batch of samples, the empirical mutual information Îθ(x;y) is a good approximation to Iθ(x;y). This enables us to optimize Iθ(x;y) in a theoretically justifiable way that is amenable to stochastic gradient descent with minibatches of data. Proposition 2. Let y be a discrete random variable that belongs to some categorical class C. Assume the marginal probabilities of the true and the predicted labels are bounded below, i.e. pθ(y), p̂θ(y) ∈ [1/(CK2), 1] for all y ∈ C with some constant C > 1. Then for any δ ∈ (0, 1),
P (∣∣∣Iθ(x;y)− Îθ(x;y)∣∣∣ ≤ K2 (max{logCK2 − 1, 1}+ e)√ log(2K2/δ) 2N ) ≥ 1− 2δ . (8)
Here N denotes the number of samples used to establish Îθ(x;y) according to Eq (7).
Therefore, to maximize the mutual information Iθ(x;y), we consider the following:
max Lθ(y) := Îθ(x;y). (9) Maximizing the the mutual information Iθ(x;y) provides a natural way to learn discrete categorical representations. To see this, notice that Iθ(x;y) contains two fundamental quantities, the category balance term Hθ(y) and the category separation term Hθ(y|x). In other words, maximizing Iθ(x;y) trades off uniformly assigning data over categories and seeking highly confident categorical identity for each sample x. The maximum is achieved if pθ(y|x) is deterministic while the marginal distribution pθ(y) is uniform, that is Hθ(y|x) = 0 and Hθ(y) = logK2.
Overall Objective As a summary of (3) (6) and (9), our overall objective is
β ( max θ,φ Lθ(z) + Lθ(y)−DKL [ p(z)||ΠK1k=1p(zk) ]) + Epθ(x,y,z) [log qφ(x|y, z)] .
The first three terms associate with our information maximization objective, while the last one aims at better approximation of the posterior pθ(x|y, z). A better balance between these two targets can be achieved by weighting them differently. One the other hand, the informativeness of each latent factor can be optimized through Lθ(z) and Lθ(y), while statistically independent latent continuous factors can be promoted by minimizing the total correlation term DKL [ p(z)||ΠK1k=1p(zk) ] . Therefore, trade-offs can be formalized regarding the informativeness of each latent factor, disentanglement of the representation, and better decoding quality. This motivates us to consider the following objective, let β, γ > 0,
max θ,φ LIMAE := Epθ(x,y,z) [log qφ(x|y,z)] + βLθ(y) + βLθ(z)− γDKL
[ pθ(z)||ΠK1k=1pθ(zk) ] . (10)

heading: 4 EXPERIMENTAL RESULTS
text: We compare IMAE against various VAE based approaches that are summarized in Figure 1. We would like to demonstrate that IMAE can (i) successfully learn a hybrid of continuous and discrete representations, with y matching the intrinsic categorical information ytrue well and z capturing the disentangled feature information shared across categories; (ii) outperform the VAE based models by achieving a better trade-off between representation interpretability and decoding quality. We choose the priors r(z) and r(y) to be the isotropic Gaussian distribution and uniform distribution respectively. Detailed experimental settings are provided in Appendix G.

heading: 4.1 INFORMATIVE REPRESENTATIONS YIELD BETTER INTERPRETABILITY
text: We first qualitatively demonstrate that informative representations can yield better interpretability. For the continuous representation, Figure 2 validates Proposition 1 by showing that, with roughly same amount of variance for each latent variable zk, those achieving high mutual information with the data have mean values µk(x) of the conditional probability p(zk|x) disperse across data samples and variances σk(x) decrease to small values for all data samples. As a qualitative evaluation, we traverse latent dimensions corresponding with different levels of I(x, zk). As seen in Figure 2(b)-(d), informative variables in the continuous representation have uncovered intuitive continuous factors of the variation in the data, while the factor z8 has no mutual information with the data and shows no variation. We observe the same phenomenon for the discrete representation y in Figure 2(e)&(f), which were obtained with two different values of β and γ, where the more informative one discovers matches the natural labels better. This provides further evidence for that interpretable latent factors can be attained by maximizing the mutual information between the representations and the data.

heading: 4.2 QUANTITATIVE COMPARISONS
text: In this section, we perform quantitative evaluations on MNIST (LeCun and Cortes, 2010), Fashion MNIST (Xiao et al., 2017) and dSprites (Matthey et al., 2017). We show that IMAE achieves better interpretability vs. decoding quality trade-off.
Unsupervised learning of discrete latent factor Before we present our main results, we first describe an assumption that we make on the discrete representations. For the discrete representation, a reasonable assumption is that the conditional distribution p(y|x) should be locally smooth so that the data samples that are close on their manifold should have high probability of being assigned to the same category (Agakov, 2005). This assumption is crucial for using neural networks to learn discrete representations, since it’s easy for a high capacity model to learn a non-smooth function p(y|x) that can abruptly change its predictions without guaranteeing similar data samples will be mapped to similar y. To remedy this issue, we adopt the virtual adversarial training (VAT) trick proposed by (Miyato et al., 2016) and augment Lθ(y) as follows:2
max Lθ(y) := Îθ(x;y)− Ep̂(x) [ max‖η‖≤ H (pθ(y|x); pθ(y|x + η)) ] . (11)
The second term of RHS regularizes pθ(y|x) to be consistent within the norm ball of each data sample so as to maintain the local smoothness of the prediction model. For fair comparison, we augment all four methods with VAT. As demonstrated in Appendix D, using VAT is essential for all of them except β-VAE to learn interpretable discrete representations.

heading: 4.2.1 MNIST AND FASHION MNIST
text: We start by evaluating different methods on MNIST and Fashion MNIST, for which we train over a range of β values (we set γ = 2β for IMAE).
Discrete representations For the discrete representations, by simply pushing the conditional distribution p(y|x) towards the uniform distribution r(y), β-VAE sacrifices the mutual information I(x;y) and hence struggles in learning interpretable discrete representation even with VAT. As a comparison, InfoVAE performs much better by dropping I(x;y) from ELBO. For data that are distinctive enough between categories (MNIST), with large β values InfoVAE performs well by uniformly distributing the whole data over categories through minimizing DKL(p(y)||r(y)) while simultaneously encouraging local smoothness with VAT. However, InfoVAE struggles with less distinctive data (Fashion-MNIST), where it cannot give fairly confident category separation by only
2In this paper, we set = 1 across datasets. VAT can be effectively approximated by a pair of forward and backward passes (Miyato et al., 2016).
requiring local smoothness. In contrast, IMAE achieves much better performance by explicitly encouraging confident category separation via minimizing the conditional entropy H(y|x), while using VAT to maintain local smoothness so as to prevent overfitting of neural network. Although JointVAE performs much better than β-VAE by pushing the upper bound of I(x;y) towards a progressively increasing target value Cy , we found it can easily get stuck at some bad local optima where I(x;y) is comparatively large while the accuracy is poor. A heuristic is that once JointVAE enters the local region of a local optima, progressively increasing Cy only induces oscillation within that region. 3
Informativeness, interpretability and decoding quality As illustrated in Figure 1, by using large β values, β-VAE sacrifices more mutual information between the data and its representations, which in turn (see Figure 3) results in less informative representations followed by poor decoding quality. In contrast, the other three methods can remedy this issue to different degrees, and hence attains better trade-off regarding informativeness of latent representations and decoding quality. Compared to JointVAE and InfoVAE, IMAE is more capable of learning discrete presentations over a wide range of β, γ values, which implies less overlap between the manifolds of different categories is induced. As a result, IMAE is expected to yield better decoding quality for each category. Although InfoVAE and JointVAE can also learn comparatively good discrete representations when using large and small β values respectively, the corresponding results of these two regions associate with either poor decoding quality or much lower disentanglement score (see section 4.2.2). In contrast, IMAE consistently performs well with different hyperparameters, especially in the region of interest where the decoding quality as well as the informativeness of latent representations are good enough.

heading: 4.2.2 2D SHAPES
text: In this section, we quantitatively evaluate the disentanglement capability of IMAE on dSprites where the ground truth factors of both continuous and discrete representaions are available. We use the disentanglement metric proposed by (Chen et al., 2018), which is defined in terms of the gap between the top two empirical mutual information of each latent representation factor and a ground truth factor. The disentanglement score is defined as the weighted average of the gaps. A high disentanglement score implies that each ground truth factor associates with one single representation factor that is more informative than the others, i.e., the learnt representation factors are more disentangled.4
Figure 5 shows that, with large β values, β-VAE penalizes the mutual information too much and this degrades the usefulness of representations. while all other three methods achieve higher disentanglement score with better decoding quality. For JointVAE, higher β values push the upper bound of mutual information converges to the prefixed target value, it therefore can maintain more mutual
3More results of JointVAE can be found in Appendix F. 4Although the truth discrete factor is provided, we evaluate the disentanglement quality only in terms of the continuous representations since the pixel-wise difference between different categories are very small. The results of considering the disentanglement score regrading both y and z is provided in Appendix E.
information between the data and the whole latent representations and give better decoding quality. However, the disentanglement quality is poor in this region, which implies that simply restricting the overall capacity of the latent representations is not enough for learning disentangled representations. While InfoVAE yields comparatively better disentanglement score by pushing the marginal joint distribution of the representations towards a factorial distribution harder with large values of β, the associated decoding quality and informativeness of latent representations are both poor. In contrast, IMAE is capable of achieving better trade-off between the disentanglement score and the decoding quality in the region of interest where the decoding quality as well as the informativeness are fairly good. We attribute this to the effect of explicitly seeking for statistically independent latent factors by minimizing the total correlation term in our objective.

heading: 5 CONCLUSION
text: We have proposed IMAE, a novel approach for simultaneously learning the categorical information of data while uncovering latent continuous features shared across categories. Different from VAE, IMAE starts with a stochastic encoder that seeks to maximize the mutual information between data and their representations, where a decoder is used to approximate the true posterior distribution of the data given the representations. This model targets at informative representations directly, which in turn naturally yields an objective that is capable of simultaneously inducing semantically meaningful representations and maintaining good decoding quality, which is further demonstrated by the numerical results.
Unsupervised joint learning of disentangled continuous and discrete representations is a challenging problem due to the lack of prior for semantic awareness and other inherent difficulties that arise in learning discrete representations. This work takes a step towards achieving this goal. A limitation of our model is that it pursues disentanglement by assuming or trying to encourage independent scalar latent factors, which may not always be sufficient for representing the real data. For example, data may exhibit category specific variation, or a subset of latent factors might be correlated. This motivates us to explore more structured disentangled representations; one possible direction is to encourage group independence. We leave this for future work.

heading: A PROOF OF SECTION 3
text: Balance between posterior inference fidelity and information maximization Notice that we can rewrite the mutual information between the data x and its representations as the following,
Iθ(x;y, z) = H(x) + Epθ(x,y,z) [log qφ(x|y, z)] +DKL [pθ(x|y, z)||qφ(x|y, z)] . (12) It then follows that,
Iθ(x;y, z)−DKL (pθ(x|y, z)||qφ(x|y, z)) = H(x) + Epθ(x,y,z) [log qφ(x|y, z)] (13) Since H(x) is independent of the optimization procedure, we have the following,
max βIθ(x;y, z)−DKL (pθ(x|y, z)||qφ(x|y, z)) , β > 1 ⇒ max (β − 1)Iθ(x;y, z) + Epθ(x,y,z) [log qφ(x|y, z)] (14)
where β trade-off the informativeness of the latent representation and generation fidelity.
Decomposition of Iθ(x;y, z) Let b = (z,y) denote the joint random variable consisting of the continuous random variable b and discrete random variable y.
Note that Iθ(x;y, z) = Iθ(x; b) can be written as: Iθ(x; b) = − ∫ X p(x) ∫ Z pθ(b|x) log pθ(b)dbdx + ∫ X p(x) ∫ Z pθ(b|x) log pθ(b|x)dbdx
= − ∫ Z pθ(b) log pθ(b)db + ∫ X p(x) ∫ Z pθ(b|x) log pθ(b|x)dbdx . (15)
The second term in Eq (15) has the form:∫ X p(x) ∫ Z pθ(b|x) log pθ(b|x)dbdx ϑ1= K1+1∑ k=1 ∫ X p(x) ∫ Z pθ(b|x) log pθ(bk|x)dbdx
= K1+1∑ k=1 Hθ(bk|x) , (16)
where ϑ1 follows by the assumption that pθ(b|x) is factorial. For the first term in Eq (15), we have:∫
Z pθ(b) log pθ(b)db = ∫ Z pθ(b) log
pθ(b)
ΠK1+1k=1 pθ(bk) db + K1+1∑ k=1 ∫ Z pθ(b) log pθ(bk)db
= DKL ( pθ(b)||ΠK1+1k=1 pθ(bk) ) − K1+1∑ k=1 Hθ(bk) . (17)
Substituting Eqs (16) & (17) into Eq (15) yields the result:
Iθ(x;y, z) = Iθ(x; b) = Hθ(bk)−DKL ( pθ(b)||ΠK1+1k=1 pθ(bk) ) − K1+1∑ k=1 Hθ(bk|x)
= K1+1∑ k=1 Iθ(x; bk)−DKL ( pθ(b)||ΠK1+1k=1 pθ(bk) ) = Iθ(x;y) +
K1∑ k=1 Iθ(x; zk)−DKL ( pθ(y, z)||pθ(y)ΠK1k=1pθ(zk) ) . (18)
Since y and z are assumed to be marginally independent, i.e., pθ(y; z) = pθ(y)pθ(z), then
Iθ(x;y) + K1∑ k=1 Iθ(x; zk)−DKL ( pθ(y, z)||pθ(y)ΠK1k=1pθ(zk) ) = Iθ(x;y) +
K1∑ k=1 Iθ(x; zk)−DKL ( pθ(z)||ΠK1k=1pθ(zk) ) . (19)
Proof of proposition 1
Proof. We start with computing the expectation of zk: Eθ [zk] = ∫ Zk zk ∫ X pθ(zk|x)p(x)dxdzk = ∫ X p(x) ∫ Zk zkpθ(zk|x)dzkdx
= ∫ X p(x)µk(x)dx = Ex [µk(x)] . (20)
Then the variance of zk followed as:
Varθ [zk] = ∫ Zk z2k ∫ X pθ(zk|x)p(x)dxdzk − Ex [µk(x)]2
= ∫ X p(x) ∫ Zk z2kpθ(zk|x)dzkdx− Ex [µk(x)] 2
= ∫ X p(x) [ σ2k(x) + µk(x) 2 ] dx− Ex [µk(x)]2
= Ex [ σ2k(x) ] + Varx [µk(x)] . (21)
Note that
Iθ(x; zk) = Hθ(zk)−Hθ(zk|x) , (22)
for which we have the following, Hθ(zk|x) = − ∫ X p(x) ∫ Zk pθ(zk|x) log pθ(zk|x)dzdx
= 1
2 ∫ X p(x) log ( 2πeσ2k(x) ) dx
= 1
2
( log(2πe) + Ex [ log σ2k(x) ]) . (23)
For the entropy of zk, we leverage the fact that Hθ(zk) is upper bounded by the entropy of a Gaussian distributed random variable with the same mean and variance, that is
Hθ(zk) ≤ 1
2
( log 2πe+ log ( Ex [ σ2k(x) ] + Varx [µk(x)] )) (24)
Substituting Eqs (23) & (24) into Eq (22) completes the proof.
Proof of proposition 2 Proof. Let p̂θ(y) = 1N ∑N n=1 pθ(y|xn) denote the Monte Carlo estimator of the true probability
pθ(y) = ∫ X p(x)pθ(y|x)dx = Ex [pθ(y|x)]. Note that pθ(y|x) ∈ [0, 1] for all x ∈ X , then applying the Hoeffding’s inequality for bounded random variables [Theorem 2.2.6, (Vershynin, 2018)] yields,
P (|p̂θ(y)− pθ(y)| ≥ t) = P (∣∣∣∣∣ 1N N∑ n=1 pθ(y|xn)− Ex [pθ(y|x)] ∣∣∣∣∣ ≥ t ) ≤ 2 exp ( −2Nt2 ) (25)
Let δ′ = 2 exp ( −2Nt2 ) , it then follows,
P ( |p̂θ(y)− pθ(y)| < √ log(2/δ′)
2N
) ≥ 1− δ′ (26)
Given Eq (26), we first establish the concentration results of the entropy Hp̂θ (y) with respect to the empirical distribution p̂θ(y). Assume For all y ∈ C, we have pθ(y), p̂θ(y) bounded below by 1/(CK2) for some fixed constant C > 1. This assumption is practical since the distributions of true
data and predicted data are approximately uniform and therefore pθ(y), p̂θ(y) ≈ 1/K2 for all y ∈ C. Consider the function t log t, with derivative 1 + log t ∈ [1− logCK2, 1] for t ∈ [1/(CK2), 1],
|p̂θ(y) log p̂θ(y)− pθ(y) log pθ(y)| = ∣∣∣∣∣ ∫ p̂θ(y) pθ(y) (1 + log t)dt ∣∣∣∣∣ ≤
∣∣∣∣∣ ∫ p̂θ(y) pθ(y) |1 + log t|dt ∣∣∣∣∣ ≤ ∣∣∣∣∣ ∫ p̂θ(y) pθ(y) max{logCK2 − 1, 1}dt ∣∣∣∣∣ ≤ max{logCK2 − 1, 1}|p̂θ(y)− pθ(y)| (27)
Summing over C gives∣∣∣Ĥθ(y)−Hθ(y)∣∣∣ ≤ K2 max{logCK2 − 1, 1}|p̂θ(y)− pθ(y)| . (28) Let δ = K2δ′, then Eq (26) together with Eq (28) yield the following,
P (∣∣∣Ĥθ(y)−Hθ(y)∣∣∣ < K2 max{logCK2 − 1, 1}√ log(2K2/δ) 2N ) ≥ 1− δ (29)
Next we are going to bound the divergence between Ĥθ(y|x) and Hθ(y|x) which are defined as,
Ĥθ(y|x) = − 1
N N∑ n=1 ∑ y pθ(y|xn) log pθ(y|xn),
Hθ(y|x) = − ∫ x∈X ∑ y pθ(y|x) log pθ(y|x) .
Note that h log h ∈ [−1/e, 0] for all h ∈ [0, 1], then again applying [Theorem 2.2.6, (Vershynin, 2018)] yields,
P (∣∣∣∣∣ 1N N∑ n=1 pθ(y|xn) log pθ(y|xn)− Ep(x) [pθ(y|x) log pθ(y|x)] ∣∣∣∣∣ < t ) ≤ 2 exp ( −2t2e2N ) (30)
Following the similar arguments as before, let δ′ = 2 exp ( −2t2e2N ) , then
P (∣∣∣∣∣ 1N N∑ n=1 pθ(y|xn) log pθ(y|xn)− Ep(x) [pθ(y|x) log pθ(y|x)] ∣∣∣∣∣ < √ e2 log(2/δ′) 2N ) ≤ δ′
(31)
Now let δ = K2δ′, then applying the union bound we have
|Ĥθ(y|x)−Hθ(y|x)| ≤ ∑ y∈C ∣∣∣∣∣ 1N N∑ n=1 pθ(y|xn) log pθ(y|xn)− Ep(x) [pθ(y|x) log pθ(y|x)] ∣∣∣∣∣ ≤ K2 √ e2 log(2K2/δ)
2N (32)
hold with probability 1− δ. Conclude from Eqs (29) & (32), we have∣∣∣Iθ(x;y)− Îθ(x;y)∣∣∣ ≤ ∣∣∣Hθ(y)− Ĥθ(y)∣∣∣+ ∣∣∣Hθ(y|x)− Ĥθ(y|x)∣∣∣
= K2 (max{logCK2 − 1, 1}+ e) √ log(2K2/δ)
N . (33)
hold with probability at least 1− 2δ.

heading: B APPROXIMATION OF THE MARGINAL DISTRIBUTION
text: Computing the marginal distributions of the continuous representations z and zk requires the entire dataset, e.g., pθ(z) = ∫ X pθ(z,x)dx ≈ 1 N ∑N i=1 pθ(z|x(i)). To scale up our method to large
datasets, we propose to estimate based on the minibatch data, e.g., pθ(z) ≈ 1B ∑B i=1 pθ(z|x(i)).
Now consider the entropy H(z) of z, which we approximate in the following way,
H(z) = Ez[log p(z)] ≈ 1
B B∑ i=1 log p(z(i)) = 1 B B∑ i=1 log 1 B B∑ j=1 pθ(z (i)|x(j)) . (34)
We estimate the integral of z by sampling z ∼ pθ(z|xi) and perform the Monte Carlo approximation. Although we minimize the unbiased estimator of the lower bound of the KL divergence, the term inside the logarithm is a summation of probability densities of Gaussians. In particular, we record the distribution of the variances output by our encoder and observe that the mean of the variances of the Gaussians is bounded between 0.2 and 2, which implies that the values of probability densities do not range in a large scale. Since logarithm is locally affine, we argue that our bound in (34) is tight. Other quantities involved in our objective function (10) are estimated in a similar fashion.

heading: C CONNECTIONS TO VAE
text: In VAE, they assume a generative model specified by a stochastic decoder pθ(x|z), taking the continuous representation as an example, and seek an encoder qφ(z|x) as a variational approximation of the true posterior pθ(z|x). The model is fitted by maximizing the evidence lower bound (ELBO) of the marginal likelihood,
Ex [log pθ(x)] ≥ L(x, θ, φ) = Eqφ(z|x) [log pθ(x|z)]− Ex [DKL (qφ(z|x)||r(z))] . (35)
Here the KL divergence term can be further decomposed as (Hoffman and Johnson, 2016),
Ex [DKL (qφ(z|x)||r(z))] = Iθ(x; z) + Ex [DKL (qφ(z)||r(z))] . (36)
That is, minimizing the KL divergence also penalizes the mutual information Iθ(x; z), thus reduces the amount of information z has about x. This can make the inference task qφ(z|x) hard and lead to poor reconstructions of x as well. Many recent efforts have been focused on resolving this problem by revising ELBO. Although approaches differ, it can be summarized as either dropping the mutual information term in Eq (36), or encouraging statistical independence across the dimensions of z by increasing the penalty on the total correlation term extracted from the KL divergence DKL (qφ(z)||r(z)) with respect to qφ(z). However, these approaches either result in an invalid lower bound for the VAE objective, or cannot avoid minimizing the mutual information Iθ(x; z) between the representation and the data.
In contrast, IMAE starts with a stochastic encoder pθ(z|x) and aims at maximizing the mutual information between the data x and the representations z from the very beginning. By following the constraints which are naturally implied by the objective in order to avoid degenerated solutions, IMAE targets at both informative and statistical independent representations. On the other hand, in IMAE the decoder qφ(x|z) serves as a variational approximation to the true posterior pθ(x|z). As we will show in Section 4, being able to learn more interpretable representations allows IMAE to reconstruct and generate data with better quality.
D VAT STABILIZES THE LEARNING OF CATEGORICAL REPRESENTATIONS

heading: E DISENTANGLEMENT QUALITY WITH RESPECT TO BOTH CONTINUOUS AND DISCRETE REPRESENTATIONS ON 2D SHAPES
text: See figure 7.

heading: F MORE RESULTS ON JOINTVAE
text: See figure 8.

heading: G EXPERIMENTAL SETTINGS
text: Table 1: Encoder and Decoder architecture for MNIST and Fashion MNIST.
Encoder Decoder Input vectorized 28× 28 grayscale image Input y ∈ R10 and z ∈ R10
FC. 500 BatchNorm ReLU FC. 500 ReLU FC. 2× 500 BatchNorm ReLU FC. 500 ReLU FC. 20 (µz , log σz) + 10 (py) FC. 28× 28 Sigmoid
Table 2: Encoder and Decoder architecture for dSprites.
Encoder Decoder Input vectorized 64× 64 grayscale image Input y ∈ R3 and z ∈ R10
FC. 1200 ReLU FC. 1200 ReLU FC. 1200 ReLU FC. 1200 ReLU
FC. 2× 1200 ReLU FC. 1200 ReLU FC. 20 (µz , log σz) + 3 (py) FC. 28× 28 Sigmoid
Training procedure:
• MNIST & Fashion MNIST: We use momentum to train all models. The initial learning rate is set as 1e-3, and we decay the learning rate by 0.98 every epoch.
• dSprites: We use Adam to train all models. The learning rate is set as 1e-3.

Review summaries:
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.

* 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.