Title: THE VARIATIONAL INFOMAX AUTOENCODER

Abstract: The Variational AutoEncoder (VAE) can learn simultaneously a representation and a generative model, but often, in order to generate sharper samples, a useless representation is learnt. Starting from the observation that a meaningful representation is an informative one, in this work we analyse the VAE from an information theoretic perspective. We define the concept of capacity for the variational network, and we associate the two tasks of VAE: generation and representation quality, to two information properties of the network: information of the generative model and capacity of the network. Then we suggest that an optimal generative model is the one optimising the Capacity-Constrained InfoMax (CCIM), a theoretical objective learning the maximal informative generative model while maintaining bounded the network capacity. The theoretical assumptions are confirmed by the computational experiments, where between the different families of VAE, the one optimising a variational lower bound of the CCIM generates sharper samples and learns a clustered and robust representation. We call this the Variational InfoMax AutoEncoder (VIMAE).

1. INTRODUCTION:

A common assumption in machine learning is that any visible data x ∈ X is completely described by some generative factor o, living in a smaller hidden space O, i.e. x = g(o) with g a (possibly stochastic) generative function. The aim of unsupervised representation learning research is to find a representation z of the generative factor o living in a known space Z describing, as well as o, the visible data x. This is particularly relevant because the learnt small representation z is task agnostic and, in principle, can be used as input for networks performing different tasks, leading to faster and more robust learning (generalisation property), (Rifai et al., 2011) .

Many models f φ : X → Z trying to learn such representations have been proposed (Dinh et al., 2016; Hinton et al., 2006; Maddison et al., 2017; Radford et al., 2015) , but recently in order to solve this problem it was proposed to consider a dual problem: define a priori z and find a generator map g θ , such that for any z, g θ (z) is an element of X . In particular, two families of probabilistic generative models have become dominant: Variational AutoEncoder (VAE) (Kingma & Welling, 2013; Rezende et al., 2014) and Generative Adversarial Network (GAN) (Goodfellow et al., 2014) . The common idea of the two approaches is that a good generator p θ (x|z) is the one able to generate the data that is as close as possible to the visible one, i.e. that with respect a certain metric D, the distance between the marginal p θ (x) = E p(z) [p θ (x|z) ] and the visible distribution p D (x) is minimal.

In this manuscript we restrict our attention to the VAE model, since by its architecture, it is the only one where the learnt representation, possibly from different datasets, can be used as input for networks performing different tasks, (Achille et al., 2018; Ramapuram et al., 2017) . Although VAE, by its training robustness and general good generative performance is the most popular model for representation learning, in particular cases it suffers from the uninformative representation issue: the representation is entangled and the generative model tends to be independent of z, i.e. p θ (x|z) ≈ p θ (x). As highlighted in the next section such behaviour is intrinsic in the variational loss, the Evidence Lower BOund (ELBO), encouraging a less informative representation.

Following the direction suggested in (Alemi et al., 2017; Zhao et al., 2017) we describe the VAE from an information theoretic perspective. Such description lead us to two observations: the WAE (Tolstikhin et al., 2017) and InfoVAE (Zhao et al., 2017) models are actually maximising a lower bound of the mutual information associated to a generator p θ (x|z) belonging in a certain family P, the Variational InfoMax (VIM); and given that the capacity of the network is a function of the entropy of the prior p(z), the VIM is the variational expression of the Capacity-Constrained InfoMax (CCIM) . Following the analogies between the CCIM and the Information Bottleneck (IB), the theoretical principle associated with the ELBO, we deduce that the representation quality is associated to the capacity term.

The theoretical arguments are confirmed by the performed experiments where we observe that, differently from what was argued in previous works (Higgins et al., 2017; Burgess et al., 2018) , it is possible to train a model that is able to learn good (able to generalise) representations while maintaining optimal generative performance. The main contributions of the paper are summarised in the following points:

• derivation of a variational lower bound for the maximal mutual information of a generaive model belonging in a certain family, see equation 7;

• definition and bounds estimation for the network capacity for a variational autoencoder, see equation 9;

• association of the two main properties of VAE, generation quality and good representation, to two different information concepts, respectively Mutual Information and network capacity;

• proposal of a new learning principle for unsupervised models: the Capacity-Constrained InfoMax, see equation 10, that allows both to learn a good representation while maintaining optimal generative performance.

The work is divided as follows: in the second section we describe briefly the VAE and its variants; in the third and fourth sections we describe the variational infomax method and related work. We conclude the paper with the experimental results and the final observations.

2. BACKGROUND:

The aim of this section is to describe VAE, understand principal issues of the ELBO objective and describe the three most relevant approaches to overcome such issues.

2.1. NOTATION AND PRELIMINARY DEFINITIONS:

We use calligraphic letters (i.e. X ) for sets, capital letters (i.e. X) for random variables, and lower case letters (i.e. x) for their samples. With abuse of notation we denote both the probability and the corresponding density with the lower case letters (i.e. p(x)).

KL divergence Given two random distributions p(x) and q(x), the Kullback-Leibler (KL) divergence

EQUATION

is an (intuitive) measure of the distance between the distributions p and q.

Mutual Information and Capacity Given a channel Z → X with X and Z random variables, jointly distributed according to p(x, z) and with marginals p(x) and p(z). The mutual information

I(X, Z) = D KL (p(x, z)||p(x)p(z)),

is a measure of the reduction of uncertainty in X due to the knowledge of Z, and the capacity

C(X, Z) = sup p(z)∈P I(X, Z)

is the maximal information that can be shared for a fixed generator p(x|z).

2.2. VARIATIONAL AUTOENCODER:

From now on let us assume that the unknown distribution of the data p(x) coincides with the empirical one p D (x), and that the distribution of the latent representation p(z) is known. In this context the VAE is a model solving the following optimisation problem: find the generative model p θ (x, z) ∈ P θ , specified by the parameters θ of the associated neural network, maximising the ELBO objective

EQUATION

a lower bound of the unfeasible-to-compute marginal likelihood E p(x) [log p θ (x)]. The ELBO objective is optimised by a regularised autoencoder, with encoder and decoder parametetrising, respectively, the inference and generative distributions, q φ (z|x) and p θ (x|z), with φ ∈ Φ, θ ∈ Θ and regulariser defined by the rate term D KL (q φ (z|x)||p(z)), measuring the excess number of bits required to encode samples from the encoder using the optimal code designed for p(z).

2.3. UNINFORMATIVE REPRESENTATION ISSUE:

As underlined in the introduction, the main issue of VAE is that the representations are not really informative of the input data and in the worst case, it is learned a Z-independent generative model p θ (x|z) = p θ (x). Such issues are intrinsic in the ELBO objective, equation 2, that can reach the optimum when D KL (q φ (z|x)||p(z)) = 0 (Zhao et al., 2017) . The latter case means that the representation is completely uninformative, indeed the rate term, which can be rewritten as

D KL (q φ (z|x)||p(z)) = I q (X, Z) + D KL (q φ (z)||p(z)),

is a penalty on the encoding information, and is zero when I q (X, Z) = 0, with q φ (z|x) = q φ (z) = p(z), i.e. when q φ does not encode any information about the input x.

We now describe the three most relevant models that try to overcome the uninformative representation issue.

InfoVAE In (Zhao et al., 2017) the InfoVAE family of models was proposed, a generalisation of the VAE model optimising the objective

-αI q (X, Z) -λD KL (q φ (z)||p(z)) + E p(x) [E q(z|x) [log p(x|z)]],

with α and λ two real positive hyper-parameters.

The main advantage of this definition is that it is possible to consider separately the two components of the rate term. In particular, in (Zhao et al., 2017) it was observed that by eliminating the information penalty (α = 0), the generative performance of the model improves and the representation results are more informative.

β-VAE In (Higgins et al., 2017) , starting from the observation that the optimal case is rare, but most of the learned features by VAE are not disentangled, it is proposed an opposite approach: put a high penalty to the rate term, in order to constrain the model to learn the most informative property of the data, and then have a disentangled representation of the data. The β-VAE family is a particular case of InfoVAE where α = λ 1. This idea, that at first sight looks counter-intuitive, is based on the observation that by the additive property of the KL-divergence

EQUATION

pushing the penalty associated with the rate is equivalent to penalising the informativeness of most features, leaving few features containing the relevant information. Starting from a bits-back coding argument, a similar conclusion was derived in (Chen et al., 2016) .

Minimal rate bound In order to avoid the null rate issue, and maintain semantic informative representation, in (Alemi et al., 2017) was suggested to optimise the following variant of ELBO:

EQUATION

where R 0 > 0, defines both the lower bound of the rate and the upper bound of the encoding information I q (X, Z). If, as we will see below, the objective in equation 4 is a generalisation of both InfoVAE (R 0 = max θ I θ (X, Z)) and β-VAE (R 0 = 0), the correct R 0 is unknown and its optimisitation is not an easy task.

3 THE MODEL

3.1. THE VARIATIONAL INFOMAX:

Assuming the distribution associated to the two random variables p(x) and p(z) is known, the InfoMax objective is defined as: find the joint distribution p θ (x, z)

∈ P θ := {p θ (x, z) : E p(z) [p θ (x|z)] = p(x), E p(x) [p θ (z|x)] = p(z)} maximising the mutual information I θ (X, Z) = D KL (p θ (x, z)||p(x)p(z)), i.e. find θ * ∈ Θ s.t. I θ * ≥ I θ for any θ ∈ Θ.

