Title: THE MUTUAL AUTOENCODER: CONTROLLING INFORMATION IN LATENT CODE REPRESENTATIONS

Abstract: learn probabilistic latent variable models by optimizing a bound on the marginal likelihood of the observed data. Beyond providing a good density model a VAE model assigns to each data instance a latent code. In many applications, this latent code provides a useful high-level summary of the observation. However, the VAE may fail to learn a useful representation when the decoder family is very expressive. This is because maximum likelihood does not explicitly encourage useful representations and the latent variable is used only if it helps model the marginal distribution. This makes representation learning with VAEs unreliable. To address this issue, we propose a method for explicitly controlling the amount of information stored in the latent code. Our method can learn codes ranging from independent to nearly deterministic, while benefiting from decoder capacity. Thus, we decouple the choice of decoder capacity and the latent code dimensionality from the amount of information stored in the code.

1. INTRODUCTION:

Latent variable models are a powerful approach to generative modeling of complicated distributions. In a latent variable model we model a distribution over observables x ∈ X through a hierarchical model, p θ (x) = Z p θ (x|z) p(z) dz.

(1)

In practice, we learn the model parameters θ in (1) using maximum likelihood estimation (MLE),

EQUATION

The recent variational autoencoder (VAE) method, (Kingma & Welling, 2013; Rezende et al., 2014) , provides tractable lower bounds to (2) for deep latent models like (1).

The model (1) has another interesting property: although the latent variables z are never observed, they provide a high-level summary of the observation x. Therefore z could serve as a powerful representation in a number of machine learning tasks (Bengio et al., 2013) . Because (1) can be learned from unlabeled data, we can in principle use latent variable models for unsupervised representation learning, an important building block in artificial intelligence systems.

But when does unsupervised learning lead to useful representations? We argue that trusting (2) alone is not enough to ensure that z stores useful information. 1 In particular, the amount of information stored in z depends on the expressiveness of the model p θ (x|z) with respect to the true data distribution (Chen et al., 2017) . In practice this has made the VAE approach difficult to work with in important applications such as natural language processing and for modeling discrete data.

We propose a solution in the form of explicit control of information flow in latent variable models. This is illustrated in Figure 1 : optimizing log p θ (x) as in VAEs will lead to a model p which achieves a certain mutual information between x and z but this amount of information is difficult to predict and in fact may be zero (Chen et al., 2017; Zhao et al., 2017) . Our mutual autoencoder (MAE) approach forces information flow by ensuring that the estimated model pm achieves a userspecified mutual information M = I pm . Therefore, we precisely control the amount of bits stored in the representation but leave the organization and use of this information to be learned. To control information in this way requires novel algorithms and the rest of the paper discusses our procedure in more detail. 

2. BACKGROUND: VARIATIONAL AUTOENCODERS:

Consider a latent variable model with a data variable x ∈ X and a latent variable z ∈ Z, p θ (z, x) = p(z) p θ (x|z). Given the data x 1 , . . . , x n , we would like to train the model by maximizing the marginal log-likelihood,

EQUATION

where q(x) denotes the empirical distribution of x, q(x) = 1 n n i=1 δ xi (x). However, the integral in (3) is intractable in many applications of interest. The idea behind variational methods is to instead maximize a lower bound L(p θ , q) to the log-likelihood, where

L(p, q) := E x∼q E z∼q(z|x) [log p(x|z)] -D KL (q(z|x) p(z)) .

(ELBO)

Any choice of q(z|x) gives a valid lower bound L(p, q) ≤ L with a gap of

E x∼q [D KL (q(z|x) p(z|x))].

VAEs (Kingma & Welling, 2013; Rezende et al., 2014) replace the per-instance posteriors q(z|x) by an inference network q θ (z|x) that is trained together with p θ (x|z) to jointly maximize L(p θ , q θ ).

For the inference network q θ (z|x), this is equivalent to minimizing D KL (q θ (z|x) p θ (z|x)), so one can think of q θ (z|x) as an approximate posterior for p(z|x).

