Title: MIM: MUTUAL INFORMATION MACHINE

Abstract: We introduce the Mutual Information Machine (MIM), an autoencoder framework for learning joint distributions over observations and latent states. The model formulation reflects two key design principles: 1) symmetry, to encourage the encoder and decoder to learn different factorizations of the same underlying distribution; and 2) mutual information, to encourage the learning of useful representations for downstream tasks. The objective comprises the Jensen-Shannon divergence between the encoding and decoding joint distributions, plus a mutual information regularizer. We show that this can be bounded by a tractable cross-entropy loss between the true model and a parameterized approximation, and relate this to maximum likelihood estimation and variational autoencoders. Experiments show that MIM is capable of learning a latent representation with high mutual information, and good unsupervised clustering, while providing NLL comparable to VAE (with a sufficiently expressive architecture).

1. INTRODUCTION:

Mutual information is a natural indicator of the quality of a learned representation (Hjelm et al., 2019) , along with other characteristics, such as the compositionality of latent factors that are expected to be useful in downstream tasks, like transfer learning (Bengio et al., 2017) . Mutual information is, however, computationally difficult to estimate for continuous high-dimensional random variables. As such, it can be hard to optimize when learning latent variable models (Belghazi et al., 2018; Chen et al., 2016a) . This paper formulates a new class of probabilistic autoencoder model that is motivated by two key design principles, namely, the maximization of mutual information, and the symmetry of the encoder-decoder components. Symmetry captures our desire for both the encoder and decoder to effectively and consistently model the underlying observation and latent domains. This is particularly useful for downstream tasks in which either one or both of the encoder and decoder play a central role. These properties are formulated in terms of the symmetric Jensen-Shannon Divergence between the encoder and decoder, combined with an objective term to maximize mutual information. We refer to the resulting model as the mutual information machine, or MIM.

We contrast MIM with models trained using (approximate) maximum likelihood, the canonical example being the variational auto-encoder, or VAE (Kingma and Welling, 2013; Rezende et al., 2014) . The VAE comprises a probabilistic decoder and an approximate encoder, learned via optimization of an evidence-based lower bound (ELBO) on the log marginal data distribution. In contrast to MIM it is asymmetric in its formulation, and while often producing excellent representations, VAEs sometimes produce pathological results in which the encoder, or approximate posterior, conveys relatively little information between observations and latent states. This behavior, often referred to as posterior collapse,results in low mutual information between observations and inferred latent states (Bowman et al., 2015; Chen et al., 2016b; Razavi et al., 2019; van den Oord et al., 2016; 2017) .

We formulate the MIM model, and a learning algorithm that minimizes an upper bound on the desired loss. The resulting objective can be viewed as a symmetrized form of KL divergence, thereby closely related to the asymmetric KL objective of the VAE. This also enables direct comparisons to the VAE in terms of posterior collapse, mutual information, data log likelihood, and clustering. Experiments show that MIM offers favourable mutual information, better clustering in the latent representation, and similar reconstruction, at the expense of sampling quality and data log likelihood when compared to VAE with the same architecture. We also demonstrate that for a sufficiently powerful architecture, MIM can match sampling quality and log likelihood of a VAE with the same architecture.

2. VARIATIONAL AUTOENCODERS:

VAE learning entails optimization of a variational lower bound on the log-marginal likelihood of the data, log P(x), to estimate the parameters θ of an approximate posterior q θ (z|x) over latent states z (i.e., the encoder) and a corresponding decoder, p θ (x|z) (Kingma and Welling, 2013; Rezende et al., 2014) . A prior over the latent space, P(z), often assumed to be an isotropic Gaussian, serves as a prior for q θ (z|x) in the evidence-lower-bound (ELBO) on the marginal likelihood: log P(x) ≥ E z∼q θ (z|x) [ log p θ (x|z) ] -D KL (q θ (z|x) P(z)). Here, we use the notation P(x) and P(z) to emphasize that these priors are given, and that we can draw random samples from them, but not necessarily evaluate the log-likelihood of samples under them. In what follows we often refer to them as anchors to further emphasize their role.

With amortized posterior inference, we take expectation over the observation distribution, P(x), to obtain the VAE objective:

R VAE (θ) = E x∼P(x) E z∼q θ (z|x) [ log p θ (x|z) ] -D KL (q θ (z|x) P(z))

= E x∼P(x),z∼q θ (z|x) [ log p θ (x|z) + log P(z) -log q θ (z|x) ] ,

(1) Gradients of Eqn. (1) are estimated through MC sampling from q θ (z|x) with reparameterization, yielding unbiased low-variance gradient estimates (Kingma and Welling, 2013; Rezende et al., 2014) .

VAEs are normally thought of as maximizing a lower bound on the data log-likelihood, however it can also be expressed as minimizing the divergence between two joint distributions over x and z. To see this, we first subtract log P(x) from (1), which does not change the gradients of the objective with respect to θ. We then negate the result, as we will be performing minimization. This yields L VAE (θ) = -R VAE (θ) + E x∼P(x) [ log P(x) ] = D KL (q θ (z|x)P(x) p θ (x|z)P(z)) .

(2) The VAE optimization is therefore equivalent to minimizing the KL divergence between an encoding distribution q θ (z|x)P(x) and a decoding distribution p θ (x|z)P(z).

3. SYMMETRY AND MUTUAL INFORMATION:

Our goal is to find a consistent encoder-decoder pair, representing a joint distribution over the observation and latent domains, with high mutual information between observations and latent states. By consistent, we mean that the encoding and decoding distributions q θ (z|x) P(x) and p θ (x|z) P(z), define the same joint distribution. Effectively, we aim to estimate an undirected graphical model with two valid factorizations. We note that consistency is achievable in the VAE when the approximate posterior q θ (z|x) is capable of representing the posterior under the decoding distribution p θ (z|x). In the general case, however, consistency is not usually achieved.

In contrast to the asymmetric divergence between encoding and decoding distributions in the VAE objective (2), here we consider a symmetric measure, namely, the well-known Jensen-Shannon divergence (JSD),

EQUATION

where M S is an equally weighted mixture of the encoding and decoding distributions; i.e.,

EQUATION

In addition to encoder-decoder consistency, to learn useful latent representations we also want high mutual information between x and z. Indeed, the link between mutual information and representation learning has been explored in recent work (Hjelm et al., 2019; Belghazi et al., 2018; Chen et al., 2016a) . Here, to emphasize high mutual information, we add a particular regularizer of the form

R H (θ) = 1 2 H(p θ (x|z) P(z)) + H(q θ (z|x) P(x)) .

(5) This is the average of the joint entropy over x and z according to the encoding and decoding distributions. This is related to mutual information by the identity H(x, z) = H(x) + H(z) -I(x; z). That is, minimizing joint entropy encourages the minimization of the marginal entropy and maximization of the mutual information. In addition to encouraging high mutual information, one can show that this particular regularizer has a deep connection to JSD and the entropy of M S , i.e., JSD(θ) + R H (θ) = H(M S ) .

(6) The derivation for Eqn. ( 6) is given in Appendix A.1.

4. MUTUAL INFORMATION MACHINE:

Section 3 formulates a loss function (6) that reflects our desire for model symmetry and high mutual information. This objective is difficult to optimize directly since we do not know how to evaluate log P(x) in the general case (i.e., we do not have an exact closed-form expression for P(x)). As a consequence, we introduce parameterized approximate priors, q θ (x) and p θ (z), to derive tractable bounds on the penalized Jensen-Shannon divergence. This is similar in spirit to VAEs, which introduce a parameterized approximate posterior. These parameterized priors, together with the conditional encoder and decoder, q θ (z|x) and p θ (x|z), comprise a new pair of joint distributions, q θ (x, z) ≡ q θ (z|x) q θ (x), and p θ (x, z) ≡ p θ (x|z) p θ (z).

These new joint distributions allow us to formulate a new, tractable loss that bounds H(M S ); i.e.,

EQUATION

where M θ = 1 2 p θ (x, z) + q θ (x, z) , and H(M S , M θ ) denotes the cross-entropy between M S and M θ . We refer to L CE as the cross-entropy loss. It aims to match the model prior distributions to the anchors, while also minimizing H(M S ). A key advantage of this formulation is that the cross-entropy loss can be trained by Monte Carlo sampling from the anchor distributions with reparameterization (Kingma and Welling, 2013; Rezende et al., 2014) .

At this stage it might seem odd to introduce a parametric prior for P(z). Indeed, setting it directly is certainly an option. Nevertheless, in order to achieve consistency between p θ (x, z) and q θ (x, z) it can be advantageous to allow p θ (z) to vary. Essentially, we trade-off latent prior fidelity for increased model consistency. We provide more insights about this in Appendix E.3.

One issue with L CE is that, while it will try to enforce consistency between the model and the anchored distributions, i.e., p θ (x, z) ≈ p θ (x|z)P(z) and q θ (x, z) ≈ q θ (z|x)P(x), it will not directly try to achieve model consistency: p θ (x, z) ≈ q θ (x, z). To remedy this, we bound L CE using Jensen's inequality, i.e.,

EQUATION

Equation ( 8) gives us the loss function for the Mutual Information Machine (MIM). It is an average of cross entropy terms between the mixture distribution M S and the model encoding and decoding distributions respectively. To see that this encourages model consistency, it can be shown that L MIM is equivalent to L CE plus a non-negative model consistency regularizer; i.e.,

EQUATION

The non-negativity of

EQUATION

)

One can conclude from (10) that the regularizer R MIM is zero only when the two joint model distributions, q θ (x, z) and p θ (x, z), are identical under fair samples from the joint sample distribution M S (x, z). In practice we find that encouraging model consistency also helps stabilize learning.

To understand the MIM objective in greater depth, we find it helpful to express L MIM as a sum of fundamental terms that provide some intuition for its expected behavior. In particular, as derived in the Appendix A.3,

EQUATION

The first term in Eqn. (11) encourages high mutual information between observations and latent states.

The second shows that MIM directly encourages the model priors to match the anchor distributions. Indeed, the KL term between the data anchor and the model prior is the maximum likelihood objective.

The third term encourages consistency between the model distributions and the anchored distributions, in effect fitting the model decoder to samples drawn from the anchored encoder (cf. VAE), and, via symmetry, fitting the model encoder to samples drawn from the anchored decoder (both with reparameterization). As such, MIM can be seen as simultaneously training and distilling a model distribution over the data into a latent variable model. The idea of distilling density models has been used in other domains, e.g., for parallelizing auto-regressive models (Oord et al., 2017) .

In summary, the MIM loss provides an upper bound on the joint entropy of the observation and latent states under the mixture distribution M S : L MIM (θ) ≥ H M S (x) + H M S (z) -I M S (x; z). Through the MIM loss and the introduction of the parameterized model distribution M θ , we are pushing down on the entropy of the anchored mixture distribution M S , which is the sum of marginal entropies minus the mutual information. Minimizing the MIM bound yields consistency of the model encoder and decoder, and high mutual information under M S between observations and latent states.

5. EXPERIMENTS:

In what follows we examine MIM empirically, with the VAE as a baseline. We consider synthetic datasets and well-known image datasets, namely MNIST (LeCun et al., 1998) , Fashion-MNIST (Xiao et al., 2017a) and Omniglot (Lake et al., 2015) . All models were trained using Adam optimizer Kingma and Ba (2014) with a learning rate of 10 -3 , and a mini-batch size of 128. Following Alemi et al. (2017) , we anneal the loss to stabilize the optimization. To this end we linearly increase β from 0 to 1 in the following expression for a number of 'warm-up' epochs:

EQUATION

)

Training continues until the loss (i.e., with β = 1) on a held-out validation set has not improved for the same number of epochs as the warm-up steps (i.e., defined per experiment). We have found the number of epochs to convergence of MIM learning to be between 2 to 5 times greater than a VAE with the same architecture. A complete description of MIM learning is provided in Appendix C.