Since the definition via KL divergence is computationally intractable, it is necessary to re-write the mutual information as

I θ (X, Z) = h θ (X) -h θ (X|Z), (5) where h θ (X) = -E p θ (x) [log p θ (x)] is the entropy of X, and h θ (X|Z) = -E p θ (x,z) [log p θ (x|z)] is the conditional entropy h θ (X|Z). Since p θ (x, z) ∈ P θ the entropy h θ (X) = h(X)

is constant, and in order to maximise the mutual information it is sufficient to minimise the conditional entropy.

Excluding some special cases (Bell & Sejnowski, 1997) , minimising the conditional entropy is unfeasible, so it is necessary to consider an associated variational problem: for any q φ (z|x) such that q φ (z) = p(z) and φ ∈ Φ = Θ, learn the generative model p θ (x|z) minimising the reconstruction accuracy term

E p(x) [E q φ (z|x) [log(p θ (x|z))]]

. Indeed, the following variational objective:

EQUATION

is a lower bound of I θ * (X, Z) and is maximal when q(z|x) = p θ (z|x) = p θ * (z|x), (see the Appendix).

Unfortunately, the formulation in equation 6 is still unfeasible to compute, because it requires that q φ (z) = p(z), but by the butterfly architecture of the autoencoder, q φ (z) tends to be uniformly distributed on the space Z. For this reason, the model is trained maximising the following relaxed form:

EQUATION

where it is introduced a term D(q φ (z)||p(z)) encouraging the empirical distribution q φ (z) to be close, according to the metric D, to p(z). In the following we assume D = D KL , and in order to avoid any confusion the variational autoencoder trained maximising equation 7 will be dubbed VIMAE.

Encoding channel In VAE we observed that an uninformative representation was caused by the non-informativeness of the encoding map q φ (z|x). Since from equation 7 it is not clear how q φ (z|x) behaves, we consider an equivalent representation, (see the Appendix):

EQUATION

From equation 8 we see that the infomax objective, equation 7, can be read as a composition of three sub-objectives: find a generative model p θ (x|z), with marginal resembling the visible distribution p(x) (first term); maximise the (unbounded) variational mutual information (third term); and learn an inferred distribution q φ (z, x) close to the generative model p θ (x, z). Then the optimum is obtained by q φ (x, z) = p θ (x, z) such that I θ (X, Z) is maximal, confirming the validity of the approximation made above.

3.2. CHANNEL CAPACITY:

In a channel with variational mutual information I θ,φ as defined in equation 6 the (variational) capacity C θ,φ (X, Z), is defined as

EQUATION

)

If P is the space of all distributions on Z, the capacity of the network coincides with the variational mutual information of the model trained minimising only the reconstruction loss. In the latter case, as observed above, it is not guaranteed to learn the generator with p θ (z) = p(z). This is because, given two equally informative generative models p θ (x|z) and p θ (x|z), with encoder respectively q φ (z|x) and q φ (z|x), with h(q φ (z)) < h(q φ (z)), then I θ,φ < I θ ,φ . From such observation we deduce that bounding the entropy of the representation Z is the way to bound the capacity without penalising the encoding information, and that the penalty introduced in equation 7 is actually a bound to the capacity itself. So, the VIM objective can be defined as a variational approximation of the Capacity-Constrained InfoMax:

EQUATION

which given a set of equally informative generators, learns the one having the minimal capacity. The idea of the CCIM is similar to the idea of the IB, from which was derived the β-VAE: constrain the capacity of the network in order to learn only the relevant features of the input data. The difference with the IB, max

EQUATION

lies in the second term: the network capacity instead of the encoding information. This choice allows the network to learn a good representation (small capacity) while maintaining good generative performance (high mutual information). Indeed, as shown in equation 8, the generative performance is associated to the informativeness of q φ (z|x) and p θ (x|z).

In order to test the assumption that it is sufficient to bound the entropy of Z, instead of the encoding mutual information, to learn a good representation, in the experiments (see below) we consider the cases Z is Normal (VIMAE-n) or Logistic (VIMAE-l) distributed. We choose to compare the popular Normal distribution with the Logistic one for two reasons: the Logistic has less entropy than a Gaussian distribution and because it is a common assumption in natural science to suppose that the hidden factors of the visible data are logistically distributed (Hyvärinen et al., 2009) .

Entropy vs Rate:

According to what written above the choice of the distribution p(z) is a way to bound the maximal information I θ of the generative model as well as R 0 in equation 4, indeed, if we choose R 0 = C(X, Z) the theoretical solutions of equation 4 and equation 7 coincide. Although the two approaches are similar, both defines a target I θ for the generator p θ (x|z), they differ in one fundamental aspect: the ratio between the information I θ and network capacity C. For any given choice of the prior p(z) the generative model maximising VIM is the one with I θ = C, instead for any R 0 < C the information I θ associated to the generative model learned, by equation 4, coincides with R 0 , i.e. I θ = R 0 < C. That implies an high sensitivity to the hyper-parameter R 0 for the model in equation 4. In fact, as suggested in (Mathieu et al., 2018) , a ratio I θ /C ≈ 1, guarantees an appropriate level of overlap in the latent space between latent encoding, and then a better representation quality.

4. RELATED WORK:

Autoencoder literature Autoencoder models are one of the most used family of neural networks to extract features in an unsupervised way (Bengio et al., 2013) , and their relationship with Information Theory is well-established from the first unregularised autoencoders (Baldi & Hornik, 1989) . The classical unregularised autoencoders, minimising the reconstruction loss

E p(x) [E q φ (z|x) [-log p θ (x|z)]]

, are maximising an unbounded information, i.e. they are looking for a solution in the space Pθ = {p θ : p θ (x) = p(x)}. In general, a solution in this wide space is good only for reconstruction performance because Z contains all the possible information that can be stored in the space Z, and is not robust to input noise (Vincent et al., 2008) ; but, as observed in (Grover & Ermon, 2018) , if q φ (z|x) ∼ N (µ(x), σ(x)) the model p θ (x|z) is robust to noise and is a Gaussian generator. In this context the uninformative issue is avoided, but the price to pay is the impossibility to sample directly from a prior p(z) that is not defined; indeed, the model described in (Grover & Ermon, 2018) requires running relatively expensive Markov Chain to obtain samples.

Many regularised models have been proposed, but the most well known is VAE, that minimises the expected code length of communicating x. As we observed in the previous sections, it is not guaranteed that the method finds a useful representation. Such issue can be solved both controlling the information of the model, or considering a more flexible prior p(z), (Rezende & Mohamed, 2015; Kingma et al., 2016; Dinh et al., 2016) . In this manuscript we do not consider the latter approach, because by the model structure, it requires quite expensive computation (autoregressive models are necessary) and because the learned representation is not suitable for task different than the generation such as lifelong context, for which the standard VAE is used with success (Achille et al., 2018) .

The objective in equation 7 was firstly derived in (Tolstikhin et al., 2017) and (Zhao et al., 2017) .

Particularly relevant is the derivation in (Tolstikhin et al., 2017) because it allows us to describe an informative model p θ (x, z) as the one minimising the transport cost between the original and generated data.

Finally, we underline that in case we wish to consider a Jensen-Shannon divergence in equation 7 it is necessary to consider an adversarial network model, discriminating the true samples z ∼ p(z) from the fake sampled by q φ (z) (Goodfellow et al., 2014) . In the latter case the obtained model is equivalent to the Adversarial AutoEncoder (Makhzani et al., 2015) . We conclude by remarking that in all the cases cited above the Infomax objective was never maximised using a prior p(z) different from a Gaussian.

Information theoretic literature Information theory is strongly related with neural networks, and not only with autoencoders. Originally the InfoMax objective was applied to a self-organised system with a single hidden layer, (Bell & Sejnowski, 1997; Linsker, 1989) where the bound in the capacity was given by the numbers of hidden neurons. More recently, the (naive) InfoMax has given way to a new information-theoretic principle: the Information-Bottleneck (Tishby et al., 2000) . The idea of this principle is that a feed-forward neural network trained for task T tends to learn a minimal sufficient representation of the data, maximising the following objective:

EQUATION

Although it was shown that in the general case this principle does not hold true (Saxe et al., 2018), the principle was used as a regularisation technique with success both in unsupervised (Alemi et al., 2017; Higgins et al., 2017) and supervised (Alemi et al., 2016) settings. We observe that the VIM, equation 7, and IB, equation 12, coincide in the case of a deterministic encoder, where the encoding information is the entropy of Z.

5. EXPERIMENTS:

Here we empirically evaluate the VIMAE. The section is divided into three parts: in the first part we compare the ability of the described models to infer the representation, z ∼ N (0, I). Such an experiment is to evaluate the entropy of Z ∼ q(Z) and then, as observed in section 3.2, an indirect way to estimate capacity of the network (C(X, Z) ∝ h(Z)). In the second part we evaluate the reconstruction and generative performance of the models. Indeed, the combination of the two tasks is estimation of the mutual information of the generative model I θ (X, Z), see equation 8.

In the third part we evaluate the robustness to noise and generalisation property of the learnt representation, observing that an informative model with small capacity is the best one for these tasks.

In all the described experiments, the divergence D KL (q(z)||p(z)) in equation 7 is approximated via the Maximum Mean Discrepancy (Zhao et al., 2017 ) defined as:

EQUATION

where H k is the Reproducing Kernel Hilbert Space associated to a positive definite kernel k(•, •) : Z × Z → R + . Moreover, by difficulties to compute the objective equation 4, as suggested in (Alemi et al., 2017) we decided to optimise a β-VAE, denoted β A -VAE, with β < 1; in order to avoid any confusion, the original version proposed in (Higgins et al., 2017) with β 1 will be renamed β H -VAE. From the 2d hidden distribution learnt by the different methods in figure 1 , we observe that the ELBO-based models are not able to learn a distribution q(z) that is close to the prior p(z); and, as we will see in the following experiments, β H -VAE, the only ELBO-model having comparing results to the VIMAE, has to penalise drastically the decoding information.

Thanks to this experiment we see that, in general, the ELBO-based models are not seeking for a small capacity network, since h(q(z)) h(p(z)), and that the divergence penalty introduced in equation 7 is a bound for the entropy h(q(z)), and then of the channel capacity C.

5.2. THE MODEL INFORMATION:

The experiments in these final sections were performed with the same settings and autoencoder models used in (Tolstikhin et al., 2017) , an architecture similar to the DCGAN (Radford et al., 2015) with batch normalization (Ioffe & Szegedy, 2015) (more details given in the Appendix). We consider four data-sets: MNIST and CIFAR10, two standard data-sets with ground-truth labels; Omniglot, a data-set of 1623 characters from 50 alphabets, 30 training and 20 evaluation, where each character appears 80 times, to evaluate the informativeness of the model and the quality of the learned representation; in the Appendix, we also consider CelebA (Liu et al., 2015) , consisting of roughly of 203k faces of 64 × 64 resolution, in order to compare the generative quality of the pictures. After considering many parameters for β H , β A and λ, we choose, in accordance with what was suggested in (Tolstikhin et al., 2017) , β H = λ = 10, and β A = 0.2 for MNIST and Omniglot and β = λ = 100 and β A = 0.4 for CelebA and CIFAR10 experiments.

The goal of this section is to evaluate the informativeness of the learnt generative model p θ (x|z). From what was described above, the reconstruction loss or the generation quality alone are not reliable metrics, because the reconstruction loss is an estimation of the variational mutual information I θ,φ (X, Z), and the generation quality is an estimation of D KL (p θ (x)||p(x)) that, as observed in section 2.3, it is possible to minimise with an uninformative generator. But, according to equation 8 the combination of the two task performances is a good empirical estimation of I θ (X, Z); indeed, by generative experiments, we require that q(z) = p(z), so the KL divergence term in equation 8 comes for free.

Reconstruction and generative performances:

According to what was asserted above, we observe from figure 2 (top) and figure 3 that an high penalty on the rate term (small encoding information) in the ELBO based model coincides with poor reconstruction. Indeed, the β H -VAE, is the one with worst reconstruction, and instead β A -VAE is the model that behaves similarly to VIMAE, that is optimising an encoding information-free objective. Such behaviour is consistent with the rate term associated to each model, see caption of each figure, where we observe that VIMAE, have the highest rate and in particular this is twice the one associated to β A -VAE, confirming that VIMAE learns a maximal informative decoder, and that β A -VAE is only theoretically analogous to VIMAE.

The models that we are considering are defined as generative models, so giving a sample z ∼ p(z) they should be able to generate a new data x similar to the original one. As we observe from the (Liu et al., 2015) , see the Appendix, where it is observed that if the generative difference between the two VIMAEs is small, the difference with the VAE counterparts is high. Such behaviour is in agreement with what was observed until now: the ELBO based model does not learn a good generative network, and the good reconstruction is simply associated to a large entropy of Z, (encoding information) instead of an informative generative model p θ (x|z).

Table 1 : NLL for generated samples on Omniglot (smaler is better)

Method NLL VAE 1224 β H -VAE 1254 β A -VAE 1228 VIMAE-n 1190 VIMAE-l 1223

Similar experiments were conducted with the MNIST dataset and are discussed in the appendix.

5.3. GENERALISATION PROPERTY:

We defined a good representation as the one containing the relevant properties of the visible data and able to generalise from the task for which was trained. In order to evaluate such quality, following the approach proposed in (Rifai et al., 2011) , we evaluate the accuracy of an SVM directly trained on the learned features of the data. Proceeding as in (Zhao et al., 2017) , we train the M1+TSVM (Kingma et al., 2014) and use the semi-supervised performance over 1000 (100 for Omniglot) samples as an approximate metric to verify the relevance and the quality of the learned representation. In order to evaluate the robustness of the learned features, we performed the same algorithm on the representation associated to corrupted data, i.e. z ∼ q(z|x + ν), considering two types of noise: Gaussian and mask. In the Gaussian case, we add to each pixel a ν value sampled from N (0, σ 2 ) with σ ∈ {0.2, 0.3, 0.4}, and in the masking case a fraction ν of the elements is forced to be 0: each pixel is masked according to a Bernoulli distribution B(p), p ∈ {0.2, 0.5}. Higher classification performance suggests that the learned representation contains the relevant information and, in case of corrupted input data, that it is robust. In the Omniglot case by the challenge of the task (the test alphabet was never seen in the training) we consider a 5-character data-set, split into 300 (60 × 5) for training and 100 for evaluation.

