Understanding the Limitations of Variational Mutual Information Estimators

Variational approaches based on neural networks are showing promise for estimating mutual information ( MI ) between high dimensional variables . However , they can be difficult to use in practice due to poorly understood bias/variance tradeoffs . We theoretically show that , under some conditions , estimators such as MINE exhibit variance that could grow exponentially with the true amount of underlying MI . We also empirically demonstrate that existing estimators fail to satisfy basic self-consistency properties of MI , such as data processing and additivity under independence . Based on a unified perspective of variational approaches , we develop a new estimator that focuses on variance reduction . Empirical results on standard benchmark tasks demonstrate that our proposed estimator exhibits improved biasvariance trade-offs on standard benchmark tasks . 1 INTRODUCTION . Mutual information ( MI ) estimation and optimization are crucial to many important problems in machine learning , such as representation learning ( Chen et al. , 2016 ; Zhao et al. , 2018b ; Tishby & Zaslavsky , 2015 ; Higgins et al. , 2018 ) and reinforcement learning ( Pathak et al. , 2017 ; van den Oord et al. , 2018 ) . However , estimating mutual information from samples is challenging ( McAllester & Statos , 2018 ) and traditional parametric and non-parametric approaches ( Nemenman et al. , 2004 ; Gao et al. , 2015 ; 2017 ) struggle to scale up to modern machine learning problems , such as estimating the MI between images and learned representations . Recently , there has been a surge of interest in MI estimation with variational approaches ( Barber & Agakov , 2003 ; Nguyen et al. , 2010 ; Donsker & Varadhan , 1975 ) , which can be naturally combined with deep learning methods ( Alemi et al. , 2016 ; van den Oord et al. , 2018 ; Poole et al. , 2019 ) . Despite their empirical effectiveness in downstream tasks such as representation learning ( Hjelm et al. , 2018 ; Velicković et al. , 2018 ) , their effectiveness for MI estimation remains unclear . In particular , higher estimated MI between observations and learned representations do not seem to indicate improved predictive performance when the representations are used for downstream supervised learning tasks ( Tschannen et al. , 2019 ) . In this paper , we discuss two limitations of variational approaches to MI estimation . First , we theoretically demonstrate that the variance of certain estimators , such as MINE ( Belghazi et al. , 2018 ) , could grow exponentially with the ground truth MI , leading to poor bias-variance trade-offs . Second , we propose a set of self-consistency tests over basic properties of MI , and empirically demonstrate that all considered variational estimators fail to satisfy critical properties of MI , such as data processing and additivity under independence . These limitations challenge the effectiveness of these methods for estimating or optimizing MI . To mitigate these issues , we propose a unified perspective over variational estimators treating variational MI estimation as an optimization problem over ( valid ) density ratios . This view highlights the role of partition functions estimation , which is the culprit of high variance issues in MINE . To address this issue , we propose to improve MI estimation via variance reduction techniques for partition function estimation . Empirical results demonstrate that our estimators have much better bias-variance trade-off compared to existing methods on standard benchmark tasks . 2 BACKGROUND AND RELATED WORK . 2.1 NOTATIONS . We use uppercase letters to denote a probability measure ( e.g. , P , Q ) and corresponding lowercase letters to denote its density1 functions ( e.g. , p , q ) unless specified otherwise . We use X , Y to denote random variables with separable sample spaces denoted as X and Y respectively , and P ( X ) ( or P ( Y ) ) to denote the set of all probability measures over the Borel σ-algebra on X ( or Y ) . Under Q ∈ P ( X ) , the p-norm of a function r : X → R is defined as ‖r‖p : = ( ∫ |r|pdQ ) 1/p with ‖r‖∞ = limp→∞‖r‖p . The set of locally p-integrable functions is defined as Lp ( Q ) : = { r : X → R | ‖r‖p < ∞ } . The space of probability measures wrt . Q is defined as ∆ ( Q ) : = { r ∈ L1 ( Q ) | ‖r‖1 = 1 , r ≥ 0 } ; we also call this the space of “ valid density ratios ” wrt . Q . We use P Q to denote that P is absolutely continuous with respect to Q . We use ÎE to denote an estimator for IE where we replace expectations with sample averages . 2.2 VARIATIONAL MUTUAL INFORAMTION ESTIMATION . The mutual information between two random variables X and Y is the KL divergence between the joint and the product of marginals : I ( X ; Y ) = DKL ( P ( X , Y ) ‖P ( X ) P ( Y ) ) ( 1 ) which we wish to estimate using samples from P ( X , Y ) ; in certain cases we may know the density of marginals ( e.g . P ( X ) ) . There are a wide range of variational approaches to variational MI estimation . Variational information maximization uses the following result ( Barber & Agakov , 2003 ) : Lemma 1 ( Barber-Agakov ( BA ) ) . For two random variables X and Y : I ( X ; Y ) = sup qφ { EP ( X , Y ) [ log qφ ( x|y ) − log p ( x ) ] = : IBA ( qφ ) } ( 2 ) where qφ : Y → P ( X ) is a valid conditional distribution over X given y ∈ Y and p ( x ) is the probability density function of the marginal distribution P ( X ) . Another family of approaches perform MI estimation through variational lower bounds to KL divergences . For example , the Mutual Information Neural Estimator ( MINE , Belghazi et al . ( 2018 ) ) applies the following lower bound to KL divergences ( Donsker & Varadhan , 1975 ) . Lemma 2 ( Donsker-Varadahn ( DV ) ) . ∀P , Q ∈ P ( X ) such that P Q , DKL ( P‖Q ) = sup T∈L∞ ( Q ) { EP [ T ] − logEQ [ eT ] = : IMINE ( T ) } . ( 3 ) One could set P = P ( X , Y ) and Q = P ( X ) P ( Y ) , T as a parametrized neural network ( e.g . Tθ ( x , y ) parametrized by θ ) , and obtain the estimate by optimizing the above objective via stochastic gradient descent over mini-batches . However , the corresponding estimator ÎMINE ( where we replace the expectations in Eq . ( 3 ) with sample averages ) is biased , leading to biased gradient estimates ; Belghazi et al . ( 2018 ) propose to reduce bias via estimating the partition function EQ [ eT ] with exponential moving averages of mini-batches . The variational f -divergence estimation approach ( Nguyen et al. , 2010 ; Nowozin et al. , 2016 ) considers lower bounds on f -divergences which can be specialize to KL divergence , and subsequently to mutual information estimation : Lemma 3 ( Nyugen et al . ( NWJ ) ) . ∀P , Q ∈ P ( X ) such that P Q , DKL ( P‖Q ) = sup T∈L∞ ( Q ) { EP [ T ] − EQ [ eT−1 ] = : INWJ ( T ) } ( 4 ) and DKL ( P‖Q ) = INWJ ( T ) when T = log ( dP/ dQ ) + 1 . 1In the remainder of the paper , we slightly overload “ density ” for discrete random variables . The supremum over T is a invertible function of the density ratio dP/ dQ , so one could use this approach to estimate density ratios by inverting the function ( Nguyen et al. , 2010 ; Nowozin et al. , 2016 ; Grover & Ermon , 2017 ) . The corresponding mini-batch estimator ( denoted as ÎNWJ ) is unbiased , so unlike MINE , this approach does not require special care to reduce bias in gradients . Contrastive Predictive Coding ( CPC , van den Oord et al . ( 2018 ) ) considers the following objective : ICPC ( fθ ) : = EPn ( X , Y ) [ 1 n n∑ i=1 log fθ ( xi , yi ) 1 n ∑n j=1 fθ ( xi , yj ) ] ( 5 ) where fθ : X × Y → R≥0 is a neural network parametrized by θ and Pn ( X , Y ) denotes the joint pdf for n i.i.d . random variables sampled from P ( X , Y ) . CPC generally has less variance but is more biased because its estimate does not exceed log n , where n is the batch size ( van den Oord et al. , 2018 ; Poole et al. , 2019 ) . While one can further reduce the bias with larger n , the number of evaluations needed for estimating each batch with fθ is n2 , which scales poorly . To address the high-bias issue of CPC , Poole et al . proposed an interpolation between ICPC and INWJ to obtain more fine-grained bias-variance trade-offs . 3 VARIATIONAL MUTUAL INFORMATION ESTIMATION AS OPTIMIZATION OVER DENSITY RATIOS . In this section , we unify several existing methods for variational mutual information estimation . We first show that variational mutual information estimation can be formulated as a constrained optimization problem , where the feasible set is ∆ ( Q ) , i.e . the valid density ratios with respect to Q. Theorem 1 . ∀P , Q ∈ P ( X ) such that P Q we have DKL ( P‖Q ) = sup r∈∆ ( Q ) EP [ log r ] ( 6 ) where the supremum is achived when r = dP/ dQ . We defer the proof in Appendix A . The above argument works for KL divergence between general distributions , but in this paper we focus on the special case of mutual information estimation . For the remainder of the paper , we use P to represent the short-hand notation for the joint distribution P ( X , Y ) and use Q to represent the short-hand notation for the product of marginals P ( X ) P ( Y ) . 3.1 A SUMMARY OF EXISTING VARIATIONAL METHODS . From Theorem 1 , we can describe a general approach to variational MI estimation : 1 . Obtain a density ratio estimate – denote the solution as r ; 2 . Project r to be close to ∆ ( Q ) – in practice we only have samples from Q , so we denote the solution as Γ ( r ; Qn ) , where Qn is the empirical distribution of n i.i.d . samples from Q ; 3 . Estimate mutual information with EP [ log Γ ( r ; Qn ) ] . We illustrate two examples of variational mutual information estimation that can be summarized with this approach . In the case of Barber-Agakov , the proposed density ratio estimate is rBA = qφ ( x|y ) /p ( x ) ( assuming that p ( x ) is known ) , which is guaranteed to be in ∆ ( Q ) because EQ [ qφ ( x|y ) /p ( x ) ] = ∫ qφ ( x|y ) /p ( x ) dP ( x ) dP ( y ) = 1 , ΓBA ( rBA , Qn ) = rBA ( 7 ) for all conditional distributions qφ . In the case of MINE / Donsker-Varadahn , the logarithm of the density ratio is estimated with Tθ ( x , y ) ; the corresponding density ratio might not be normalized , so one could apply the following normalization for n samples : EQn [ eTθ/EQn [ eTθ ] ] = 1 , ΓMINE ( e Tθ , Qn ) = e T θ /EQn [ eTθ ] ( 8 ) where EQn [ eTθ ] ( the sample average ) is an unbiased estimate of the partition function EQ [ eTθ ] ; ΓDV ( e Tθ , Qn ) ∈ ∆ ( Q ) is only true when n→∞ . Similarly , we show ICPC is a lower bound to MI in Corollary 2 , Appendix A , providing an alternative proof to the one in Poole et al . ( 2019 ) . These examples demonstrate that different mutual information estimators can be obtained in a procedural manner by implementing the above steps , and one could involve different objectives at each step . For example , one could estimate density ratio via logistic regression ( Hjelm et al. , 2018 ; Poole et al. , 2019 ; Mukherjee et al. , 2019 ) while using INWJ or IMINE to estimate MI . While logistic regression does not optimize for a lower bound for KL divergence , it provides density ratio estimates between P and Q which could be used for subsequent steps .

Summary: This work summarizes the existing methods of mutual information estimation in a variational inference framework and describes the limitations in terms of bias-variance tradeoffs. Further, the authors care about the self-consistency, namely, independence, data processing, and additivity, which are properties of both entropy and differential entropy. Further, density ratio clipping is proposed to lessen a high-variance problem in estimating a partition function.