Being a stochastic mapping of data x to a latent code z, the inference network q θ (z|x) is sometimes called the encoder and, by the same analogy, the generator network p θ (x|z) is called the decoder.

3. VAE OBJECTIVE IS INSUFFICIENT FOR REPRESENTATION LEARNING:

A major appeal of the VAE framework is the ability to learn meaningful latent codes z from unlabeled or only weakly labeled data. Despite numerous promising results on image and video datasets (see e.g. Higgins et al. (2017) , Bouchacourt et al. (2017) , or Goyal et al. (2017) ), the application of VAEs to text has proven challenging. Specifically, Bowman et al. (2015) found that a VAE with an LSTM-based decoder fails to learn a useful latent code when trained naively -the approximate posterior collapses to the prior, q θ (z|x) = p(z), for all inputs x, leading to a near independent relationship between x and z.

A number of works have addressed this problem, mostly treating it as an optimization issue (Sønderby et al., 2016; Kingma et al., 2016; Yeung et al., 2017) . However, as Chen et al. (2017) point out, even if one could optimize exactly, the model would still learn trivial latent codes when using a high-capacity decoder such as an LSTM.

The reason for this is the VAE objective. Note that the log-likelihood (3) is only a function of the marginal distribution p θ (x), whereas representation is an aspect of the joint distribution p θ (x, z).

That is, we approximately optimize a marginal quantity in the hope of producing the desired effect on the joint distribution. This approach is unreliable, although it works when the marginal imposes strong constraints on the joint distribution, such as when the decoder has limited capacity and the model is forced to use the latent structure to reach a high likelihood. However, in the high-capacity regime, e.g. when p θ (x|z) = q(x) is close to achievable, the task of density estimation becomes disconnected from the goal of representation learning (Huszár, 2017) .

We propose an alternative optimization problem that better captures the representation objective.

Our idea is to explicitly control the amount of information stored in the latent code, as measured by the mutual information. In the next sections, we derive our procedure in detail and show that it enables representation learning with powerful decoders.

4. CONTROLLING INFORMATION IN LATENT CODE REPRESENTATIONS:

Our goal is to learn a deep latent variable model p θ (z, x) = p(z) p θ (x|z), while precisely controlling the coupling between the latent code z and the output x. We formalize our goal in the following constrained optimization problem,

EQUATION

EQUATION

where M ≥ 0 is a scalar constant denoting the desired mutual information, z) ], between x and z as encoded in the model p θ (z, x). When M is close to zero, the code z is uninformative about x, whereas a large value of M approaching the entropy as maximum possible value, H p θ (x), should lead to a deterministic relation between x and z. Intermediate values of M will lead to lossy codes z which recover a compressed representation of structure in q(x) that can be most efficiently captured by p θ (x|z).

I p θ (x, z) = E (x,z)∼p θ [log p θ (x,z) p θ (x) p θ (

5. THE MUTUAL AUTOENCODER (MAE):

We now describe our idea for approximately solving the problem (4-5). We call our approach the mutual autoencoder.

Our method relies on two results from the literature: 1) exact penalty functions (Zangwill, 1967) to accommodate the equality constraint (5) into the estimation problem; 2) the variational infomax bound, (Barber & Agakov, 2003) , to approximate the intractable mutual information I p θ (x, z).

We briefly describe each of them and their role in the mutual autoencoder.

5.1. EXACT PENALTY FUNCTIONS:

To approximately solve (4-5), we resort to a classic method for constrained nonlinear optimization, the penalty function method. First proposed by Zangwill (1967) , this method uses a positive penalty constant C > 0 to relax (4-5) to the unconstrained problem

EQUATION

In ( 6), any deviation of I p θ (x, z) from M is penalized linearly. For values of C large enough, an unconstrained maximum of (6) recovers a feasible local maximum of (4-5). In particular, we do not need to take C all the way to infinity, but a finite value is sufficient as shown by Han & Mangasarian (1979) , and the magnitude of this value is determined by the unknown optimal Lagrange multiplier for (5), see Bertsekas (1999, Proposition 4.3.1) . In practice, we observe that values of C in the range [0.1, 10] work best.

