A Data-Efficient Mutual Information Neural Estimator for Statistical Dependency Testing

1 INTRODUCTION . Mutual Information ( MI ) is an important , theoretically grounded measure of similarity between random variables . MI captures general , non-linear , statistical dependencies between random variables . MI estimators that estimate MI from samples are important tools widely used in not only subjects such as physics and neuroscience , but also machine learning ranging from feature selection and representation learning to explaining decisions and analyzing generalization of neural networks . Existing studies on MI estimation between general random variables focus on deriving asymptotic lower bounds and approximations to MI under infinite data , and techniques for reducing estimator bias such as bias correction , improved signal modeling with neural networks and tighter lower bounds . Widely used approaches include the k-NN-based KSG estimator ( Kraskov et al. , 2004 ) and the variational lower-bound-based Mutual Information Neural Estimator ( MINE ) family ( Belghazi et al. , 2018 ; Poole et al. , 2018 ) . Despite the empirical and asymptotic bias improvements , MI estimation has not seen wide adoption . The challenges are two-fold . First , the analysis of dependencies among variables - let alone any MI analyses for scientific studies - requires not only an MI estimate , but also confidence intervals ( Holmes & Nemenman , 2019 ) around the estimate to quantify uncertainty and statistical significance . Existing MI estimators , however , do not provide confidence intervals . As low probability events may still carry a significant amount of information , the MI estimates could vary greatly given additional observations ( Poole et al. , 2018 ) . Towards providing upper and lower bounds of true MI under limited number of observations , existing MI lower bound techniques assume infinite data and would need further relaxations when a limited number of observations are provided . Closest to our work , Belghazi et al . ( 2018 ) studied the lower bound of the MINE estimator under limited data , but it involves bounds on generalization error of the signal model and would not yield useful confidence intervals for realistic datasets . Second , practical MI estimators should be insensitive to the choice of hyperparameters . An estimator should return a single MI estimate with its confidence interval irrespective of the type of the data and the number of observations . For learning-based approaches , this means that the model design and optimization hyperparameters need to not only be determined automatically but also taken into account when computing the confidence interval . Towards addressing these challenges , our estimator , DEMINE , introduces a predictive MI lower bound for limited samples that enables statistical dependency testing under practical dataset sizes . Our estimator builds on top of the MINE estimator family , but performs cross-validation to remove the need to bound generalization error . This yields a much tighter lower bound agnostic to hyperparameter search . We automatically selected hyperparameters through hyperparameter search , and a new cross-validation meta-learning approach is developed , based upon few-shot meta-learning , to automatically decide initialization of model parameters . Meta-overfitting is strongly controlled through task augmentation , a new task generation approach for meta-learning . With these improvements , we show that DEMINE enables practical statistical testing of dependency for not only synthetic datasets but also for real world functional Magnetic Resonance Imaging ( fMRI ) data analysis capturing nonlinear and higher-order brain-to-brain coupling . Our contributions are summarized as follows : 1 ) A data-efficient Mutual Information Neural Estimator ( DEMINE ) for statistical dependency testing ; 2 ) A new formulation of meta-learning using Task Augmentation ( Meta-DEMINE ) ; 3 ) Application to real life , data-scarce applications ( fMRI ) . 2 RELATED WORK . 2.1 MI ESTIMATION . A widely used approach for estimating MI from samples is using k-NN estimates , notably the KSG estimator ( Kraskov et al. , 2004 ) . Gao et al . ( 2017 ) provided a comprehensive review and studied the consistency and of asymptotic confidence bound of the KSG estimator ( Gao et al. , 2018 ) . MI estimation can also be achieved by estimating individual entropy terms through kernel density estimation ( Ahmad & Lin , 1976 ) or cross-entropy ( McAllester & Statos , 2018 ) . Despite their good performance on random variables with few dimensions , MI estimation on high-dimensional random variables remains challenging for commonly used Gaussian kernels . Fundamentally , estimating MI requires accurately modeling the random variables , where high-capacity neural networks have shown excellent performance on complex high-dimensional signals such as text , image and audio . Recent works on MI estimation have focused on developing tight asymptotic variational MI lower bounds where neural networks are used for signal modeling . The IM algorithm ( Agakov , 2004 ) introduces a variational MI lower bound , where a neural network q ( z|x ) is learned as a variational approximation to the conditional distribution P ( Z|X ) . The IM algorithm requires the entropy , H ( Z ) , andEXZ log q ( z|x ) to be tractable , which applies to latent codes of Variational Autoencoders ( VAEs ) and Generative Adversarial Networks ( GANs ) as well as categorical variables . Belghazi et al . ( 2018 ) introduces MI lower bounds MINE and MINE-f which allow the modeling of general random variables and shows improved accuracy for high-dimensional random variables , with application to improving generative models . Poole et al . ( 2018 ) introduces a spectrum of energy-based MI estimators based on MINE and MINE-f lower bounds and a new TCPC estimator inspired by Contrastive Predictive Coding ( Oord et al. , 2018 ) for the case when multiple samples from P ( Z|X ) can be drawn . Our work introduces cross-validation to the MINE-f estimator . We derive the lower bound of MINE-f under limited number of samples , and introduce meta-learning and hyperparameter search to enable practical statistical dependency testing . 2.2 GENERAL STATISTICAL DEPENDENCY TESTING . Existing works in general statistical dependency testing ( Bach & Jordan , 2002 ; Gretton et al. , 2005a ; Berrett & Samworth , 2019 ) have developed non-parametric independent criterions based on correlation and mutual information estimators equivalent to testing I ( X ; Z ) = 0 , followed by detailed bias and variance analyses . Our approach for independent testing suggest a different direction by harnessing the generalization power of neural networks and may improve test performance on complex signals . The p-values provided by our test do not involve approximated distributions and hold for small number of examples and arbitrary number of signal dimensions . As different statistical dependency testing approaches have explicit or implicit assumptions and biases that make them suitable in different situations , a fair comparison across different approaches is a challenging task . Instead , we focus on a self-contained presentation of our dependency test , and provide preliminary comparisons with a widely studied Hilbert-Schmidt independence criterion ( HSIC ) ( Gretton et al. , 2005a ) in the appendix . 2.3 META LEARNING . Meta-learning , or “ learning to learn ” , seeks to improve the generalization capability of neural networks by searching for better hyperparameters ( Maclaurin et al. , 2015 ) , network architectures ( Pham et al. , 2018 ) , initialization ( Finn et al. , 2017a ; 2018 ; Kim et al. , 2018 ) and distance metrics ( Vinyals et al. , 2016 ; Snell et al. , 2017 ) . Meta-learning approaches have shown significant performance improvements in applications such as automatic neural architecture search ( Pham et al. , 2018 ) , few-shot image recognition ( Finn et al. , 2017a ) and imitation learning ( Finn et al. , 2017b ) . In particular , our estimator benefits from the Model-Agnostic Meta-Learning ( MAML ) ( Finn et al. , 2017a ) framework which is designed to improve few-shot learning performance . A network initialization is learned to maximize its performance when fine-tuned on few-shot learning tasks . Applications include few-shot image classification and navigation . We leverage the model-agnostic nature of MAML for MI estimation between generic random variable and adopt MAML for maximizing MI lower bounds . To construct a collection of diverse tasks for MAML learning from limited samples , inspired by MI ’ s invariance to invertible transformations , we propose a task-augmentation protocol to automatically construct tasks by sampling random transformations to transform the samples . Results show reduced overfitting and improved generalization . 3 BACKGROUND . In this section , we will provide the background necessary to understand our approach1 . We define X and Z to be two random variables , P ( X , Z ) is the joint distribution , and P ( X ) and P ( Z ) are the marginal distributions over X and Z respectively . Our goal is to estimate MI , I ( X ; Z ) given independent and identically distributed ( i.i.d . ) sample pairs ( xi , zi ) , i = 1 , 2 . . . n from P ( X , Z ) . Let F = { Tθ ( x , z ) } θ∈Θ be a class of scalar functions , where θ is the set of model parameters . Let q ( x|z ) = p ( x ) e Tθ ( x , z ) E ( x , z ) ∼PXZ e Tθ ( x , z ) . Results from previous works ( Belghazi et al. , 2018 ; Poole et al. , 2018 ) show that the following energy-based family of lower bounds of MI hold for any θ : I ( X ; Z ) ≥ E ( x , z ) ∼PXZ log q ( x|z ) p ( x ) = E ( x , z ) ∼PXZTθ ( x , z ) − Ex∼PX logEz∼PZ e Tθ ( x , z ) , IEB1 ≥ E ( x , z ) ∼PXZTθ ( x , z ) − logEx∼PX , z∼PZ e Tθ ( x , z ) , IMINE ≥ E ( x , z ) ∼PXZTθ ( x , z ) − Ex∼PX , z∼PZ e Tθ ( x , z ) + 1 , IMINE-f , IEB ( 1 ) where , E is the expectation over the given distribution . Based on IMINE , the MINE estimator I ( X , Z ) ∧ n is defined as in Eq.2 . Estimators for IEB1 , IMINE-f and IEB can be defined similarly . I ( X ; Z ) ∧ n = sup θ∈Θ 1 n n∑ i=1 Tθ ( xi , zi ) − log 1 n2 n∑ i=1 n∑ j=1 eTθ ( xi , zj ) . ( 2 ) With infinite samples to approximate expectation , Eq.2 converges to the lower bound I ( X , Z ) ∧ ∞ = supθ∈Θ IMINE . Note that the number of samples n needs to be substantially more than the number of model parameters d = |θ| to guarantee that Tθ ( X , Y ) does not overfit to the samples ( xi , zi ) , i = 1 , 2 . . . n and overestimate MI . Formally , the sample complexity of MINE is defined as the minimum number of samples n in order to achieve Eq.3 , Pr ( |I ( X , Z ) ∧ n − I ( X , Z ) ∧ ∞| ≤ ) ≥ 1− δ . ( 3 ) Specifically , MINE proves that under the following assumptions : 1 ) Tθ ( X , Z ) is L-Lipschitz ; 2 ) Tθ ( X , Z ) ∈ [ −M , M ] , 3 ) { θi ∈ [ −K , K ] , ∀i ∈ 1 , . . . , d } , the sample complexity of MINE is given by Eq.4 . n ≥ 2M 2 ( d log ( 16KL √ d/ ) + 2dM + log ( 2/δ ) ) 2 . ( 4 ) For example , a neural network with dimension d = 10 , 000 , M = 1 , K = 0.1 and L = 1 , achieving a confidence interval of = 0.1 with 95 % confidence ( δ = 0.05 ) would require n ≥ 18 , 756 , 256 samples . This is achievable for synthetic example generated by GANs like that studied in Belghazi et al . ( 2018 ) . For real data , however , the cost of data acquisition for reaching statistically significant 1We follow the same notations in Belghazi et al . ( 2018 ) . We encourage the review of Belghazi et al . ( 2018 ) ; Poole et al . ( 2018 ) for a detailed understanding of IMINE , IEB1 , and IEB . estimation can be prohibitively expensive . Our approach instead uses the MI lower bounds specified in Eq.1 from a prediction perspective , inspired by cross-validation . Our estimator , DEMINE , improves sample complexity by disentangling data for lower bound estimation from data for learning a generalizable Tθ ( X , Z ) . DEMINE enables high-confidence MI estimation on small datasets .

Summary: This manuscript studies mutual-information estimation, in particular variational lower bounds, and focuses on reducing their sample complexity. The first contribution is based on adapting the MINE energy-based MI estimator family to out-of-sample testing. MINE involves fitting a very flexible parametric form of the distribution, such as a neural network, to the data to derive a mutual information lower bound. The present work separates the data fitting from the mutual information evaluation to decrease sample complexity, the argument being that the function class is no longer a limiting factor to sample complexity of the mutual information estimation. The second contribution uses meta learning to decrease the sample complexity required to fit the neural network, creating a family of tasks derived from the data with data transformation that do not modify the mutual information. The approaches are demonstrated on synthetic data as well as fMRI data, to detect significant inter-subject dependencies in time-series of neural responses.