From the classification scores listed in tables 2-4, we see that the ELBO-based model learnt good representations for clean data, but not when corrupted data is given as input. This is particularly clear in the Bernoulli case, that is a noise different from the one seen in the training. Particularly relevant is the behaviour of the two VIMAEs: they are comparable in the cases of clean data and small noise, but the one with big capacity, VIMAE-n, suffers in large noise, while the one with small capacity, VIMAE-l, is the most robust and in some challenging cases, see table 2, the noise helps to improve the model accuracy. Such a result is consistent with the idea that a small capacity network is learning the relevant factors of the input data, that are the only ones robust to the input noise.

6. CONCLUSION:

We observe, via an information theoretic description of VAE, that it is possible to learn a good generative model while maintaining a meaningful hidden representation, and that goal can be reached by optimising the CCIM, an objective that separates out the two properties of a network: the generative information and its capacity. We compare both theoretically and computationally the CCIM objective with the IB and its variation (Alemi et al., 2017) , observing that the CCIM optimal solution, in particular cases, coincides with the IB one, but that in the practical case the difference is substantial.

Future work includes the extension of the VIMAE model to a lifelong task, a scenario where it is necessary to learn both an informative representation while maintaining bounded the capacity of the network, in order to avoid the catastrophic forgetting issue (Achille et al., 2018) .

RELATIONSHIP BETWEEN ENCODING, DECODING AND VARIATIONAL INFORMATION:

Defined q(z, x) := q φ (z|x)p(x) the encoding distribution and p θ (x, z) := p θ (x|z)p(z) the decoding one. The encoding I q (X, Z), decoding I θ (X, Z) and the variational I θ,φ (X, Z) information are defined respectively:

I q (X, Z) = D KL (q(z, x)||q(z)p(x)) = h(X) -h q (X|Z) I θ (X, Z) = D KL (p θ (z, x)||p(z)p(x)) = h(X) -h θ (X|Z) I θ,φ (X, Z) = h(X) -E q(z,x) [-log p θ (x|z)]

Assuming θ * ∈ Θ is the parameter associated to the maximal decoding information, I θ (X, Z) ≤ I θ * (X, Z) for any θ ∈ Θ, it follows that for any q φ (x, z) ∈ P θ , i.e. for any q φ (z) = E p(x) [q φ (x|z)] = p(z) and φ ∈ Φ ⊂ Θ,

I θ * (X, Z) ≥ I q (X, Z).

Then a lower bound of I q is a lower bound of I θ * . By property of KL-divergence we have that for any p θ (x|z) the following relationship holds:

EQUATION

From equation 14 and the definition of the variational information I θ,φ we deduce that:

I θ * (X, Z) ≥ I q (X, Z) ≥ I θ,φ (X, Z)

We conclude observing that if Θ = Φ, at optimum the three information terms above are equal, and then q φ (x, z) = p θ (x, z).

Observation: If Φ ⊂ Θ, the variational mutual information can be at most equal to I q * , the maximal I q , but it is not guaranteed that I q * = I θ * .

DERIVATION OF EQUATION 8:

In (Zhao et al., 2017) , it is observed that equation 7 can be written as follows:

V IM θ,φ = -D KL (p(x)||p θ (x)) -E p(x) [D KL (q φ (z|x)||p θ (z|x))]- -(λ -1)D KL (q φ (z)||p(z)) + I q (X, Z).

From equation above, to verify that equation 8 is correct, it is sufficient to show that

EQUATION

The equation 15 follows by the property of the autoencoder and equation 15.

More precisely, by equation 14 we have that

I θ,φ (X, Z) = I q (X, Z) -E q(z) D KL (q φ (x|z)||p θ (x|z))

and, by AutoEncoder architecture, as observed in section 3, p θ (z) = E p(x) [p θ (z|x)] = q(z). Then the following equation holds

E q(z) D KL (q φ (x|z)||p θ (x|z)) = E p(x) D KL (q φ (z|x)||p θ (z|x)).

Indeed,

E q(z) D KL (q φ (x|z)||p θ (x|z)) = q(z) q(x|z) log q(x|z) p θ (x|z) dxdz = q(x)q(z|x) log q(x|z)q(z) p θ (x|z)(z) dxdz = p(x) q(z|x) log q(z|x) p θ (z|x) dxdz = = E p(x) D KL (q φ (z|x)||p θ (z|x)).

