Improving Mutual Information Estimation with Annealed and Energy-Based Bounds

Mutual information ( MI ) is a fundamental quantity in information theory and machine learning . However , direct estimation of mutual information is intractable , even if the true joint probability density for the variables of interest is known , as it involves estimating a potentially high dimensional log partition function . In this work , we view mutual information estimation from the perspective of importance sampling . Since naive importance sampling with the marginal density as a proposal requires exponential sample complexity in the true mutual information , we propose Multi-Sample Annealed Importance Sampling ( AIS ) bounds on mutual information to bridge between the marginal and joint distributions . In settings where the full joint distribution is available , we provide lower and upper bounds that can tightly estimate large values of MI in our experiments . In settings where only a single marginal distribution is known , we improve upon existing variational methods by directly optimizing a tighter lower bound on MI , using energy-based training to estimate gradients and multi-sample AIS for evaluation . Our methods are particularly suitable for evaluating MI in deep generative models , since explicit forms for the marginal or joint densities are often available . We evaluate our bounds on estimating the MI of VAEs and GANs trained on the MNIST and CIFAR datasets , and showcase significant gains over existing bounds in these challenging settings with high ground truth MI . 1 INTRODUCTION . Mutual information ( MI ) is among the most general measures of dependence between two random variables . Among many other applications in machine learning , mutual information has been used for both training ( Alemi et al. , 2016 ; 2018 ; Chen et al. , 2016 ; Zhao et al. , 2018 ) and evaluating ( Alemi & Fischer , 2018 ; Huang et al. , 2020 ) generative models . Furthermore , a recent success in neural network function approximation has encouraged a wave of variational or contrastive methods for representation learning ( van den Oord et al. , 2018 ) or MI estimation from samples alone ( Belghazi et al. , 2018 ; Poole et al. , 2019 ) . However , McAllester & Stratos ( 2020 ) have shown strong theoretical limitations on any estimator which does not have access to an analytical form of at least one marginal distribution . In light of these limitations , we consider MI estimation in settings where a single marginal or the full joint distribution are available . In this work , we propose novel Multi-Sample Annealed Importance Sampling ( AIS ) bounds which can tightly estimate large values of MI when the full joint distribution is known . We motivate these bounds by viewing MI estimation from the perspective of importance sampling , which also suggests novel variational approaches for the case where only a single marginal is known . First , our Generalized IWAE ( GIWAE ) bounds highlight how multi-sample contrastive estimation can complement variational learning to improve MI lower bounds , and includes INFONCE ( van den Oord et al. , 2018 ; Poole et al. , 2019 ) and importance-weighted autoencoder ( IWAE ) ( Burda et al. , 2016 ; Sobolev & Vetrov , 2019 ) bounds as special cases . We also propose MINE-AIS , which uses an improved energy-based lower bound compared with Mutual Information Neural Estimation ( MINE ) ( Belghazi et al. , 2018 ) and corresponds to the infinite-sample limit of the GIWAE framework . We propose an efficient training procedure using Markov Chain Monte Carlo ( MCMC ) and utilize our Multi-Sample AIS evaluation bounds to show notable improvement over finite-sample methods for lower bounding MI . In Fig . 1 , we summarize the MI bounds discussed in this paper and the relationships between them . 1.1 PROBLEM SETTING . The mutual information between two random variables x and z with joint distribution p ( x , z ) is I ( x ; z ) = Ep ( x , z ) [ log p ( x , z ) p ( x ) p ( z ) ] = H ( x ) −H ( x|z ) = Ep ( x , z ) [ log p ( x|z ) ] − Ep ( x ) [ log p ( x ) ] , ( 1 ) whereH ( x|z ) denotes the conditional entropy −Ep ( x , z ) log p ( x|z ) . We primarily focus on bounds that assume either a single marginal distribution or the full joint distribution are available . A natural setting where the full joint distribution is available is estimating the mutual information in deep generative models between the latent variables , with a known prior z ∼ p ( z ) , and data x ∼ p ( x ) simulated from the model ( Alemi & Fischer , 2018 ) . 1 Settings where a single marginal is available appear in simulation-based inference ( Cranmer et al. , 2020 ) , where information about input parameters θ is known and a simulator can generate x for a given θ , but the likelihood p ( x|θ ) is intractable . While sampling from the posterior p ( z|x ) for an arbitrary x is often intractable , we can obtain a single posterior sample for x ∼ p ( x ) , z ∼ p ( z|x ) in cases where samples from the joint distribution of interest p ( x ) p ( z|x ) are available . Throughout this paper , we will refer to bounds which involve only a single posterior sample as practical , and those involving multiple posterior samples as impractical . When the conditional p ( x|z ) is tractable to sample and evaluate , simple Monte Carlo sampling provides an unbiased , low variance estimate of the conditional entropy term in Eq . 1 . In this case , the difficulty of MI estimation reduces to estimating the log partition function , for which importance sampling ( IS ) based methods are among the most well studied and successful solutions . 2 MULTI-SAMPLE BOUNDS ON MUTUAL INFORMATION . In this section , we review single- and multi-sample MI bounds from the perspective of importance sampling . This provides insight into existing bounds and will provide foundation for our primary contributions in Sec . 3 and 4 . 2.1 SIMPLE IMPORTANCE SAMPLING BOUNDS . For an arbitrary conditional distribution qθ ( z|x ) , we can obtain lower and upper bounds on log p ( x ) : Eqθ ( z|x ) [ log p ( x , z ) qθ ( z|x ) ] ≤ log p ( x ) ≤ Ep ( z|x ) [ log p ( x , z ) qθ ( z|x ) ] ( 2 ) where the left-hand side is the familiar Evidence Lower Bound ( ELBO ) with a gap of DKL [ qθ ( z|x ) ||p ( z|x ) ] = log p ( x ) − ELBO ( qθ ) . We refer to the right-hand side as the Evidence Upper Bound ( EUBO ) , with gap DKL [ p ( z|x ) ||qθ ( z|x ) ] = EUBO ( qθ ) − log p ( x ) . Throughout this paper , we will derive bounds which correspond to extended-state space versions of Eq . 2 ( see App . A ) . By plugging the importance sampling lower and upper bounds on log p ( x ) from Eq . 2 into Eq . 1 , we obtain upper and lower bounds on mutual information as follows IBAL ( qθ ) = Ep ( x , z ) [ log qθ ( z|x ) p ( z ) ] ≤ I ( x ; z ) ≤ Ep ( x ) qθ ( z|x ) [ log qθ ( z|x ) p ( x , z ) ] −H ( x|z ) = IBAU ( qθ ) ( 3 ) The left hand side of Eq . 3 is the well-known Barber-Agakov ( BA ) bound ( Barber & Agakov , 2003 ; Agakov , 2006 ) , which we denote by IBAL ( qθ ) = I ( x ; z ) −DKL [ p ( z|x ) ||qθ ( z|x ) ] . We refer to the right hand side as the BA upper bound IBAU ( qθ ) = I ( x ; z ) +DKL [ qθ ( z|x ) ||p ( z|x ) ] . 1An alternative , ‘ encoding ’ MI between the real data and the latent code is often of interest in representation learning ( see App . F ) , but can not be directly estimated using our methods due to the unavailability of densities for the data distribution pd ( x ) or the ‘ aggregated posterior ’ q ( z ) = ∫ pd ( x ) q ( z|x ) dx . Relationship with CLUB When the prior is used as the proposal , IBAU ( p ( z ) ) corresponds to the Contrastive Log-Ratio Upper Bound ( CLUB ) ( Cheng et al. , 2020 ) with known p ( x|z ) . 2.2 IMPORTANCE WEIGHTED AUTOENCODER . For a given variational distribution qθ ( z|x ) , IWAE bounds ( Burda et al. , 2016 ; Sobolev & Vetrov , 2019 ) improve upon simple importance sampling by considering multi-sample proposal and target distributions in an extended state space ( Domke & Sheldon , 2018 ; Cremer et al. , 2017 ) . Consider a proposal qIWAEPROP with K independent samples from the variational distribution qθ ( z|x ) . The extended state space target pIWAETGT is a mixture distribution involving a single sample from the posterior p ( z|x ) or joint p ( x , z ) distribution and K − 1 samples from qθ ( z|x ) qIWAEPROP ( z ( 1 : K ) |x ) = K∏ k=1 qθ ( z ( k ) |x ) pIWAETGT ( x , z ( 1 : K ) ) = 1 K K∑ s=1 p ( x , z ( s ) ) K∏ k=1 , k 6=s qθ ( z ( k ) |x ) . ( 4 ) The log ratio log p IWAE TGT ( x , z ( 1 : K ) ) qIWAEPROP ( z ( 1 : K ) |x ) reduces to familiar ratio appearing in the IWAE objective , while the ratio of normalization constants corresponds to log p ( x ) . As in App . A , taking expectations under the proposal yields a lower bound , whereas taking expectations under the target yields an upper bound Eqθ ( z ( 1 : K ) |x ) [ log 1 K K∑ i=1 p ( x , z ( k ) ) qθ ( z ( k ) |x ) ] ︸ ︷︷ ︸ ELBOIWAE ( qθ , K ) ≤ log p ( x ) ≤ Ep ( z ( 1 ) |x ) qθ ( z ( 2 : K ) |x ) [ log 1 K K∑ i=1 p ( x , z ( k ) ) qθ ( z ( k ) |x ) ] ︸ ︷︷ ︸ EUBOIWAE ( qθ , K ) . ( 5 ) Note that the indices z ( 1 : K ) denote a set of independent draws from a given distribution . As for the ELBO and EUBO , the gap in the lower and upper bounds areDKL [ qIWAEPROP ||pIWAETGT ] andDKL [ pIWAETGT ||qIWAEPROP ] , respectively . See App . B for detailed derivations . As discussed in Sec . 1.1 , lower and upper bounds on the log partition function translate to upper and lower bounds , respectively , on MI when p ( x|z ) is known . We refer to these bounds arising from IWAE estimation as IIWAEL ( qθ , K ) and IIWAEU ( qθ , K ) . For example , IIWAEL ( qθ , K ) can be written as I ( x ; z ) ≥ IIWAEL ( qθ , K ) : = Ep ( x ) p ( z ( 1 ) |x ) qθ ( z ( 2 : K ) |x ) [ log p ( x|z ( 1 ) ) 1 K ∑K i=1 p ( x , z ( k ) ) qθ ( z ( k ) |x ) ] ( 6 ) While it is well-known that increasing K leads to tighter IWAE bounds ( Burda et al. , 2016 ; Sobolev & Vetrov , 2019 ) , we characterize the logarithmic rate of bias reduction for IIWAEL ( qθ , K ) in the following proposition , with proof in App . B.2 . Proposition 2.1 . IWAE bounds on MI improve upon the BA bounds with the following relationships : IBAL ( qθ ) ≤ IIWAEL ( qθ , K ) ≤ IBAL ( qθ ) + logK , IIWAEU ( qθ , K ) ≤ IBAU ( qθ ) . ( 7 )

Summary: This work combines the technique of Annealed Importance Sampling (AIS) with Importance Weighted Autoencoder (IWAE) bounds on mutual information to obtain new and significantly tighter bounds. They examine the situations where the joint distribution $p(x, z)$ is known via $p(z)$ and $p(x | z)$, and the situation where only $p(z)$ is known. In the former case, they obtain exact bounds (or asymptotically exact given large enough sample sizes), and in the latter case they find approximate bounds by learning an energy function to approximate $p(x |z)$. Their experiments show the utility of their proposed method. In particular, they obtain impressively tight lower and upper bounds for complex latent variable generative models (VAE and GAN) on complex data such as CIFAR-10.