(Code is available from https://www.dropbox.com/s/idnls2layat77sj/MIM-master.zip?dl=1 ).

5.1. POSTERIOR COLLAPSE IN LOW DIMENSIONAL DATA:

Before turning to empirical results, it is useful to briefly revisit similarities and differences between MIM and the canonical VAE formulation. To that end, one can show from Eqn. (1) that the VAE loss can be expressed in a form that bears similarity to the MIM loss in Eqn. (8). In particular, following the derivation in Appendix D,

EQUATION

where M VAE S (x, z) = q θ (z|x) P(x), and the entropy in the last term is taken under the sample distribution M VAE S . Like the MIM loss, the first term in Eqn. ( 13) is the average of two cross-entropy terms, between a sample distribution and the encoding and decoding distributions. Unlike the MIM loss, these terms are asymmetric as samples are drawn only from the encoding distribution. Further, the last term in the VAE loss ( 13) is equivalent to

-H M VAE S (x) -H M VAE S (z) + I M VAE S (x; z),

and thus it will also encourage a reduction in the mutual information. We posit that this plays a significant role in the phenomena often referred to as posterior collapse, in which the variance of the variational posterior grows large and the latent embedding conveys relatively little information about the observations (e.g., see (Chen et al., 2016b) and others).

To empirically support the expression in Eqn. (13), we consider experiments with synthetic data comprising 2D observations x ∈ R 2 , with a 2D latent space, z ∈ R 2 . In 2D one can easily visualize the model and measure quantitative properties of interest (e.g., mutual information). Fig. 1 depicts data drawn from a Gaussian mixture model comprising five isotropic components with standard deviation 0.25 (top), and latents from an isotropic standard Normal (bottom). The encoder and decoder conditional distributions are Gaussian whose means and variances are regressed from the input using two fully connected layers and tanh activation. The parameterized data prior, q θ (x), is defined to be the marginal of the decoding distribution (see Eqn. ( 23) in Appendix C.1), and the model prior p θ (z) is defined to be the anchor, so the only model parameters are those of the encoder and decoder (see Appendix C for details). We can therefore learn models with the MIM and VAE objectives, but with the same architecture and parameterization. In this section we used a warm-up To quantify this behavior Fig. 2 shows mutual information, the average negative log-likelihood (NLL) of test points under the model q θ , the mean reconstruction error of test points, and 5-NN classification performance1 (predicting which of 5 GMM components each test point was drawn from), all as functions of the number of hidden units. Following Belghazi et al. (2018) , we estimate mutual information using the KSG estimator (Kraskov et al., 2004; Gao et al., 2016) , based on 5-NN neighborhoods. The auxiliary classification task provides a proxy for representation quality. Mutual information, reconstruction RMSE, and classification accuracy for test data under the MIM model are better than with VAE models, at the expense of poorer data log likelihoods for MIM. In this case MIM learning finds near-invertible mappings with a sufficiently powerful architecture.

VAE MIM VAE MIM VAE MIM (a) h ∈ R 5 (b) h ∈ R 20 (c) h ∈ R 500

To further investigate MIM learning in low dimensional data, we project 784D images from Fashion-MNIST (Xiao et al., 2017b ) onto a 20D linear subspace using PCA (capturing 78.5% of total variance).

The training and validation sets had 50,000 and 10,000 images respectively. Fig. 3 summarizes the results, with MIM producing high mutual information and classification accuracy, at all but very low We conclude that the VAE is prone to posterior collapse for a wide range of models' expressiveness and latent dimensionality, with latent embeddings exhibiting low mutual information. In contrast, MIM was empirically robust to posterior collapse. In this regard, we note that several papers have described ways to mitigate posterior collapse in VAE learning, e.g., by lower bounding, or annealing the KL divergence term in the VAE objective (Alemi et al., 2017; Razavi et al., 2019) , or by limiting the expressiveness of the decoder (e.g., Chen et al. (2016b) ). We posit that MIM does not suffer from this problem as a consequence of the objective design principles that encourage high mutual information between observations and the latent representation.

5.2. MIM LEARNING IN HIGH DIMENSIONAL IMAGE DATA:

We next consider MIM and VAE learning with image data (Fashion-MNIST, MNIST, Omniglot).

Unfortunately, with high dimensional data we cannot reliably compute mutual information (Belghazi et al., 2018) . Instead, for model assessment we focus on log-likelihood, reconstruction, and the quality of random samples. In doing so we also explore multiple architectures, including the top performing models from Tomczak and Welling (2017), namely, convHVAE (L = 2) and PixelHVAE (L = 2), with Standard (S) priors 2 , and VampPrior (VP) priors 3 . The VP pseudo-inputs are initialized with training data samples. All the experiments below use the same experimental setup as in Tomczak and Welling (2017), and the same latent dimensionality z ∈ R 80 . Here we also demonstrate that a powerful prior (e.g., PixelHVAE (VP)) allows MIM to learn models with competitive NLL performance.

2 p θ (z) = N (z; µ = 0, σ = I), a standard Normal distribution, where I is the identity matrix. Sampling from an auto-regressive decoder (e.g., PixelHVAE) is very slow. During training, we find that we can also learn effectively with a sampling distribution comprising just the encoding distribution, i.e., P(x) q θ (z|x), rather than the mixture of the encoder and decoder. We refer to this particular MIM variant as asymmetric MIM (or A-MIM). (For details, see Sec. C.3).

3 p θ (z) = 1 K K k=1 q θ (z|u k ),

We report quantitative results in Table 1a for test negative log-likelihood (NLL). One can see that VAE models result in better NLL, but with a rather small gap for the more expressive models (i.e., PixelHVAE (VP)). We also show qualitative results for the most expressive models (i.e., PixelHVAE). The poor NLL and hence poor sampling for MIM with a weak prior model can be explained by the tightly clustered latent representation (e.g., Fig. 1 ). A more expressive, learnable prior can capture such clusters more accurately, and as such, also produces good samples (e.g., VampPrior). In other words, while VAE opts for better NLL and sampling at the expense of lower mutual information, MIM provides higher mutual information at the expense of the NLL for a weak prior, and comparable NLL and sampling with more expressive priors.

5.3. CLUSTERING AND CLASSIFICATION:

Finally, following Hjelm et al. (2019) , we consider an auxiliary classification task as a further measure of the quality of the learned representations. We opted for K-NN classification, being a non-parametric method which relies on semantic clustering in latent space, without any additional training. Given representations learned above in Sec. 5.2, a simple 5-NN classifier5 was applied to test data to predict one of 10 classes for MNIST and Fashion-MNIST.

Table 1b shows that in all but one case, MIM yields more accurate classification results. We attribute the performance difference to higher mutual information of MIM representations. Figure 5 provides a qualitative visualization of the latent clustering, for which t-SNE (van der Maaten and Hinton, 2008)) was used to project the latent space down to 2D. One can see that MIM learning tends to cluster classes in the latent representation more tightly, while VAE clusters are more diffuse and overlapping, consistent with the results in Table 1b .

6. RELATED WORK:

Given the vast literature on generative models, here we only touch on the major bodies of work related to MIM, beyond the strong connection to VAEs.

As mentioned above, mutual information, together with disentanglement, is considered to be a cornerstone for useful representations (Belghazi et al., 2018; Hjelm et al., 2019) . Normalizing flows (Rezende and Mohamed, 2015; Dinh et al., 2014; 2016; Kingma and Dhariwal, 2018; Ho et al., 2019) directly maximizes mutual information by restricting the architecture to be invertible and tractable. This, however, requires the latent dimension to be the same as the dimension of the observations (i.e., no bottleneck). As a consequence, normalizing flows are not well suited to learning a concise representation of high dimensional data (e.g., images). Here, MIM often yields mappings that are approximately invertible, with high mutual information and low reconstruction errors. Bornschein et al. (2015) share some of the same design principles as MIM, i.e., symmetry and encode/decoder consistency. However, their formulation models the joint density in terms of the geometric mean between the encoder and decoder, for which one must compute an expensive partition function. Pu et al. (2017) focuses on minimizing symmetric KL, but must use an adversarial learning procedure, while MIM can be minimized directly.

GANs (Goodfellow et al., 2014) , which focus mainly on decoder properties, without a proper inference model, have been shown to minimize JSD between the data anchor P(x) and the model generative process q θ (x) (i.e., the marginal of the decoding distribution in MIM terms). In particular, prior work recognizes the importance of symmetry in learning generative models with reference to symmetric discriminators on x and z (Donahue et al., 2016; Dumoulin et al., 2017; Bang and Shim, 2018) . In contrast, here we target JSD between the joint encoding and decoding distributions, together with a regularizer to encourage high mutual information.

7. CONCLUSIONS:

We introduce a new representation learning framework, named the mutual information machine (MIM), that defines a generative model which directly targets high mutual information (i.e., between the observations and the latent representation), and symmetry (i.e., consistency of encoding and decoding factorizations of the joint distribution). We derive a variational bound that enables the maximizion of mutual information in the learned representation for high dimensional continuous data, without the need to directly compute it. We then provide a possible explanation for the phenomena of posterior collapse, and demonstrate that MIM does not suffer from it. Empirical comparisons to VAEs show that MIM learning leads to higher mutual information and better clustering (and classification) in the latent representation, given the same architecture and parametrization. In addition, we show that MIM can provide reconstruction error similar to a deterministic auto-encoder, when the dimensionality of the latent representation is equal to that of the observations. Such behaviour can potentially allow approximate invertibility when the dimensionality differs, with a stochastic mapping that is defined through consistency and high mutual information.

In future work, we intend to focus on utilizing the high mutual information mapping provided by MIM, by exploiting the clustered latent representation to further improve the resulting generative model.

A.1 REGULARIZED JSD AND ENTROPY:

First we develop the relation in Equation ( 6), between Jensen-Shannon divergence of the encoder and decoder, the average joint entropy of the encoder and decoder, and the joint entropy of the mixture distribution M S .

The Jensen-Shannon divergence with respect to the encoding distribution q θ (z|x) P(x) and the decoding distribution p θ (x|z) P(z) is defined as

