MIM: Mutual Information Machine

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 , logP ( 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 : logP ( x ) ≥ Ez∼qθ ( z|x ) [ log pθ ( x|z ) ] − DKL ( 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 : RVAE ( θ ) = Ex∼P ( x ) [ Ez∼qθ ( z|x ) [ log pθ ( x|z ) ] −DKL ( qθ ( z|x ) ‖P ( z ) ) ] = Ex∼P ( x ) , z∼qθ ( z|x ) [ log pθ ( x|z ) + logP ( 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 logP ( 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 LVAE ( θ ) = −RVAE ( θ ) + Ex∼P ( x ) [ logP ( x ) ] = DKL ( 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 ) , JSD ( θ ) = 1 2 ( DKL ( pθ ( x|z ) P ( z ) ‖MS ) +DKL ( qθ ( z|x ) P ( x ) ‖MS ) ) , ( 3 ) whereMS is an equally weighted mixture of the encoding and decoding distributions ; i.e. , MS = 1 2 ( pθ ( x|z ) P ( z ) + qθ ( z|x ) P ( x ) ) . ( 4 ) 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 RH ( θ ) = 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 ofMS , i.e. , JSD ( θ ) + RH ( θ ) = H ( MS ) . ( 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 logP ( 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 ( MS ) ; i.e. , LCE ( θ ) ≡ H ( MS , Mθ ) = DKL ( MS ‖Mθ ) +H ( MS ) ≥ H ( MS ) , ( 7 ) whereMθ = 12 ( pθ ( x , z ) + qθ ( x , z ) ) , and H ( MS , Mθ ) denotes the cross-entropy betweenMS andMθ . We refer to LCE as the cross-entropy loss . It aims to match the model prior distributions to the anchors , while also minimizing H ( MS ) . 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 LCE 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 LCE using Jensen ’ s inequality , i.e. , LMIM ( θ ) ≡ 1 2 ( H ( MS , qθ ( x , z ) ) +H ( MS , pθ ( x , z ) ) ) ≥ LCE ( θ ) . ( 8 ) Equation ( 8 ) gives us the loss function for the Mutual Information Machine ( MIM ) . It is an average of cross entropy terms between the mixture distributionMS and the model encoding and decoding distributions respectively . To see that this encourages model consistency , it can be shown that LMIM is equivalent to LCE plus a non-negative model consistency regularizer ; i.e. , LMIM ( θ ) = LCE ( θ ) + RMIM ( θ ) . ( 9 ) The non-negativity of RMIM is a simple consequence of LMIM ( θ ) ≥ LCE ( θ ) in ( 8 ) . One can further show ( see Appendix A.2 ) that RMIM ( θ ) satisfies RMIM ( θ ) = 1 2 ( DKL ( MS ‖ pθ ( x , z ) ) +DKL ( MS ‖ qθ ( x , z ) ) ) −DKL ( MS ‖Mθ ) . ( 10 ) One can conclude from ( 10 ) that the regularizer RMIM 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 MS ( 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 LMIM as a sum of fundamental terms that provide some intuition for its expected behavior . In particular , as derived in the Appendix A.3 , LMIM ( θ ) = RH ( θ ) + 1 4 ( DKL ( P ( z ) ‖ pθ ( z ) ) +DKL ( P ( x ) ‖ qθ ( x ) ) ) + 1 4 ( DKL ( qθ ( z|x ) P ( x ) ‖ pθ ( x , z ) ) +DKL ( pθ ( x|z ) P ( z ) ‖ qθ ( z , x ) ) ) ( 11 ) 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 distributionMS : LMIM ( θ ) ≥ HMS ( x ) +HMS ( z ) − IMS ( x ; z ) . Through the MIM loss and the introduction of the parameterized model distributionMθ , we are pushing down on the entropy of the anchored mixture distributionMS , 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 underMS between observations and latent states .

Summary: This paper defines a new learning objective of an autoencoder framework for learning joint distributions over observations and latent states, where the objective is the joint entropy of  M_s, an equally weighted mixture of the encoding and decoding distributions. Since it involves the marginal distribution which is intractable, it goes further to introduce an upper bound which is the cross-entropy between M_s and M_theta. To maintain consistency, it introduce an upper bound again, which is the average of cross entropy between the mixture distribution M_S and the model encoding and decoding distribution. Through comprehensive experiments, the effectiveness is validated.