FURTHER DETAILS ON EXPERIMENTS:

In all the experiments in section 5.2 we considered the latent space Z = R d , for all the models we choose the prior p(z) to be a Gaussian with zero mean and identity covariance, only in VIMAE-l we choose the prior p(z) to be a logistic with mean zero and identity variance. We choose p θ (x|z) to be similar to DCGAN with batch normalization and q φ (z|x) to be a convolutional deep neural network. The entire models are trained end to end by Adam (Kingma & Ba, 2014) with α = 10 -3 , β 1 = 0.5, β 2 = 0.999. We considered a deterministic decoder and we approximate the reconstruction loss with the L 2 loss, i.e. E p(x) [E q(z|x) [-log p θ (x|z)]] = x -x g 2 2 , with x g indicating the generated datum. In VIMAE case, while training we were adding a pixel-wise Gaussian noise truncated at 0.01 to all the images before feeding them to encoder, in order to make the encoder random. In VAE and β-VAE case, instead we used the standard reparameterization trick (Kingma & Welling, 2013) .

In the following we describe the data-sets considered and the associated neural networks, we follow the same description given in (Tolstikhin et al., 2017) since we used the same neural nets.

MNIST AND OMNIGLOT:

MNIST is a data-set containing 70k grey-scale handwritten digits and associated labels of resolution 28 × 28, subdivided in three subsets: train (50k), validation (10k) and test (10k). Omniglot is a data-set containing 1623 different handwritten characters of resolution 28 × 28 from 50 different alphabets. Each of the 1623 characters was drawn online via Amazon's Mechanical Turk by 20 different people. In the same fashion as done in (Vinyals et al., 2016) we considered an augmented version where each character is rotated respectively by 90, 180, 270 degrees, in this way each character appears 80 times. The Omniglot data-set although has the same resolution of the MNIST, for this reason we use the same network, it is more challenging because, it is more entropic then MNIST, in fact the classes move from 10 to 1623 and the test classes are never seen in the training.

We choose Z = R 8 , and β = λ = 10, we used mini-batches of size 100 and trained the model for 80 epochs. Both encoder and decoder used fully convolutional architectures with 4 × 4 convolutional filters. Encoder:

x ∈ R 28×28 → Conv 128 → BN → ReLu → Conv 256 → BN → ReLu → Conv 512 → BN → ReLu → Conv 1024 → BN → ReLu Decoder: z ∈ R 8 → FC 7×7×1024 → FSConv 512 → BN → ReLu → FSConv 256 → BN → ReLu → FSConv 1

Where Conv k stands for a convolution with k filters, FSConv k for the fractional strided convolution with k filters, BN for batch normalization, ReLU for the rectified linear units, and FC k for the fully connected layer mapping to R k . All the convolutions in the encoder used vertical and horizontal strides 2 and SAME padding.

CELEBA AND CIFAR10:

CelebA is a data-set with 202 599 faces images. We preprocessed the images by first taking a 140 × 140 center crops and then resizing to the 64 × 64 resolution and we consider the last 20k images as test subset. CIFAR10 is a dataset consisting of of 60k 32 × 32 colour images in 10 classes, with 6k images per class. There are 50k training images and 10k test images.

:



Figures and Tables:

Figure 1: 2-d learned representations. Under each plot the model name is followed by the respective MMD value.

Figure 2: Test reconstruction (top) and random generative samples (bottom) of the different methods with Omniglot (right). In test reconstructions, the odd rows are the original data. Each model is denoted by its name and the reconstruction • 2 and rate term.

Semi-supervised classification CIFAR10.