JSD(θ) = 1 2 (D KL (p θ (x|z) P(z) M S ) + D KL (q θ (z|x) P(x) M S )) = 1 2 H(p θ (x|z) P(z), M S ) -H(p θ (x|z) P(z))

+ H(q θ (z|x) P(x), M S ) -H(q θ (z|x) P(x))

Where M S = 1 2 (p θ (x|z) P(z) + q θ (z|x) P(x)) is a mixture of the encoding and decoding distributions. Adding on the regularizer R H (θ) = 1 2 (H(p θ (x|z) P(z)) + H(q θ (z|x) P(x))) gives

JSD(θ) + R H (θ) = 1 2 (H(p θ (x|z) P(z), M S ) + H(q θ (z|x) P(x), M S )) = H(M S ) A.2 MIM CONSISTENCY

Here we discuss in greater detail how the learning algorithm encourages consistency between the encoder and decoder of a MIM model, beyond the fact that they are fit to the same sample distribution.

To this end we expand on several properties of the model and the optimization procedure.

A.2.1 MIM CONSISTENCY REGULARIZER:

In what follows we derive the form of the MIM consistency regularizer, R MIM (θ), given in Equation ( 9). Recall that we define M θ = 1 2 (p θ (x, z) + p θ (x, z)). We can show that L MIM is equivalent to L CE plus a regularizer by taking their difference.

R MIM (θ) = L MIM (θ) -L CE (θ) = 1 2 (H(M S , p θ (x, z)) + H(M S , p θ (x, z))) -H(M S , M θ ) = 1 2 (D KL (M S p θ (x, z)) + H(M S ) + D KL (M S p θ (x, z)) + H(M S )) -D KL (M S M θ ) -H(M S ) = 1 2 (D KL (M S p θ (x, z)) + D KL (M S p θ (x, z))) -D KL (M S M θ )

where R MIM (θ) is non-negative, and is zero only when the encoding and decoding distributions are consistent (i.e., they represent the same joint distribution). To prove that R MIM (θ) ≥ 0 we now construct Equation ( 9) in terms of expectation over a joint distribution, which yields

R MIM (θ) = 1 2 (H(M S , p θ (x, z)) + H(M S , p θ (x, z))) -H(M S , M θ ) = E x,z∼M S - 1 2 log p θ (x, z) - 1 2 log p θ (x, z) + log 1 2 (q θ (x, z) + p θ (x, z)) = E x,z∼M S -log q θ (x, z) • p θ (x, z) + log 1 2 (q θ (x, z) + p θ (x, z)) = E x,z∼M S -log q θ (x, z) • p θ (x, z) 1 2 (q θ (x, z) + p θ (x, z)) ≥ 0

where the inequality follows Jensen's inequality, and equality holds only when q θ (x, z) = p θ (x, z) (i.e., encoding and decoding distributions are consistent).

A.2.2 SELF-CORRECTING GRADIENT:

One important property of the optimization follows directly from the difference between the gradient of the upper bound L MIM and the gradient of the cross-entropy loss L CE . By moving the gradient operator into the expectation using reparametrization, one can express the gradient of L MIM (θ) in terms of the gradient of log M θ and the regularization term in Equation ( 9). That is, with some manipulation one obtains

EQUATION

which shows that for any data point where a gap q θ > p θ exists, the gradient applied to p θ grows with the gap, while placing correspondingly less weight on the gradient applied to q θ . The opposite is true when q θ < p θ . In both case this behaviour encourages consistency between the encoder and decoder. Empirically, we find that the encoder and decoder become reasonably consistent early in the optimization process.

A.2.3 NUMERICAL STABILITY:

Instead of optimizing an upper bound L MIM , one might consider a direct optimization of L CE . Earlier we discussed the importance of the consistency regularizer in L MIM . Here we motivate the use of L MIM from a numerical perspective point of view. In order to optimize L CE directly, one must convert log q θ and log p θ to q θ and p θ . Unfortunately, this is has the potential to produce numerical errors, especially with 32-bit floating-point precision on GPUs. While various tricks can reduce numerical instability, we find that using the upper bound eliminates the problem while providing the additional benefits outlined above.

A.2.4 TRACTABILITY:

As mentioned earlier, there may be several ways to combine the encoder and decoder into a single probabilistic model. One possibility we considered, as an alternative to M θ in Equation ( 7), is

EQUATION

where β = √ q θ p θ dx dz is the partition function, similar to Bornschein et al. (2015) . One could then define the objective to be the cross-entropy as above with a regularizer to encourage β to be close to 1, and hence to encourage consistency between the encoder and decoder. This, however, requires a good approximation to the partition function. Our choice of M θ avoids the need for a good value approximation by using reparameterization, which results in unbiased low-variance gradient, independent of the accuracy of the approximation of the value.

A.3 MIM LOSS DECOMPOSITION:

Here we show how to break down the L MIM into the set of intuitive components given in Equation (11). To this end, first note the definition of L MIM :

EQUATION

We will focus on the first half of Equation ( 16) for now,

EQUATION