We will show experimentally that the above penalty approach is highly effective. However, to leverage the approach, we also need to approximate the intractable mutual information I p θ (x, z).

5.2. VARIATIONAL INFORMATION MAXIMIZATION (INFOMAX INEQUALITY):

The mutual information I p θ (x, z) is intractable because of the complex structure in the conditional distribution p θ (x|z). To overcome this difficulty, we leverage the variational infomax inequality of Barber & Agakov (2003) .

Theorem 1 (Variational infomax bound, Barber & Agakov (2003) ). For any two random variables x and z distributed according to the joint distribution with a density p(x, z) and for any conditional density function r(z|x) we have

EQUATION

Proof. See Appendix B on page 12.

Equality in ( 7) is attained for r(z|x) = p(z|x), so we can write

EQUATION

In our method, we constrain the inner search over r to a parametric class, yielding the following lower bound approximation Îp (x, z) ≤ I p (x, z), to the mutual information:

EQUATION

5.3. THE MUTUAL AUTOENCODER:

We are now in the position to combine the variational infomax bound (9) with the variational likelihood bound (ELBO) to define the following mutual autoencoder objective L m (p, q) to be maximized over the distributions p and q.

L m (p, q) := L(p, q) -C Îp (x, z) -M (10)

EQUATION

We now discuss how to realize this objective via gradient-based methods. The quantity H p (z) is typically a constant because we fix the prior p(z) to a simple distribution such as a multivariate normal distribution. We assume the latent variable is continuous, that is z ∈ Z with Z = R k , so we can use the reparametrization trick (Kingma & Welling, 2013; Rezende et al., 2014) to compute an unbiased and low-variance estimate of the gradients of the expectation over q(z|x).

The difficult part is the derivative of the term max ω E (x,z)∼p [log r ω (z|x)] with respect to p. Denoting by r * p the optimal r ω corresponding to p, one can derive by REINFORCE-style reasoning (Williams, 1992) that

EQUATION

)

There are two difficulties here: 1) The last term in (13) that arises from the dependence of r * p on p is difficult to compute. 2) Evaluating r * p requires solving a separate optimization program in ω, which is too expensive to do for each gradient evaluation.

To address the first issue, we note that in the non-parametric limit, as r * p (z|x) ≈ p(z|x), the problematic term vanishes. We therefore ignore it, which is equivalent to treating r * p as independent of p during back-propagation. For the second issue, we keep a running estimate of the optimal r * p and perform a single gradient update to it whenever p is updated. This gives rise to the practical procedure described in Algorithm 1. A more detailed derivation of the method can be found in Appendix A.1. Note that unlike (13), the gradient of the term E (x,z)∼p θ [log r ω (z|x)] with respect to θ can be efficiently approximated via the reparametrization trick (if x is continuous) or REINFORCE (if x is discrete). We provide a derivation of the necessary gradients in Appendix A.2.

6. DISCRETE DATA REQUIRES FLEXIBLE ENCODER DISTRIBUTIONS:

Before presenting experimental results, we discuss the choice of encoder distribution, which we find to be critical to making our method work. More generally, it turns out that flexible encoder distributions are essential for VAE representation learning of discrete data. Below we give the formal statement and discuss its practical implications.

Theorem 2. Consider a VAE model applied to discrete data. Assume that both the prior p(z) and the approximate posterior q(z|x) are Gaussian and that the decoder p(x|z) is powerful enough to model the true marginal q(x). Then the independent configuration given by q(z|x) = p(z) and p(x|z) = q(x) is the only global optimum of the VAE objective.

Proof. Maximizing the VAE objective L(p, q) is equivalent to minimizing the Kullback-Leibler divergence between the joint distributions q and p, D KL (q(x) q(z|x) p(z) p(x|z)). The VAE optimum is attained when the KL divergence is zero i.e. when q(x, z) = p(x, z) almost everywhere. The independent solution clearly satisfies this.

