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 paper investigates annealed importance sampling (AIS) techniques and multi-sample bounds to improve mutual information estimation. Assuming the knowledge of extract likelihood (joint or marginal), this paper shows the mutual information bound can be tightened via exploiting a series of intermediate distributions interpolating between a (convenient) proposal and the joint density. An energy-based perspective was presented and discussed, revealing interesting insights. The effectiveness of this proposal is demonstrated with generative modeling applications on real-world datasets.