Regularized Mutual Information Neural Estimation

1 INTRODUCTION . Identifying a relationship between two variables of interest is one of the great linchpins in mathematics , statistics , and machine learning ( Goodfellow et al. , 2014 ; Ren et al. , 2015 ; He et al. , 2016 ; Vaswani et al. , 2017 ) . Not surprisingly , this problem is closely tied to measuring the relationship between two variables . One of the fundamental approaches is information theory-based measurement , namely the estimation of mutual information ( MI ) . Recently , Belghazi et al . ( 2018 ) proposed a neural network-based MI estimator , which is called Mutual Information Neural Estimator ( MINE ) . Due to its differentiability and applicability , it motivated several researches such as various loss functions bridging the gap between latent variables and representations ( Chen et al. , 2016 ; Belghazi et al. , 2018 ; Oord et al. , 2018 ; Hjelm et al. , 2019 ) , and methodologies identifying the relationship between input , output and hidden variables ( Tishby & Zaslavsky , 2015 ; Shwartz-Ziv & Tishby , 2017 ; Saxe et al. , 2019 ) . Although many have shown the computational tractability and its usefulness , many intriguing questions about the MI estimator itself remain unanswered . • How does the neural network inside MINE behave when estimating MI ? • Why does MINE loss diverge in some cases ? Where does the instability originate from ? • Can we make a more stable estimate on small batch size settings ? • Why does the value of each term in MINE loss are shifting even after the estimated MI converges ? Are there any side effects of this phenomenon ? This study attempts to answer these questions by designing a synthetic dataset to interpret network outputs . Through keen observation , we dissect the Donsker-Varadhan representation ( DV ) one by one and conclude that the instability and the drifting are caused by the interrelationship between the stochastic gradient descent based optimization and the theoretical properties of DV . Based on these insights , we extend DV to draw out a novel lower bound for MI estimation , which mitigates the aforementioned problems , and circumvents the batch size limitation by maintaining the history of network outputs . We furthermore look into the L2 regularizer form of our bound in detail and analyze how various hyper-parameters impact the estimation of MI and its dynamics during the optimization process . Finally , we demonstrate that our method , called ReMINE , performs favorably against other existing estimators in multiple settings . 2 RELATED WORKS . Definition of Mutual Information The mutual information between two random variablesX and Y is defined as I ( X ; Y ) = DKL ( PXY ||PX ⊗ PY ) = EPXY [ log dPXY dPX⊗Y ] ( 1 ) where PXY and PX ⊗ PY are the joint and the marginal distribution , respectively . DKL is the Kullback-Leibler ( KL ) divergence . Without loss of generality , we consider PXY and PX ⊗ PY as being distributions on a compact domain Ω ⊂ Rd . Variational Mutual Information Estimation Recent works on MI estimation focus on training a neural network to represent a tight variational MI lower bound , where there are several types of representations . Although these methods are known to have statistical limitations ( McAllester & Stratos , 2018 ) , their versatility is widely employed nonetheless ( Hjelm et al. , 2019 ; Veličković et al. , 2018 ; Polykovskiy et al. , 2018 ; Ravanelli & Bengio , 2018 ; François-Lavet et al. , 2019 ) . One of the most commonly used is the Donsker-Varadhan representation , which is first used in Belghazi et al . ( 2018 ) to estimate MI through neural networks . Lemma 1 . ( Donsker-Varadhan representation ( DV ) ) I ( X ; Y ) = sup T : Ω→R EPXY [ T ] − log ( EPX⊗PY [ eT ] ) . ( 2 ) where both the expectations EPXY [ T ] and EPX⊗PY [ eT ] are finite . However , as the second term in Eq . ( 2 ) leads to biased gradient estimates with a limited number of samples , MINE uses exponential moving averages of mini-batches to alleviate this problem . To further improve the sampling efficiency of MINE , Lin et al . ( 2019 ) proposes DEMINE that partitions the samples into train and test sets . Other representations based on f-measures are also proposed by Nguyen et al . ( 2010 ) ; Nowozin et al . ( 2016 ) , which produce unbiased estimates and hence eliminating the need for additional techniques . Lemma 2 . ( Nguyen , Wainwright , and Jordan representation ( NWJ ) ) I ( X ; Y ) = sup T : Ω→R EPXY [ T ] − EPX⊗PY [ eT−1 ] , ( 3 ) where the bound is tight when T = log ( dP/dQ ) + 1 . Nevertheless , if MI is too big , estimators exhibit large bias or variation ( McAllester & Stratos , 2018 ; Song & Ermon , 2020 ) . To balance in between , Poole et al . ( 2019 ) design a new estimator Iα that interpolates Contrastive Predictive Coding ( Oord et al. , 2018 ) and NWJ . Yet , these methods concentrate on various stabilization techniques rather than revealing the dynamics inside the black box . In this paper , we focus on the DV representation and provide intuitive understandings of the inner mechanisms of neural network-based estimators . Based on the analysis , we introduce a new regularization term for MINE , which can effectively remedy its weaknesses theoretically and practically . 3 HOW DOES MINE ESTIMATE ? . Before going any further , we first observe the statistics network output in MINE during the optimization process using our novel synthetic dataset , and identify and analyze the following phenomena : • Drifting phenomenon ( Fig . 1a ) , where estimates of EPXY [ T ] and log ( EPX⊗PY [ eT ] ) drifts in parallel even after the MI estimate converges . • Exploding network outputs ( Fig . 1d ) , where smaller batch sizes cause the network outputs to explode , but training with larger batch size reduces the variance of MI estimates ( Fig . 2a ) . • Bimodal distribution of the outputs ( Fig . 2b ) , where the network not only classifies input samples but also clusters the network outputs as the MI estimate converges . Based on these observations , we analyze the inner workings of MINE , and understand how batch size affects MI estimation . 3.1 EXPERIMENT SETTINGS . Dataset . We designed a one-hot discrete dataset with uniform distribution U ( 1 , N ) to estimate I ( X ; X ) = logN with MINE , while easily discerning samples of joint distribution X , X from marginal distribution X ⊗ X . Additionally , we use one-hot representation to increase the input dimension , resulting in more network weights to train . In this paper , we used N = 16 . Network settings . We designed a simple statistics network T with a concatenated vector of dimension N × 2 = 32 as input . We pass the input through two fully connected layers with ReLU activation by widths : 32−256−1 . The last layer outputs a single scalar with no bias and activation . We used stochastic gradient descent ( SGD ) with learning rate 0.1 to optimize the statistics network . 3.2 OBSERVATIONS . We can observe the drifting phenomenon in Fig . 1a , where the statistics of the network output are adrift even after the convergence of MINE loss . The analysis for this phenomenon will be covered in more detail with theoretical results in Section 4 . This section will focus extensively on the relationship between batch size and logsumexp , and the classifying nature of MINE loss . Batch size limitation . MINE estimates in a batch-wise manner , i.e. , MINE uses samples inside a single batch when estimating EPXY [ T ] and log ( EPX⊗PY [ eT ] ) . Consider the empirical DV Î ( X ; Y ) = sup Tθ : Ω→R E ( n ) P̂XY [ Tθ ] − logE ( n ) ̂PX⊗PY [ e Tθ ] , ( 4 ) where E ( n ) P̂ is an empirical average associated with batch size n. Therefore , the variance of Î ( X ; Y ) increases as the batch size decreases . The observation in Fig . 2a is consistent with the batch size limitation problem ( McAllester & Stratos , 2018 ; Song & Ermon , 2020 ) , which shows that MINE must have a batch size proportional to the exponential of true MI to control the variance of the estimation . Exploding network outputs . We can understand the output explosion problem in detail by comparing Fig . 1b and Fig . 1d . During optimization , network outputs of joint samples get increased by the first term of Eq . ( 4 ) , where the inverse of the batch size is multiplied to the gradient of each network output . On the other hand , the output of marginal samples get decreased by the second term of Eq . ( 4 ) which concentrates the gradient to the maximum output . Note that the second term is dominated by the maximum network output due to logsumexp , which is a smooth approximation of max . As a single batch is sampled from the true underlying distribution , joint case samples may or may not exist . If it exists , then the joint sample output dominates the term , and its output gets de- creased accordingly , while other non-joint sample outputs also get slightly decreased . In summary , the second term acts as an occasional restriction for the increase of joint sample network outputs.1 The second term imposes a problem when the batch size is not large enough . With reduced sample size , joint samples that dominate the second term are getting rare . For the case where joint sample does not exist , marginal sample network outputs decrease much faster compared to the opposite case , and joint sample network outputs are more rarely restricted ; thus network outputs diverge in both directions ( Fig . 1d ) , and the second term vibrates between two extreme values depending on whether the joint case occurred ( Fig . 1c ) . This obviously leads to numerical instability and estimation failure . Bimodal distribution of the outputs . We furthermore observed network outputs directly , as both averaging terms of DV can inhibit the observation of how the statistics network acts on each sample . From the neural network viewpoint , whether each sample is realized from the joint or the marginal distribution is not distinguishable for the joint cases in marginal samples . Therefore , the statistics network has no means but to return the same output value , as it can be seen in Fig . 1b , indicating that the network can only separate joint and non-joint cases . This approach provides a clue that the network is solving a classification task , isolating joint samples from marginal samples , although the statistics network is only provided with samples from the joint and marginal distribution . We observed the distribution of network outputs in detail , on the case where only the marginal samples are fed to the statistics network in Fig . 2b . It stands to reason that the network outputs follow a particular distribution , as the network output estimates a log-likelihood ratio between joint and marginal distribution with an added constant ( Lemma 3 ) . Through this , we can view the estimated MI as a sample average ; hence Fig . 2a resembles the Gaussian noise by Central Limit Theorem ( CLT ) . Let us continue by concentrating on each network output . There is no distinction between the loglikelihood ratio of the samples in the same class for the one-hot discrete dataset : j for the joint case and j for the non-joint case . This explains the classifying nature of the statistics network , and there have to be exactly two clusters in Fig . 2b . Also , as j becomes −∞ , exp ( j ) nears 0 , and exp ( j ) is few magnitudes bigger than exp ( j ) ( see Fig . 2b ) . As mentioned above , few joint cases dominate the second term , so the second term becomes inherently noisier than the first term . Note that the effectiveness of conventional methods , such as applying exponential moving average to the second term ( Belghazi et al. , 2018 ) or clipping the network output values to restrict the magnitude of network outputs ( Song & Ermon , 2020 ) can also be understood with the analysis above . In addition , we can not interpret the network outputs directly as the log-likelihood ratio due to unregularized outputs , or the drifting problem . We will look into this fundamental limitation of MINE in more detail in the next section . 1Loosely speaking , the first term slowly increases a lot of joint samples network outputs , in contrary to the second term which quickly decreases a few joint sample network outputs .

Summary: This paper attempts to answer the four questions raised from the mutual information estimator. To this end, this paper investigates why the MINE succeeds or fails during the optimization on a synthetic dataset. Based on the observations and discussions, the paper then proposes a novel lower bound to regularize the neural networks and alleviate the problems of MINE.