To show that all other configurations are suboptimal, let p, q be such that q(z|x) = p(z) for some value of x with q(x) > 0. Then the implied marginal q(z) = x∈X q(z|x) q(x) is a finite Gaussian mixture with at least one nonstandard component, so it cannot equal p(z) and the KL divergence is positive.

Our theorem shows that not only is the VAE agnostic to representation learning, it in fact encourages the independent solution. The Gaussianity assumption introduces a gap due to q(z) not being able to fit p(z). We have observed this to be a major practical concern. The gap tends to grow with the dimension of the latent space as well as the amount of encoded information. We include 1D plots illustrating the issue in Appendix C.

In our experiments, we mitigate the issue by employing a resampling-based approximate posterior of Cremer et al. (2017) based on importance weighted autoencoders (Burda et al., 2015) .

7. EXPERIMENTS:

We show that the mutual autoencoder can learn latent codes ranging from independent to nearly deterministic. First, we consider two toy examples (a continuous and a discrete one), where we can visualize important quantities. Then we show promising results on text data, building a VAE model of movie reviews.

7.1. SPLITTING THE NORMAL:

We compare our method to the variational autoencoder on the task of modeling the one-dimensional standard normal distribution, q(x) = N (x| 0, 1). Our goal here is to a) show in a minimal setting the issue of representation collapse in VAEs, and b) show that our method overcomes this problem and in fact can encode any pre-specified amount of information.

We now describe the experimental setup. The prior p(z) is assumed to be a one-dimensional standard normal. Both the decoder p θ (x|z) and encoder q θ (z|x) are modeled as normal, with means and log-variances parametrized by a three-layer fully connected network. We train a VAE and several instances of the MAE with different values of the mutual information target M .

One can think of the task as splitting the normal q(x) into an infinite mixture of normals p(x|z) p(z)dz. There exists a trivial optimum of the VAE objective p θ (x|z) = q(z|x) = N (x| 0, 1) that ignores the latent code. Indeed, this is the solution recovered by the VAE and the MAE with M = 0, corresponding to the top row of Figure 2a .

When we set M > 0, though, our method learns a non-trivial representation, as can be seen in Figure 2a . As we increase M , the conditionals p θ (x|z) become more peaked and carry more information about z. 

7.2. CATEGORICAL EXAMPLE:

VAEs are particularly prone to ignoring the latent code when modeling discrete data. Here we show that the MAE is an effective solution even in this setting, and is able to learn a wide range of codes.

We consider synthetic data drawn from q(x) = Unif[10]. We let the latent variable again be a onedimensional standard normal. The decoder is a three-layer fully connected network with a 10-way softmax output; the encoder is given by 10 separate normal distributions, one for each possible value of x.

As in the continuous case, the decoder p θ (x|z) can easily reach an optimum of the VAE objective by outputting the marginal, Unif[10], irrespective of z. What makes the discrete case more difficult is that any distribution over X can be represented independently of the latent z. Such a scenario is shown in the top row of Figure 2b ; this is the model learned by the VAE and the MAE with M = 0.

For higher values of M , the MAE avoids this problem. In Figure 2b , we see that as M increases, the conditionals p θ (x|z) learned by MAE move from independence towards determinism.

7.3. MOVIE REVIEWS:

We now demonstrate the effectiveness of the mutual autoencoder on real text data.

We consider a sentence modeling task using a subset of the IMDB movie review dataset (Diao et al., 2014) . We split each review into sentences and extract those of length 8 words or shorter. We train several instances of the MAE with different M , using a variant of the bidirectional LSTM (Schuster & Paliwal, 1997; Hochreiter & Schmidhuber, 1997) as the encoder and a standard LSTM as the decoder. The conditioned value of the latent variable z is fed to the decoder LSTM at each step.

