Mutual Information Gradient Estimation for  Representation Learning

1 INTRODUCTION . Mutual information ( MI ) is an appealing metric widely used in information theory and machine learning to quantify the amount of shared information between a pair of random variables . Specifically , given a pair of random variables x , y , the MI , denoted by I ( x ; y ) , is defined as I ( x ; y ) = Ep ( x , y ) [ log p ( x , y ) p ( x ) p ( y ) ] , ( 1 ) where E is the expectation over the given distribution . Since MI is invariant to invertible and smooth transformations , it can capture non-linear statistical dependencies between variables ( Kinney & Atwal , 2014 ) . These appealing properties make it act as a fundamental measure of true dependence . Therefore , MI has found applications in a wide range of machine learning tasks , including feature selection ( Kwak & Choi , 2002 ; Fleuret , 2004 ; Peng et al. , 2005 ) , clustering ( Müller et al. , 2012 ; Ver Steeg & Galstyan , 2015 ) , and causality ( Butte & Kohane , 1999 ) . It has also been pervasively used in science , such as biomedical sciences ( Maes et al. , 1997 ) , computational biology ( Krishnaswamy et al. , 2014 ) , and computational neuroscience ( Palmer et al. , 2015 ) . Recently , there has been a revival of methods in unsupervised representation learning based on MI . A seminal work is the InfoMax principle ( Linsker , 1988 ) , where given an input instance x , the goal of the InfoMax principle is to learn a representation Eψ ( x ) by maximizing the MI between the input and its representation . A growing set of recent works have demonstrated promising empirical performance in unsupervised representation learning via MI maximization ( Krause et al. , 2010 ; Hu et al. , 2017 ; Alemi et al. , 2018b ; Oord et al. , 2018 ; Hjelm et al. , 2019 ) . Another closely related work is the Information Bottleneck method ( Tishby et al. , 2000 ; Alemi et al. , 2017 ) , where MI is used to limit the contents of representations . Specifically , the representations are learned by extracting taskrelated information from the original data while being constrained to discard parts that are irrelevant to the task . Several recent works have also suggested that by controlling the amount of information between learned representations and the original data , one can tune desired characteristics of trained models such as generalization error ( Tishby & Zaslavsky , 2015 ; Vera et al. , 2018 ) , robustness ( Alemi et al. , 2017 ) , and detection of out-of-distribution data ( Alemi et al. , 2018a ) . Despite playing a pivotal role across a variety of domains , MI is notoriously intractable . Exact computation is only tractable for discrete variables , or for a limited family of problems where the probability distributions are known . For more general problems , MI is challenging to analytically compute or estimate from samples . A variety of MI estimators have been developed over the years , including likelihood-ratio estimators ( Suzuki et al. , 2008 ) , binning ( Fraser & Swinney , 1986 ; Darbellay & Vajda , 1999 ; Shwartz-Ziv & Tishby , 2017 ) , k-nearest neighbors ( Kozachenko & Leonenko , 1987 ; Kraskov et al. , 2004 ; Pérez-Cruz , 2008 ; Singh & Póczos , 2016 ) , and kernel density estimators ( Moon et al. , 1995 ; Kwak & Choi , 2002 ; Kandasamy et al. , 2015 ) . However , few of these mutual information estimators scale well with dimension and sample size in machine learning problems ( Gao et al. , 2015 ) . In order to overcome the intractability of MI in the continuous and high-dimensional settings , Alemi et al . ( 2017 ) combines variational bounds of Barber & Agakov ( 2003 ) with neural networks for the estimation . However , the tractable density for the approximate distribution is required due to variational approximation . This limits its application to the general-purpose estimation , since the underlying distributions are often unknown . Alternatively , the Mutual Information Neural Estimation ( MINE , Belghazi et al . ( 2018 ) ) and the Jensen-Shannon MI estimator ( JSD , Hjelm et al . ( 2019 ) ) enable differentiable and tractable estimation of MI by training a discriminator to distinguish samples coming from the joint distribution or the product of the marginals . In detail , MINE employs a lower-bound to the MI based on the Donsker-Varadhan representation of the KL-divergence , and JSD follows the formulation of f-GAN KL-divergence . In general , these estimators are often noisy and can lead to unstable training due to their dependence on the discriminator used to estimate the bounds of mutual information . As pointed out by Poole et al . ( 2019 ) , these unnormalized critic estimators of MI exhibit high variance and are challenging to tune for estimation . An alternative low-variance choice of MI estimator is Information Noise-Contrastive Estimation ( InfoNCE , Oord et al . ( 2018 ) ) , which introduces the Noise-Contrastive Estimation with flexible critics parameterized by neural networks as a bound to approximate MI . Nonetheless , its estimation saturates at log of the batch size and suffers from high bias . Despite their modeling power , none of the estimators are capable of providing accurate estimation of MI with low variance when the MI is large and the batch size is small ( Poole et al. , 2019 ) . As supported by the theoretical findings in McAllester & Statos ( 2018 ) , any distribution-free high-confidence lower bound on entropy requires a sample size exponential in the size of the bound . More discussions about the bounds of MI and their relationship can be referred to Poole et al . ( 2019 ) . In summary , existing estimators first approximate MI and then use these approximations to optimize the associated parameters . For estimating MI based on any finite number of samples , there exists an infinite number of functions , with arbitrarily diverse gradients , that can perfectly approximate the true MI at these samples . However , these approximate functions can lead to unstable training and poor performance in optimization due to gradients discrepancy between approximate estimation and true MI . Estimating gradients of MI rather than estimating MI may be a better approach for MI optimization . To this end , to the best of our knowledge , we firstly propose the Mutual Information Gradient Estimator ( MIGE ) in representation learning . In detail , we estimate the score function of an implicit distribution , ∇x log q ( x ) , to achieve a general-purpose MI gradient estimation for representation learning . In particular , to deal with high-dimensional inputs , such as text , images and videos , score function estimation via Spectral Stein Gradient Estimator ( SSGE ) ( Shi et al. , 2018 ) is computationally expensive and complex . We thus propose an efficient high-dimensional score function estimator to make SSGE scalable . To this end , we derive a new reparameterization trick for the representation distribution based on the lower-variance reparameterization trick proposed by Roeder et al . ( 2017 ) . We summarize the contributions of this paper as follows : • We propose the Mutual Information Gradient Estimator ( MIGE ) for representation learning based on the score function estimation of implicit distributions . Compared with MINE and MINE-f , MIGE provides a tighter and smoother gradient estimation of MI in a highdimensional and large-MI setting , as shown in Figure 1 of Section 4 . • We propose the Scalable SSGE to alleviate the exorbitant computational cost of SSGE in high-dimensional settings . • To learn meaningful representations , we apply SSGE as gradient estimators for both InfoMax and Information Bottlenck , and have achieved improved performance than their corresponding competitors . 2 SCALABLE SPECTRAL STEIN GRADIENT ESTIMATOR . Score estimation of implicit distributions has been widely explored in the past few years ( Song et al. , 2019 ; Li & Turner , 2017 ; Shi et al. , 2018 ) . A promising method of score estimation is the Stein gradient estimator ( Li & Turner , 2017 ; Shi et al. , 2018 ) , which is proposed for implicit distributions . It is inspired by generalized Steins identity ( Gorham & Mackey , 2015 ; Liu & Wang , 2016 ) as follows . Steins identity . Let q ( x ) be a continuously differentiable ( also called smooth ) density supported on X ⊆ Rd , and h ( x ) = [ h1 ( x ) , h2 ( x ) , . . . , hd′ ( x ) ] T is a smooth vector function . Further , the boundary conditions on h is q ( x ) h ( x ) = 0 , ∀x ∈ ∂X if X is compact , or lim x→∞ q ( x ) h ( x ) = 0 if X = Rd . ( 2 ) Under this condition , the following identity can be easily checked using integration by parts , assuming mild zero boundary conditions on h , Eq [ h ( x ) ∇x log q ( x ) T +∇xh ( x ) ] = 0 . ( 3 ) Here h is called the Stein class of q ( x ) if Steins identity Eq . ( 3 ) holds . Monte Carlo estimation of the expectation in Eq . ( 3 ) builds the connection between ∇x log q ( x ) and the samples from q ( x ) in Steins identity . For modeling implicit distributions , Motivated by Steins identity , Shi et al . ( 2018 ) proposed Spectral Stein Gradient Estimator ( SSGE ) for implicit distributions based on Stein ’ s identity and a spectral decomposition of kernel operators where the eigenfunctions are approximated by the Nyström method . Below we briefly review SSGE . More details refer to Shi et al . ( 2018 ) . Specifically , we denote the target gradient function to estimate by g : X → Rd : g ( x ) = ∇x log q ( x ) . The ith component of the gradient is gi ( x ) = ∇xi log q ( x ) . We assume g1 , . . . , gd ∈ L2 ( X , q ) . { ψj } j≥1 denotes an orthonormal basis of L 2 ( X , q ) . We can expand gi ( x ) into the spectral series , i.e. , gi ( x ) = ∑∞ j=1 βijψj ( x ) . The value of the j th eigenfunction ψj at x can be approximated by the Nyström method ( Xu et al. , 2015 ) . Due to the orthonormality of eigenfunctions { ψj } j≥1 , there is a constraint under the probability measure q ( . ) : ∫ ψi ( x ) ψj ( x ) q ( x ) dx = δij , where δij = 1 [ i = j ] . Based on this constraint , we can obtain the following equation for { ψj } j≥1 : ∫ k ( x , y ) ψ ( y ) q ( y ) dy = µψ ( x ) , ( 4 ) where k ( . ) is a kernel function . The left side of the above equation can be approximated by the Monte Carlo estimate using i.i.d . samples x1 , ... , xM from q ( . ) : 1MKψ ≈ µψ , where K is the Gram Matrix and ψ = [ ψ ( x1 ) , . . . , ψ ( xM ) ] > . We can solve this eigenvalue problem by choose the J largest eigenvalues λ1 ≥ · · · ≥ λJ for K. uj denotes the eigenvector of the Gram matrix . The approximation for { ψj } j≥1 can be obtained combined with Eq . ( 4 ) as following : ψj ( x ) ≈ ψ̂j ( x ) =√ M λj ∑M m=1 ujmk ( x , x m ) . Furthermore , based on the orthonormality of { ψj } j≥1 , we can easily obtain βij = −Eq∇xiψj ( x ) . By taking derivative both sides of Eq . ( 4 ) , we can show that : µj∇xiψj ( x ) = ∇xi ∫ k ( x , y ) ψj ( y ) q ( y ) dy = ∫ ∇xik ( x , y ) ψj ( y ) q ( y ) dy . ( 5 ) Then we can estimate as following : ∇̂xiψj ( x ) ≈ 1 µjM M∑ m=1 ∇xik ( x , xm ) ψj ( xm ) . ( 6 ) Finally , by truncating the expansion to the first J terms and plugging in the Nyström approximations of { ψj } j≥1 , we can get the score estimator : ĝi ( x ) = J∑ j=1 β̂ijψ̂j ( x ) , β̂ij = − 1 M M∑ m=1 ∇xi ψ̂j ( xm ) . ( 7 ) In general , representation learning for large-scale datasets is usually costly in terms of storage and computation . For instance , the dimension of images in the STL-10 dataset is 96 × 96 × 3 ( i.e. , the vector length is 27648 ) . This makes it almost impossible to directly estimate the gradient of MI between the input and representation . To alleviate this problem , we introduce random projection ( RP ) ( Bingham & Mannila , 2001 ) to reduce the dimension of x . We briefly review RP . More details refer to Bingham & Mannila ( 2001 ) . RP projects the original d-dimensional data into a k-dimensional ( k < < d ) subspace . Concretely , let matrix Xd×N denotes the original set of N d-dimensional data , the projection of the original data XRPk×N is obtained by introducing a random matrix Rk×d whose columns have unit length , as follows ( Bingham & Mannila , 2001 ) , XRPk×N = Rk×dXd×N . After RP , the Euclidean distance between two original data vectors can be approximated by the Euclidean distance of the projective vectors in reduced spaces : ‖x1 − x2‖ ≈ √ d/k ‖Rx1 −Rx2‖ , ( 8 ) where x1 and x2 denote the two data vectors in the original large dimensional space . Based on the principle of RP , we can derive a Salable Spectral Stein Gradient Estimator , which is an efficient high-dimensional score function estimator . One can show that the RBF kernel satisfies Steins identity ( Liu & Wang , 2016 ) . Shi et al . ( 2018 ) also shows that it is a promising choice for SSGE with a lower error bound . To reduce the computation of the kernel similarities of SSGE in high-dimensional settings , we replace the input of SSGE with a projections obtained by RP according to the approximation of Eq . ( 8 ) for the computation of the RBF kernel .

Summary: This paper proposes MIGE---a novel estimator of the mutual information (MI) gradient, based on estimating the score function of an implicit distribution. To this end, the authors employ the spectral Stein gradient estimator (SSGE) and propose its scalable version based on random projections of the original input. The theoretical advantages of the method are presented using a toy experiment with correlated Gaussian random variables, where both the mutual information and its gradient can be computed analytically. In this setting, MIGE provides gradient estimates that are less biased and smoother than baselines. The method is also evaluated on two more complicated tasks: unsupervised representation learning on Cifar-10 and CIfar-100 via DeepInfoMax (DIM) and classification on MNIST with Information Bottleneck (IB), where MIGE outperforms all baselines by a significant margin.