Table content: <table><tr><td/><td/><td/><td/><td colspan="2">accuracy (%)</td><td/><td/><td/><td/><td/></tr><tr><td/><td/><td colspan="2">Method</td><td colspan="4">ν = 0 N (0, 0.3 2 ) B(0.2)</td><td/><td/><td/></tr><tr><td/><td/><td>VAE</td><td/><td>30</td><td>25</td><td>16</td><td/><td/><td/><td/></tr><tr><td/><td/><td colspan="2">β H -VAE</td><td>29</td><td>26</td><td>19</td><td/><td/><td/><td/></tr><tr><td/><td/><td colspan="2">β A -VAE</td><td>31</td><td>31</td><td>18</td><td/><td/><td/><td/></tr><tr><td/><td/><td colspan="2">VIMAE-n</td><td>29</td><td>28</td><td>23</td><td/><td/><td/><td/></tr><tr><td/><td/><td colspan="2">VIMAE-l</td><td>32</td><td>34</td><td>23</td><td/><td/><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td colspan="5">Table 4: Semi-supervised classification,</td></tr><tr><td colspan="6">Table 3: Semi-supervised classification, MNIST.</td><td colspan="5">Omniglot (random sampling: 20%).</td></tr><tr><td/><td/><td colspan="2">accuracy (%)</td><td/><td/><td/><td colspan="2">accuracy (%)</td><td/><td/></tr><tr><td>Method</td><td colspan="5">ν = 0 ν = N (0, σ 2 ) ν = B(p)</td><td colspan="5">ν = 0 ν = N (0, σ 2 ) ν = B(p)</td></tr><tr><td/><td/><td>0.2</td><td>0.4</td><td colspan="2">0.2 0.5</td><td/><td>0.2</td><td>0.4</td><td colspan="2">0.2 0.5</td></tr><tr><td>VAE</td><td>80</td><td>77</td><td>70</td><td>72</td><td>52</td><td>22</td><td>22</td><td>17</td><td>22</td><td>16</td></tr><tr><td>β H -VAE</td><td>92</td><td>86</td><td>82</td><td>91</td><td>84</td><td>21</td><td>21</td><td>22</td><td>19</td><td>17</td></tr><tr><td>β A -VAE</td><td>93</td><td>66</td><td>13</td><td>85</td><td>65</td><td>22</td><td>22</td><td>21</td><td>21</td><td>24</td></tr><tr><td>VIMAE-n</td><td>93</td><td>92</td><td>86</td><td>92</td><td>86</td><td>22</td><td>23</td><td>24</td><td>22</td><td>22</td></tr><tr><td>VIMAE-l</td><td>93</td><td>92</td><td>88</td><td>92</td><td>87</td><td>24</td><td>23</td><td>20</td><td>23</td><td>22</td></tr></table>

Peer review (iclr2020):
The paper develops an information-theoretic training scheme for Variational Auto-Encoders (VAEs). This scheme is tailored for addressing the well-known disentanglement problem of VAEs where an over-capacity encoder sometimes manages to both maximize data fit and shrink the KL divergence between the approximate posterior and prior to zero. Consequently, the latent representations of the observations become independent, making them unusable for any downstream task. The method developed in Section 3 and proposed explicitly in Eqs 7 to 11 is novel per se, though not groundbreaking. Figs 2 and 3 are only few examples manually chosen from rather simple tasks. In the absence of a quantitative evaluation metric, they are not informative. In the outcomes of the same runs, there might exist counterexamples where the vanilla VAE generate perceptively more appealing reconstructions than VIMAE. As a minor comment, I think the paper unnecessarily complicates the presentation of the idea. The development of notions such as f-Divergence, links to MMD, the constrained optimization setting in Eq 5 etc., do not serve to the main story line, but only distracts a mind. I would prefer a brief intro on VAEs and directly jumping into the proposed method. Last but not least, the statement of technical novelty in Section 3 is extremely encrypted. What exactly is *absolutely* novel there and what is prior art? Can the authors give a to-the-point answer to this question during the rebuttal? The paper has two fundamental weaknesses: i) The paper misses a key reference, which addresses the same representation disentanglement problem using the same information-theoretic approach, only with some minor technical divergences: Alemi et al., Fixing a Broken ELBO, ICML, 2018. This paper is a must-cite, plus a key baseline. This submission cannot be treated as a contribution without showing an improvement on top of this extremely closely related work. The problem is the same, the solution is almost the same, and the theoretical implications of the solution are also the same. ii) The paper builds the entire story line around the statement that the representation disentanglement problem is caused by the properties of the ELBO formula and attempts to fix it by developing a new inference technique. This is attitude largely overlooks simpler explanations. For instance, the vanilla VAE assumes a mean field q(z|x) across the z dimensions. This makes q(z|x) fall largely apart from p(x) which definitely does not factorize that way. Earlier work has shown substantial improvements on the quality of q(z|x) when structured variational inference techniques are used, such as normalizing flows. Furthermore, there has also emerged techniques that can very tightly approximate p(x) in closed-form using the change of variables formula without performing any variational inference at all. No need to say, a tight approximation on p(x) means the same for p(z|x) simply using the Bayes rule. This leaves no room for further improvements by information-theoretic inference alternatives. For instance see: Dinh et al., Density Estimation Using Real NVP, ICLR, 2017 The presence of such two strong alternative remedies to the problem addressed by this work makes its fundamentals shaky. The only way to get over this situation is to provide a thorough comparison against these methods, which is obviously missing. --- Post-rebuttal: Thanks to authors for all the effort they put on the updated version. I appreciate that the paper now provides a quantitative comparison. However, my point (ii) remains fully unaddressed and I fully share the theoretical concerns raised by Reviewer 1. All in all, I keep my position on the below the threshold side with a weak reject.