To evaluate the information content of a given model's latent code, we perform a simple reconstruction experiment: we use the model to encode a random subset of the training data and decode the obtained representations back into sentence space. A model with a highly informative latent code should be able to approximately reconstruct the input. We consider two metrics: 1) the number of exactly reconstructed sentences, and 2) the number of matching words between the input sentence and the reconstruction.

Figure 3 shows the results for MAE models trained with different M . The graph shows a simple monotonic relation: as M increases, so does the amount of encoded information, which in turn is reflected by the model's ability to autoencode. The MAE provides a powerful way of controlling this behavior by setting M .

In Table 1 we show sample sentence reconstructions of a VAE and two MAE models. As expected, the VAE fails to learn a useful latent code, q VAE θ (z|x) ≈ p(z), so its 'reconstructions' are random samples from the decoder model. At the other extreme, the MAE with a high value of the target mutual information (last row) learns a close-to-deterministic code and reconstructs the input sentence with high fidelity, at the expense of sample diversity. MAEs trained with intermediate values of M learn to 'paraphrase' input sentences.

To further inspect the learned representations, we interpolate between sentences in the latent space. An example is shown in Table 2 . We see that the model picks up on syntactic elements such as specific word choice or sentence length and that the generated sentences are mostly grammatical (subject to limitations of the training data).

Input:

there are many great scenes of course .

VAE:

and he knows it too . terrific performances from all three stars . this could have been a classic .

MAE (M = 5):

there are things that i liked . as a whole it works pretty well . there were a few good performances too .

MAE (M = 10):

there were many scenes like that . there are many great scenes of course . there were many scenes like that . there are many great scenes of course . there are other good people as well . they so look real its incredible . they like different , not good . they like different , not good . these people talk way too much . die hard gets a few character . i love these kind of movies . i left out a character . i love those films ) . 

8. RELATED WORK:

To our knowledge, the first work to notice the failures of VAE-type models when combined with high-capacity decoders was (Bowman et al., 2015) . Since then, a number of authors have reported similar difficulties and many fixes have been proposed.

Probably the most popular of these is 'KL cost annealing ' (Bowman et al., 2015) or 'warm-up' (Sønderby et al., 2016) , whereby the KL term in the VAE objective gets a multiplicative weight that gradually increases from zero to one during training. Another class of approaches is based on carefully limiting the capacity of the generative model class, as in (Chen et al., 2017) , (Yang et al., 2017) , or (Yeung et al., 2017) . Finally, some works make use of auxiliary objectives: Kingma et al. (2016) introduce a 'free bits' constraint on the minimum value of the KL term, whereas (Semeniuta et al., 2017) add a reconstruction term of an intermediate network layer. Zhao et al. (2017) propose a more drastic change to the objective, either replacing the KL term by one derived from a different divergence or removing it completely.

Our approach has close connections to the work of Chen et al. (2016) , who also use the variational infomax bound to learn meaningful representations. Their work is based on the GAN framework of Goodfellow et al. (2014) and hence may be difficult to apply to text. The work of Hu et al. (2017) can also be seen as employing (an approximation of) mutual information to enforce coupling between the latent and output variable, although they do not state it explicitly.

9. CONCLUSION:

In this work, we address the problem of representation collapse in VAEs with high-capacity decoders. We argue that the VAE objective is not sufficient to encourage learning of useful representations and propose to explicitly control the amount of information stored in the latent code. We present new techniques for solving the resulting constrained optimization problem. We show experimentally that our method is highly effective at encoding the user-specified number of bits, and can thus learn a continuum of representation spaces with varying properties.

Our method has the advantage of being principled, preserving the probabilistic interpretation of VAEs. Also, it is independent of specific architectural choices; it can be combined with any decoder distribution and benefits from its capacity. We find optimization to be easy. On the downside, the method is considerably slower to train due to the extra mutual information term.

In future work, we aim to extend the method to the semi-supervised setting. Making use of a supervisory signal, we would be able to control not only the amount, but also the kind of information stored in the latent code. This would be an important step towards fully controllable representation learning.