It will be more clear to write out the first term of Equation ( 17 + H(q θ (z|x) P(x)) + H(q θ (z|x) P(x), p θ (x|z)

EQUATION

= 1 4 H(p θ (x|z) P(z), q θ (z|x) P(x)) + H(q θ (z|x) P(x), p θ (x|z) P(z))

+ 2R H (θ) -R H (θ) = 1 4 D KL (p θ (x|z) P(z) q θ (z|x) P(x)) + D KL (q θ (z|x) P(x) p θ (x|z) P(z)) + 4R H (θ) -R H (θ) = 1 4 (D KL (p θ (x|z) P(z) q θ (z|x) P(x)) + D KL (q θ (z|x) P(x) p θ (x|z) P(z))) = 1 2 SKL(θ)

From Equation ( 21), if we add the regularizer R H (θ) and combine terms, we get

1 2 SKL(θ) + R H (θ) = 1 2 (H(M S , q θ (z|x) P(x)) + H(M S , p θ (x|z) P(z)))

Interestingly, we can write this in terms of KL divergence,

1 2 SKL(θ) + R H (θ) = 1 2 (D KL (M S q θ (z|x) P(x)) + D KL (M S p θ (x|z) P(z))) + H(M S ) = 1 2 (D KL (M S q θ (z|x) P(x)) + D KL (M S p θ (x|z) P(z))) + JSD(θ) + R H (θ)

which gives the exact relation between JSD and SKL.

1 2 SKL(θ) = 1 2 (D KL (M S q θ (z|x) P(x)) + D KL (M S p θ (x|z) P(z))) + JSD(θ) = 1 2 (D KL (M S q θ (z|x) P(x)) + D KL (M S p θ (x|z) P(z))) + 1 2 (D KL (q θ (z|x) P(x) M S ) + D KL (p θ (x|z) P(z) M S ))

C LEARNING

Here we provide a detailed description of MIM learning, with algorithmic pseudo-code. In addition we offer practical considerations regarding the choice of priors' parameterization , and gradient estimation. The empirical upper bound objective is expressed in terms of two cross-entropy terms (cf. ( 8)). Given N fair samples, {x i , z i } N i=1 drawn from the anchored (sample) distribution, M S (x, z) in (4), the empirical loss is

EQUATION

where samples from M S comprises equal numbers of points from p θ (x|z) P(z) and q θ (z|x) P(x). Samples from the anchors, P(x) and , P(z), are treated as external observations; i.e., we assume we can sample from them but not necessarily evaluate the density of points under the anchor distributions.

Algorithm 1 specifies the corresponding training procedure. The algorithm makes no assumptions on the form of the parameterized distributions (e.g., discrete, or continuous). In practice, for gradientbased optimization, we would like an unbiased gradient estimator without the need to accurately approximate the full expectations per se (i.e., in the cross entropy terms). This is particularly important when dealing with high dimensional data (e.g., images), where it is computationally expensive to estimate the value of the expectation. We next discuss practical considerations for the continuous case and the discrete case.

Algorithm 1 MIM learning of parameters θ Require: Samples from anchors P(x), P(z) 1: while not converged do 2:

D dec ← {x i , z i ∼ p θ (x|z)P(z)} N/2 i=1

3:

D enc ← {x j , z j ∼ q θ (z|x)P(x)} N/2 j=1 4: D ← D dec D enc 5:

# See definition of LMIM in Eq. ( 22)

6:

L MIM (θ) ≈ LMIM (θ; D) 7:

# Minimize loss 8:

∆θ ∝ -∇ θ LMIM (θ; D) 9: end while

C.1 MIM PARAMETRIC PRIORS:

There are several effective ways to parameterize the priors. For the 1D experiments below we model p θ (z) using linear mixtures of isotropic Gaussians. With complex, high dimensional data one might also consider more powerful models (e.g., autoregressive, or flow-based priors). Unfortunately, the use of complex models typically increases the required computational resources, and the training and inference time. As an alternative we use for image data the vampprior Tomczak and Welling (2017), which models the latent prior as a mixture of posteriors, i.e., p θ (z) = K k=1 q θ (z|x = u k ) with learnable pseudo-inputs {u k } K k=1 . This is effective and allows one to avoid the need for additional parameters (see Tomczak and Welling (2017) for details on vampprior's effect over gradient estimation).

Another useful model with high dimensional data, following Bornschein et al. (2015) , is to define q θ (x) as the marginal of the decoding distribution; i.e.,

EQUATION

)

Like the vampprior, this entails no new parameters. It also helps to encourage consistency between the encoding and decoding distributions. In addition it enables direct empirical comparisons of VAE learning to MIM learning, because we can then use identical parameterizations and architectures for both. During learning, when q θ (x) is defined as the marginal (23), we evaluate log q θ (x) with a single sample and reparameterization. When z is drawn directly from the latent prior:

log q θ (x) = log E z∼p θ (z) [ p θ (x|z) ] ≈ log p θ (x|z) .

When z is drawn from the encoder, given a sample observation, we use importance sampling:

log q θ (x) = log E z∼q θ (z|x) p θ (x|z) p θ (z) q θ (z|x) ≈ log p θ (x|z) + log p θ (z) -log q θ (z|x)

Algorithm 2 provides algorithm details with the marginal prior.

C.2 GRADIENT ESTIMATION:

Optimization is performed through minibatch stochastic gradient descent. To ensure unbiased gradient estimates of LMIM we use the reparameterization trick Kingma and Welling (2013); Rezende et al.

(2014) when taking expectation with respect to continuous encoder and decoder distributions, q θ (z|x) and p θ (x|z). Reparameterization entails sampling an auxiliary variable ∼ p( ), with known p( ), followed by a deterministic mapping from sample variates to the target random variable, that is p θ (z) = g θ ( ) and q θ (z|x) = h θ ( , x) for prior and conditional distributions. In doing so we assume p( ) is independent of the parameters θ. It then follows that

∇ θ E z∼q θ (z|x) [f θ (z)] = ∇ θ E ∼p( ) [f θ (h θ ( , x))] = E ∼p( ) [∇ θ f θ (h θ ( , x))]

Algorithm 2 MIM learning with marginal q θ (x)

Require: Samples from anchors P(x), P(z)

Require: Define q θ (x) = E z∼p θ (z) [p θ (x|z)] 1: while not converged do 2:

# Sample encoding distribution 3:

D enc ← {x i , z i ∼ q θ (z|x)P(x)} N i=1 4:

# Compute objective, approximate log q θ (x) with 1 sample and importance sampling 5:

log q θ (x i ) ≈ log p θ (x i |z i ) + log p θ (z i ) -log q θ (z i |x i ) 6: LMIM (θ; D enc ) ← -1 N N i=1 (log p θ (x i |z i ) + log p θ (z i )) 7:

# Sample decoding distribution 8:

D dec ← {x i , z i ∼ p θ (x|z)P(z)} N i=1 9:

# Compute objective, approximate log q θ (x) with 1 sample 10:

log q θ (x i ) ≈ log p θ (x i |z i ) 11: LMIM (θ; D dec ) ← -1 2N N i=1 (log q θ (z i |x i ) + 2 log p θ (x i |z i ) + log p θ (z i ))

12:

# Minimize loss

13: ∆θ ∝ -∇ θ 1 2 LMIM (θ; D dec ) + LMIM (θ; D enc )

14: end while where f θ (z) is the loss function with parameters θ. It is common to let p( ) be standard normal, ∼ N (0, 1), and for z|x to be Gaussian with mean µ θ (x) and standard deviation σ θ (x), in which case z = σ θ (x) + µ θ (x), A more generic exact density model can be learned by mapping a known base distribution (e.g., Gaussian) to a target distribution with normalizing flows Dinh et al. (2014; 2016) ; Rezende and Mohamed (2015) .

In the case of discrete distributions, e.g., with discrete data, reparameterization is not readily applicable. There exist continuous relaxations that permit reparameterization (e.g., Maddison et al. (2016) ; Tucker et al. (2017) ), but current methods are rather involved in practice, and require adaptation of the objective function or the optimization process. Here we simply use the REINFORCE algorithm Sutton et al. (1999) for unbiased gradient estimates, as follows

∇ θ E z∼q θ (z|x) [f θ (z)] = E z∼q θ (z|x) [∇ θ f θ (z) + f θ (z)∇ θ log q θ (z|x)] .

A detailed derivation follows the use of the relation below,

∇ θ q θ (z|x) = q θ (z|x)∇ θ log q θ (z|x)

in order to provide unbiased gradient estimates as follows

∇ θ E z∼q θ (z|x) [f θ (z)] = ∇ θ f θ (z) q θ (z|x) dz = q θ (z|x) ∇ θ f θ (z) dz + f θ (z) ∇ θ q θ (z|x) dz = q θ (z|x) ∇ θ f θ (z) dz + f θ (z) q θ (z|x) ∇ θ log q θ (z|x) dz = E z∼q θ (z|x) [∇ θ f θ (z) + f θ (z)∇ θ log q θ (z|x)]

which facilitate the use of samples to approximate the integral.

C.3 TRAINING TIME:

Training times of MIM models are comparable to training times for VAEs with comparable architectures. One important difference concerns the time required for sampling from the decoder during training. This is particularly significant for models like auto-regressive decoders (e.g., Kingma et al. ( 2016)) for which sampling is very slow. In such cases, we find that we can also learn effectively with a sampling distribution that only includes samples from the encoding distribution, i.e., P(x) q θ (z|x), rather than the mixture. We refer to this particular MIM variant as asymmetric MIM (or A-MIM). We use it in Sec. 5.2 when working with the PixelHVAE architecture Kingma et al. (2016) .

D POSTERIOR COLLAPSE IN VAE:

Here we discuss a possible root cause for the observed phenomena of posterior collapse, and show that VAE learning can be viewed as an asymmetric MIM learning with a regularizer that encourages the appearance of the collapse. We further support that idea in the experiments in Section 5.1. As discussed earlier, VAE learning entails maximization of a variational lower bound (ELBO) on the log-marginal likelihood, or equivalently, given Equation ( 1), the VAE loss in terms of expectation over a joint distribution:

EQUATION

To connect the loss in Equation ( 24) to MIM, we first add the expectation of log P(x), and scale the loss by a factor of 1 2 , to obtain

EQUATION

where P(x) is the data distribution, which is assumed to be independent of model parameters θ and to exist almost everywhere (i.e., complementing P(z)). Importantly, because P(x) does not depend on θ, the gradients of Eqs. ( 24) and ( 25) are identical up to a multiple of 1 2 , so they share the same stationary points.

Combining IID samples from the data distribution, x i ∼ P(x), with samples from the corresponding variational posterior, z i ∼ q θ (z|x i ), we obtain a joint sampling distribution; i.e., 

M VAE S (x, z) = P(x) q θ (z|x)

The sum of the last three terms in Equation ( 26) is the negative joint entropy -H(z, x) under the sample distribution M VAE S .

Equations ( 1) and ( 26), the VAE objective and VAE as regularized cross entropy objective respectively, define equivalent optimization problems, under the assumption that P(x) and samples x ∼ P(x) do not depend on the parameters θ, and that the optimization is gradient-based. Formally, the VAE objectives (1) and ( 26) are equivalent up to a scalar multiple of 1 2 and an additive constant, namely,

H M VAE S (x).

Equation ( 26) is the average of two cross-entropy objectives ( i.e., between sample distribution M VAE S and the model decoding and encoding distributions, respectively), along with a joint entropy term (i.e., last three terms), which can be viewed as a regularizer that encourages a reduction in mutual information and increased entropy in z and x. We note that Equation ( 26) is similar to the MIM objective in Equation ( 8), but with a different sample distribution, where the priors are defined to be the anchors, and with an additional regularizer. In other words, Equation (26) suggests that VAE learning implicitly lowers mutual information. This runs contrary to the goal of learning useful latent representations, and we posit that it is an underlying root cause for posterior collapse, wherein the trained model show low mutual information which can be manifested as an encoder which matches the prior, and thus provides weak information about the latent state (e.g., see (Chen et al., 2016b ) and others). We point the reader to Section E.1 for empirical evidence for the use of a joint entropy as a mutual information regularizer.

E ADDITIONAL EXPERIMENTS:

Here we provide additional experiments that further explore the characteristics of MIM learning. Test performance for modified MIM (blue) and modified VAE (red) for the 2D GMM data with (cf. Fig. 7 ), all as functions of the number of hidden units (on x-axis). Each plot shows the mean and standard deviation of 10 experiments. Adding encoding entropy regularizer to VAE loss leads to high mutual information (i.e., prevent posterior collapse), low reconstruction error, and better classification accuracy. Subtracting sample entropy regularizer from MIM loss results in almost zero mutual information (severe collapse), which leads to poor reconstruction error and classification accuracy.

MIM-H VAE+H MIM-H VAE+H MIM-H (a) h ∈ R 5 (b) h ∈ R 20 (c) h ∈ R 500

E.1 ENTROPY AS MUTUAL INFORMATION REGULARIZER:

Here we examine the use of entropy as a mutual information regularizer. We repeat the experiment in Section 5.1 with added entropy regularizer. Figure 7 depicts the effects of an added H q θ (x, z) to the VAE loss, and a subtracted H M S (x, z) from MIM . The corresponding quantitative values are presented in Figure 8 . Adding the entropy regularizer leads to increased the mutual information, and subtracting it results in a strong posterior collapse, which in turn is reflected in the reconstruction quality. While such an experiment does not represent a valid probabilistic model, it supports our use of entropy as a regularizer (cf. Eq. ( 5)) for JSD in order to define a consistent model with high mutual information.

E.2 CONSISTENCY REGULARIZER IN L MIM:

Here we explore properties of models for 1D x and z, learned with L MIM and L CE , the difference being the model consistency regularizer R MIM (θ). All model priors and conditional likelihoods (q θ (x), q θ (z|x), p θ (z), p θ (x|z)) are parameterized as 10-component Gaussian mixture models, and Top row shows anchor P(x) (dashed), prior q θ (x) (red), and reconstruction distribution

x i ∼ P(x) → z i ∼ q θ (z|x i ) → x i ∼ p θ (x|z i ) (green).

Bottom row mirrors the top row, with anchor P(z) (dotted), prior p θ (z) (blue), and reconstruction distribution z i ∼ P(z) → x i ∼ p θ (x|z i ) → z i ∼ q θ (z|x) (yellow). In (c-d) the reconstruction is replaced with decoding from anchors z i ∼ P(z) → x i ∼ p θ (x|z i ) (green), and encoding x i ∼ P(x) → z i ∼ q θ (z|x) (yellow). In (e-f) the reconstruction is replaced with decoding from priors z i ∼ p θ (z) → x i ∼ p θ (x|z i ) (green), and encoding x i ∼ q θ (x) → z i ∼ q θ (z|x) (yellow). Top and bottom rows in Fig. 9 depict distributions in observations and latent spaces respectively. Dashed black curves are anchors, P(x) on top, and P(z) below (GMMs with up to 3 components). Learned model priors, q θ (x) and p θ (z), are depicted as red (top) and blue (bottom) curves.

Green histograms in Fig. 9 (a,b) depict reconstruction distributions, computed by passing fair samples from P(x) through the encoder to z and then back through the decoder to x. Similarly the yellow histograms shows samples from P(z) passed through the decoder and then back to the latent space. For both losses these reconstruction histograms match the anchor priors well. In contrast, only the priors that were learned with L CE loss approximates the anchor well, while the L MIM priors do not.

To better understand that, we consider two generative procedures: sampling from the anchors, and sampling from the priors.

Anchor consistency is depicted in Fig. 9(c, d ), where Green histograms are marginal distributions over x from the anchored decoder (i.e., samples from P(z)p θ (x|z) ). Yellow are marginals over z from the anchored encoders P(x)q θ (z|x). One can see that both losses results in similar quality of matching the corresponding opposite anchors.

Priors consistency is depicted in Fig. 9 (e,f), where Green histograms are marginal distributions over x from the model decoder p θ (z)p θ (x|z). Yellow depicts marginals over z from the model encoder q θ (x)q θ (z|x). Importantly, with L MIM the encoder and decoder are consistent; i.e., q θ (x) (red curve) matches the decoder marginal, while p θ (z) (blue) matches the encoder marginal. The model trained with L CE (i.e., without consistency prior) fails to learn a consistent encoder-decoder pair. We note that in practice, with expressive enough priors, L MIM will be a tight bound for L CE .

E.3 PARAMETERIZING THE PRIORS:

Here we explore the effect of parameterizing the latent and observed priors. A fundamental idea in MIM is the concept of a single model, M θ . As such, parameterizing a prior increases the global expressiveness of the model M θ . Fig. 10 depicts the utilization of the added expressiveness in order to increase the consistency between the encoding and decoding model distribution, in addition to the consistency of M θ with M S .

(a) q θ (x), p θ (z) (b) q(x), p θ (z) (c) q θ (x), p(z) (d) q(x), p(z)

Figure 10 : MIM prior expressiveness. In this experiment we explore the effect of learning a prior, where the priors q(x) and p(z) are normal Gaussian distributions. Top row shows anchor P(x) (dashed), prior q θ (x) (red), and decoding distribution

z i ∼ p θ (z) → x i ∼ p θ (x|z i ) (green).

Bottom row mirrors the top row, with anchor P(z) (dotted), prior p θ (z) (blue), and encoding distribution x i ∼ q θ (x) → z i ∼ q θ (z|x) (yellow). As can be seen, parameterizing priors affects all learned distributions, supporting the notion of optimization of a single model M θ . We point that (a) demonstrates the best consistency between the priors and corresponding generated samples, following the additional expressiveness.

E.4 EFFECT OF CONSISTENCY REGULARIZER ON OPTIMIZATION:

Here we explore whether a learned model with consistent encoding-decoding distributions (i.e., trained with L MIM ) also constitutes an optimal solution of a CE objective (i.e., trained with L CE ).

Results are depicted in Fig. 11 . In order to distinguish between the effects of the optimization from the consistency regularizer we initialize a MIM model by pre-training it with L CE loss followed by L MIM training in Fig. 11 (i), and vice verse in Fig. 11 (ii). (a-b,e-f) All trained models in Fig. 11 exhibit similarly good reconstruction (green matches dashed black). (c-d,g-h) However, only models that were trained with L MIM exhibit encoding-decoding consistency (green matches red, yellow matches blue). While it is clear that the optimization plays an important role (i.e., different initialization leads to different local optimum), it is also clear that encoding-decoding consistency is not necessarily an optimum of L CE , as depicted in a non-consistent model (h) which was initialized with a consistent model (g). Not surprisingly, without the consistency regularizer training with L CE results in better fit of priors to anchors (f) as it is utilizing the expressiveness of the parametric priors in matching the sample distribution. (i) and (ii) differ in initialization order. Odd rows: anchor P(x) (dashed), prior q θ (x) (red). Even rows: anchor P(z) (dotted), prior p θ (z) (blue). (a-b,e-f) Reconstruction x i ∼ P(x) → z i ∼ q θ (z|x i ) → x i ∼ p θ (x|z i ) (x i green, z i yellow). (c-d,g-h) Prior decoding z i ∼ p θ (z) → x i ∼ p θ (x|z i ) (green), and prior encoding x i ∼ q θ (x) → z i ∼ q θ (z|x i ) (yellow). See text for details.

F ADDITIONAL RESULTS:

Here we provide additional visualization of various MIM and VAE models which were not included in the main body of the paper.

F.1 RECONSTRUCTION AND SAMPLES FOR MIM AND A-MIM:

In what follows we show samples and reconstruction for MIM (i.e., with convHVAE architecture), and A-MIM (i.e., with PixelHVAE architecture). We demonstrate, again, that a powerful enough encoder allows for generation of samples which are comparable to VAE samples. 

:

We experimented with 1-NN,3-NN,5-NN,10-NN and found the results to be consistent.

Test data in the top row of Fig.4are binary, while reconstructions depict the probability of each pixel being 1, followingTomczak and Welling (2017)

We omitted results for k ∈ {1, 3, 10} as we find them similar.

Figures and Tables:

Figure 1: VAE and MIM models with 2D inputs, a 2D latent space, and 5, 20 and 500 hidden units. Top row: Black contours depict level sets of P(x); red dots are reconstructed test points. Bottom row: Green contours are one standard deviation ellipses of q θ (z|x) for test points. Dashed black circles depict one standard deviation of P(z). The VAE predictive variance remains high, regardless of model expressiveness, an indication of various degrees of posterior collapse, while MIM produces lower predictive variance and lower reconstruction errors, consistent with high mutual information.

Figure1depicts learned VAE and MIM models, with 5, 20 and 500 hidden units (from left to right) to control model expressiveness. Mutual information and root-mean-squared test reconstruction error are inset in the top row. For the weakest architecture (a), with 5 hidden units, VAE and MIM posterior variances are similar to the prior (MIM in one dimension), a sign of posterior collapse. As the number of hidden units increases (b,c), the VAE posterior variance remains large, preferring lower mutual information while matching the aggregated posterior to the prior. The MIM encoder produces tight posteriors, and yields higher mutual information and lower reconstruction errors.

Figure 3: MIM (blue) and VAE (red) for 20D Fashion-MNIST, with latent dimension between 2 and 20. Plots depict mean and standard deviation of 10 experiments.

Figure 4: MIM and VAE learning with the PixelHVAE (VP) architecture, applied to Fashion-MNIST, MNIST, and Omniglot (left to right). The top three rows (from top to bottom) are test data samples, VAE reconstruction, and A-MIM reconstruction. Bottom: random samples from VAE and A-MIM.

Fig. 4 depicts reconstruction 4 and sampling for Fashion-MNIST, MNIST, and Omniglot, for the top performing model (PixelHVAE (VP)), with MIM and VAE being comparable. See Appendix F for additional results.

Figure 5: MIM and VAE z embedding for Fashion MNIST (top) and MNIST (botom).

distribution in M S . With it one can then rewrite the objective in Equation (25) in terms of the cross-entropy between M VAE S and the parametric encoding and decoding distributions; i.e.,



Figure 7: Effects of entropy as a mutual information regularizer in 2D x and 2D z synthetic problem.VAE and MIM models with 2D inputs, a 2D latent space, and 5, 20 and 500 hidden units. Top row: Black contours depict level sets of P(x); red dots are reconstructed test points. Bottom row: Green contours are one standard deviation ellipses of q θ (z|x) for test points. Dashed black circles depict one standard deviation of P(z). Here we added H q θ (x, z) to VAE loss, and subtracted H M S (x, z) from MIM loss, in order to demonstrate the effect of entropy on mutual information. Posterior collapse in VAE is mitigated following the increased mutual information. MIM, on the other hand, demonstrates a severe posterior collapse as a result of the reduced mutual information (i.e., posterior matches prior over z almost perfectly). (see inset quantities).

Figure 8: Effects of entropy as a mutual information regularizer in 2D x and 2D z synthetic problem.Test performance for modified MIM (blue) and modified VAE (red) for the 2D GMM data with (cf. Fig.7), all as functions of the number of hidden units (on x-axis). Each plot shows the mean and standard deviation of 10 experiments. Adding encoding entropy regularizer to VAE loss leads to high mutual information (i.e., prevent posterior collapse), low reconstruction error, and better classification accuracy. Subtracting sample entropy regularizer from MIM loss results in almost zero mutual information (severe collapse), which leads to poor reconstruction error and classification accuracy.



Figure9: We explore the influence of consistency regularizer R θ . CE and MIM indicate the loss, L CE or L MIM (the regularized objective), respectively. Top row shows anchor P(x) (dashed), prior q θ (x) (red), and reconstruction distributionx i ∼ P(x) → z i ∼ q θ (z|x i ) → x i ∼ p θ (x|z i ) (green).Bottom row mirrors the top row, with anchor P(z) (dotted), prior p θ (z) (blue), and reconstruction distribution z i ∼ P(z) → x i ∼ p θ (x|z i ) → z i ∼ q θ (z|x) (yellow). In (c-d) the reconstruction is replaced with decoding from anchors z i ∼ P(z) → x i ∼ p θ (x|z i ) (green), and encoding x i ∼ P(x) → z i ∼ q θ (z|x) (yellow). In (e-f) the reconstruction is replaced with decoding from priors z i ∼ p θ (z) → x i ∼ p θ (x|z i ) (green), and encoding x i ∼ q θ (x) → z i ∼ q θ (z|x) (yellow). While both models offers similar reconstruction (a-b), and similar consistency w.r.t. the anchors (c-d), only MIM finds a consistent model (e-f). See text for details.

Figure9: We explore the influence of consistency regularizer R θ . CE and MIM indicate the loss, L CE or L MIM (the regularized objective), respectively. Top row shows anchor P(x) (dashed), prior q θ (x) (red), and reconstruction distributionx i ∼ P(x) → z i ∼ q θ (z|x i ) → x i ∼ p θ (x|z i ) (green).Bottom row mirrors the top row, with anchor P(z) (dotted), prior p θ (z) (blue), and reconstruction distribution z i ∼ P(z) → x i ∼ p θ (x|z i ) → z i ∼ q θ (z|x) (yellow). In (c-d) the reconstruction is replaced with decoding from anchors z i ∼ P(z) → x i ∼ p θ (x|z i ) (green), and encoding x i ∼ P(x) → z i ∼ q θ (z|x) (yellow). In (e-f) the reconstruction is replaced with decoding from priors z i ∼ p θ (z) → x i ∼ p θ (x|z i ) (green), and encoding x i ∼ q θ (x) → z i ∼ q θ (z|x) (yellow). While both models offers similar reconstruction (a-b), and similar consistency w.r.t. the anchors (c-d), only MIM finds a consistent model (e-f). See text for details.

Figure11: Effects of MIM consistency regularizer and optimization on encoding-decoding consistency. (i) and (ii) differ in initialization order. Odd rows: anchor P(x) (dashed), prior q θ (x) (red). Even rows: anchor P(z) (dotted), prior p θ (z) (blue). (a-b,e-f) Reconstructionx i ∼ P(x) → z i ∼ q θ (z|x i ) → x i ∼ p θ (x|z i ) (x i green, z i yellow). (c-d,g-h) Prior decoding z i ∼ p θ (z) → x i ∼ p θ (x|z i ) (green), and prior encoding x i ∼ q θ (x) → z i ∼ q θ (z|x i ) (yellow). See text for details.

Figure 12: MIM and VAE learning with PixelHVAE for Fashion MNIST. The top three rows (from top to bottom) are test data samples, VAE reconstruction, A-MIM reconstruction. Bottom: random samples from VAE and MIM.(c-d) We initialized all pseudo-inputs with training samples, and used the same random seed for both models. As a result the samples order is similar.

Figure 18: MIM and VAE z embedding for Fashion MNIST with convHVAE architecture.

Figure 21: A-MIM and VAE z embedding for MNIST with PixelHVAE architecture.

High dimensional image data. Quantitative results based on 10 trials per condition.

Table content: <table><tr><td/><td/><td colspan="2">convHVAE (S)</td><td/><td/><td colspan="2">convHVAE (VP)</td><td/></tr><tr><td>Dataset</td><td/><td>MIM</td><td colspan="2">VAE</td><td colspan="2">MIM</td><td>VAE</td><td/></tr><tr><td colspan="9">Fashion MNIST 272.14 ± 0.64 225.40 ± 0.05 227.61 ± 0.34 224.77 ± 0.04</td></tr><tr><td>MNIST</td><td colspan="4">126.85 ± 0.56 80.50 ± 0.05</td><td colspan="2">82.73 ± 0.08</td><td colspan="2">79.66 ± 0.06</td></tr><tr><td>Omniglot</td><td colspan="4">141.81 ± 0.32 97.94 ± 0.29</td><td colspan="4">104.10 ± 2.17 97.52 ± 0.16</td></tr><tr><td/><td/><td colspan="2">PixelHVAE (S)</td><td/><td/><td colspan="2">PixelHVAE (VP)</td><td/></tr><tr><td/><td/><td>A-MIM</td><td colspan="2">VAE</td><td colspan="2">A-MIM</td><td>VAE</td><td/></tr><tr><td/><td>MIM</td><td>VAE</td><td>MIM</td><td>VAE</td><td colspan="2">A-MIM VAE</td><td>A-MIM</td><td>VAE</td></tr><tr><td colspan="2">Fashion MNIST 0.83</td><td>0.76</td><td>0.81</td><td>0.78</td><td>0.71</td><td>0.76</td><td>0.79</td><td>0.77</td></tr><tr><td>MNIST</td><td>0.97</td><td>0.92</td><td>0.97</td><td>0.92</td><td>0.95</td><td>0.86</td><td>0.96</td><td>0.81</td></tr></table>

Peer review (iclr2020):
The authors present the Mutual Information Machine (MIM), which is essentially a latent variable model learned via Adversarially Learned Inference with regularizations that prefer learning representations with high mutual information. While the method is interesting, it does not seem to be well justified in representation learning, and the empirical results are not well-justified due to the use of weak baselines. Pros: The idea of formulating the L_MIM objective appears interesting, which uses the mixture of P_	heta and Q_	heta and minimizes an upper bound to H(x) + H(z) - I(x, z), so it is effectively minimizing a lower bound to I(x, z), under the proposed symmetric distribution. Cons: Symmetry argument: - Prior work on maximizing mutual information for latent variables focus on the MI as evaluated in the *encoder distribution*, whereas in this paper the argument is focused on symmetry, i.e. the MI estimated is based on mathcal{M}_s. - However, for practical purposes, it is unclear how this is superior than the *encoder distribution* paradigm. Eventually, we wish to use the z for some downstream tasks, which requires us to obtain z from some distribution (e.g. the encoder q(z|x)) and then run classification on this z. - Unless we sample z | x from the conditional distribution of mathcal{M}_s as proposed by the authors, it is unclear why we need this symmetry in representation learning; it seems slower to sample from the MIM distribution though (see motivation for A-MIM). - The authors did not demonstrate results that compare A-MIM and MIM on the same architecture. Failure to discuss certain related literature, such as Alemi et al. 2017 and Zhao et al. 2018. These are also *strongly related to VAEs* and related to mutual information objectives with latent variables. - Alemi et al. 2017 discussed this in the context of *rate discussion trade-off* - Zhao et al. 2018 proposed objectives to find specific trade-offs under constraints In fact, Zhao et al. 2018 considered the MIs in both P_theta and Q_theta, and one should be able to obtain a *symmetric* MI by adding these two Mis uniformly (although again, I am not sure how this is useful for downstream applications on the model). Using VAE as a baseline comparison. - It is well known that the VAE objective does not encourage maximizing mutual information, and various existing works have already been proposed to address this problem. Outperforming VAE in terms of MI is hence not very surprising. - Methods that encourage mutual information maximization under the latent variable generative modeling framework would be more suited for comparison. How does MIM compare with the methods introduced in Chen et al. 2016, Alemi et al. 2017 and Zhao et al. 2018? - In the context of ALI and BiGAN, one could also apply other approaches like MINE / InfoMAX to maximize MI. Again, there is no comparison with these approaches. - Alemi and Zhao both showed that there is a trade-off between consistency and mutual information under the same architecture (using their objectives). It seems that what MIM empirically does is also to sacrifice NLL for better MI, and not clear whether MIM gives comparable or better trade-offs. The paper would be more convincing if the authors empirically compare with other sensible baselines and justify why we should use MIM as opposed to A-MIM in representation learning. [relevant papers] Chen et al. 2016 Lossy VAEs Alemi et al. 2017 Fixing a broken ELBO Zhao et al. 2018 A Lagrangian Perspective to Latent Variable Generative Models