A REINFORCE DERIVATIONS FOR SECTION 5.3

A.1 DERIVATION AND PROPERTIES OF (13) Note that we treat p(z) as constant and ∇ p denotes the gradient with respect to p(x|z) only. 

The gradient in (26) can be estimated using Monte Carlo sampling of both the expectation and the Î(θ) term inside the sign. The sign estimate will inevitably be wrong sometimes, making our gradient estimator biased. However, note that only the norm is affected, while the direction of the gradient estimate is still correct in expectation.

:

See also the recent article(Huszár,

2017).

Figures and Tables:

Figure 1: The mutual autoencoder maximizes the likelihood log p θ (x) while constraining model family P by ensuring that the latent variable z and the observables x share a desired mutual information M = I pm . The purple lines denote level sets of the likelihood E x [log p θ (x)].

MAE splits the normal. Each row corresponds to a MAE model with the specified mutual information target M . Each curve corresponds to the scaled conditional p θ (x|z) p(z) for some fixed z ∈ {-2.1, -1.3, -0.4, +0.4, +1.3, +2.1}. (b) Categorical example, modeling the uniform distribution over 10 labels. Each row corresponds to a MAE model with the specified mutual information target M . Each colored patch corresponds to one possible outcome of x. A vertical slice at position z consists of 10 line segments whose lengths (sometimes 0) indicate the probabilities p θ (x|z).

Figure 2: The mutual autoencoder learns codes ranging from independent to deterministic.

Figure 3: Reconstruction ability grows with M .

p(x|z)) • log r * p (z|x) + p(x|z) • ∇ p log r * p (z|x)(15)Taking out p(x|z) and writing the resulting ratio as the derivative of the log gives= E z∼p x∈X p(x|z) (∇ p log p(x|z)) • log r * p (z|x) + ∇ p log r * p (z|x) ) (∇ p log p(z, x)) • log r * p (z|x) + ∇ p log r * p (z|x) (17) = E (x,z∼p) (∇ p log p(z, x)) • log r * p (z|x) + ∇ p log r * p (z|x) .(18)Next, we show that the term E (x,z∼p) ∇ p log r * p (z|x) vanishes when r * p (z|x) = p(z|x).) := H p (z) + E (x,z)∼p θ [log r ω (z|x)].(21)We will compute the gradient of C Î(θ) -M with respect to the model parameters θ.∇ θ C Î(θ) -M = C • sign Î(θ) -M • ∇ θ Î(θ) (22) = C • sign Î(θ) -M • ∇ θ E z∼p E x∼p θ (x|z) [log r ω (z|x)] (23) = C • sign Î(θ) -M • E z∼p x∈X (∇ θ p θ (x|z)) • log r ω (z|x)(24)Again making use of the derivative of the log, we get= C • sign Î(θ) -M • E z∼p x∈X p θ (x|z)(∇ θ log p θ (x|z)) • log r ω (z|x)(25) = C • sign Î(θ) -M • E (x,z)∼p θ [(∇ θ log p θ (x|z)) • log r ω (z|x)].

Sentence reconstructions. The input sentence is encoded and decoded by a given model, and the output is displayed. We show three reconstructions per model.

Table content: <table/>

Sentence interpolation, M = 10.

Table content: <table/>

Peer review (iclr2018):
 The authors propose a variational autoencoder constrained in such a way that the mutual information between the observed variables and their latent representation is constant and user specified. To do so, they leverage the penalty function method as a relaxation of the original problem, and a variational bound (infomax) to approximate the mutual information term in their objective. I really enjoyed reading the paper, the proposed approach is well motivated and clearly described. However, the experiments section is very weak. Although I like the illustrative toy problem, in that it clearly highlights how the method works, the experiment on real data is not very convincing. Further, the authors do not consider a more rigorous benchmark including additional datasets and state-of-the-art modelling approaches for text. - {.cal Z} in (1) not defined, same for .Theta.