Year: 2020

Abstract:
With the variational lower bound of mutual information (MI), the estimation of MI can be understood as an optimization task via stochastic gradient descent. In this work, we start by showing how Mutual Information Neural Estimator (MINE) searches for the optimal function T that maximizes the Donsker-Varadhan representation. With our synthetic dataset, we directly observe the neural network outputs during the optimization to investigate why MINE succeeds or fails: We discover the drifting phenomenon, where the constant term of T is shifting through the optimization process, and analyze the instability caused by the interaction between the logsumexp and the insufficient batch size. Next, through theoretical and experimental evidence, we propose a novel lower bound that effectively regularizes the neural network to alleviate the problems of MINE. We also introduce an averaging strategy that produces an unbiased estimate by utilizing multiple batches to mitigate the batch size limitation. Finally, we show that L regularization achieves significant improvements in both discrete and continuous settings.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORKS
text: 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.

heading: 3 HOW DOES MINE ESTIMATE?
text: 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.

heading: 3.1 EXPERIMENT SETTINGS
text: 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.

heading: 3.2 OBSERVATIONS
text: 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 cannot 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.

heading: 4 THE PROPOSED METHOD: REMINE
text: We look into the drifting problem of Fig. 1a in detail to introduce a novel loss with a regularizer that constrains the magnitude of network outputs and enables the use of network outputs from multiple batches. All proofs are provided in Appendix.
First, let us concentrate on the optimal function T ∗, which can be directly drawn from the DV representation. We start with the results of Belghazi et al. (2018), Lemma 3. (Optimal Function T ∗). Let P and Q be distributions on Ω ⊂ Rd. For some constant C∗ ∈ R, there exists an optimal function T ∗ = log dPdQ + C
∗ such that DKL(P||Q) = EP[T ∗] − log(EQ[eT ∗ ]).
Note that Mukherjee et al. (2020) directly utilize this fact to model the statistics network. Let us extend the result above, and show that C∗ can be any real number and T ∗ still be optimal. Lemma 4. (Family of Optimal Functions). For any constant C ∈ R, T = log dPdQ + C satisfies DKL(P||Q) = EP[T ]− log(EQ[eT ]).
This explains the drifting phenomenon we encountered in Fig. 1a. C∗ is drifting as there are no penalties on different C∗s. Drifting can be stopped by freezing the network weights (Lin et al., 2019), but is there a direct way to regularize C so that the neural network can concentrate on finding a single solution, rather than a family of solutions? Theorem 5. (ReMINE Loss Function) Let d be a distance function on R. For any constant C ′ ∈ R and function T : Ω→ R,
DKL(P||Q) = sup T :Ω→R EP[T ]− log(EQ[eT ])− d(log(EQ[eT ]), C ′),
Note that for the optimal T ∗, EPXY [T ∗] = I(X;Y ) + C ′ and log(EPX⊗Y [eT ∗ ]) = C ′. Based on Theorem 5, we propose a novel loss function by adding a new term d(log(EPX⊗PY [eT ]), C ′) that regularizes the drifting of C∗. The details of the ReMINE algorithm is as follows.
Algorithm 1: ReMINE θ ← Initialize network parameters, K ←Moving average window size, i← 0 repeat
Draw J samples from the joint distribution: (x(1)i , y (1) i ), · · · (x (J) i , y (J) i ) ∼ PXY Draw M samples from the marginal distribution: (x(1)i , y (1) i ), · · · (x (M) i , y (M) i ) ∼ PX ⊗ PY Evaluate the lower-bound: ÊPXY ← 1J ∑J j=1 Tθ(x (j) i , y (j) i ) , ÊPX⊗PY ← log( 1M ∑M m=1 e Tθ(x (m) i ,y (m) i ))
ν(θ)← ÊPXY − ÊPX⊗PY − d(ÊPX⊗PY , C) Update the statistics network parameters: θ ← θ +∇θν(θ) Estimate MI based on the last window W = [max(0, i−K + 1),min(K, i)] of size K: Î(X;Y ) = 1 J·|W | ∑ w∈W ∑J j=1 Tθ(x (j) w , y(j)w )− log( 1M·|W | ∑ w∈W ∑M m=1 e Tθ(x (m) w ,y (m) w ))
Next iteration: i← i+ 1 until convergence;
ReMINE differs from other MINE-like methods that rely on single batch estimates, as ReMINE can utilize all the previous network outputs after the convergence of Î(X;Y ) to establish a final estimate. To demonstrate why these MINE-like methods cannot use the network outputs from multiple batches, we first describe two basic averaging strategies and then show that both methods produce a biased estimate when the statistics network T is drifting. Theorem 6. (Estimation Bias caused by Drifting) The two averaging strategies below produce a biased MI estimate when the drifting problem occurs.
1. Macro-averaging (similar to that of Poole et al. (2019)): Establish a single estimate through the average of estimated MI from each batch.
2. Micro-averaging (our method): Calculate DV representation using the average of the each individual network outputs.

heading: 5 EXPERIMENTAL RESULTS OF L2 REGULARIZATION
text: Many choices exist for the distance function d when implementing ReMINE. Here, we choose d(x,C) = λ(x − C)2 and explain the rationale behind it. We also explore different choices of hyperparameters C and λ. We show that ReMINE blocks the drifting effect by restricting log(EPX⊗PY [eT ]) estimates to be C, and visualize the loss surface of ReMINE for different choices of λs. Finally, we show that ReMINE achieves better performance than state-of-the-art methods in the continuous domain, and shows comparable performance on self-consistency tests.2

heading: 5.1 EFFECTIVENESS OF L2 REGULARIZATION
text: For the sake of brevity, let batch size B, step size γ, f = E(n) P̂XY [Tθ], g = logE(n)̂PX⊗PY [e Tθ ], and C = 0. Then,∇Î(X;Y ) = ∇f −∇g− 2λg∇g = ∇f −∇g(1 + 2λg). As previously discussed, f increases each B joint sample outputs with (γ/B)∇Tθ. On the other hand, g is the approximation of the maximum marginal sample output, so the gradient becomes close to γ(1 + 2λg)∇Tθ for the maximum output. We can break down the dynamics of g into two parts, depending on whether joint samples exist in g. If not, maximum marginal sample output gets reduced with step size γ(1+2λg), which is, unlike MINE, adaptively adjusted by the size of g. Hence, ReMINE regularizes the maximum output to be centered around − 12λ , preventing the network outputs from diverging to −∞. If any joint sample exists, its network output will be big enough to dominate g. Step size is also γ(1 + 2λg), which regularizes the joint sample network outputs more strongly as it increases. This, too, helps to avoid the output explosion in the +∞ side, as shown in Fig. 3.
Impact of C. Fig. 4a shows the impact of changing C on MI estimation, with the same settings as Fig. 1a. We observed that the newly added regularizer penalizes the log(EPX⊗PY [eT ]) term to converge towards C as expected, without losing the ability to estimate MI.
Impact of λ. As we observed in Section 3.2, the network outputs of joint and non-joint cases converge to j and j, respectively. Using this, we visualize the effect of the proposed method by drawing the loss surface of MINE and ReMINE for the one-hot discrete dataset in Fig. 5.
LMINE = EPXY [T ]− log(EPX⊗PY [eT ]) = j − log(pej + (1− p)ej) (5) LReMINE = EPXY [T ]− log(EPX⊗PY [eT ])− d(log(EPX⊗PY [eT ]), C) (6)
= j − log(pej + (1− p)ej)− λ(log(pej + (1− p)ej)− C)2. (7)
We again observe the drifting phenomena, as the loss surface has a plateau of equally palatable solutions. Regularization term successfully warps the loss surface, so that it has a single solution. As λ increases, the loss surface becomes steeper, resulting in sporadic spikes for each gradient step.
2We release our code in

heading: 5.2 REMINE IN THE CONTINUOUS DOMAIN
text: 20-D Gaussian dataset. We sampled (x, y) from d-dimensional correlated Gaussian dataset where X ∼ N(0, Id) and Y ∼ N(ρX, (1 − ρ2)Id) given the correlation parameter 0 ≤ ρ < 1, which is taken from Belghazi et al. (2018) to test ReMINE on continuous random variables. The true MI for the dataset is I(X,Y ) = −d2 log(1− ρ 2). In these experiments, we set batch size to 64.
Network settings. We consider a joint architecture, which concatenates the inputs (x, y), and then passes through three fully connected layers with ReLU activation (excluding the output layer) by widths 40 − 256 − 256 − 1, same as the network used in Poole et al. (2019). We used Adam optimizer with learning rate 5× 10−4, β1 = 0.9 and β2 = 0.999.
Comparison to state-of-the-arts. As mentioned in Fig. 4b, we can remove the log(EPX⊗PY [eT ]) term by choosing C = 0. As discussed in Section 3.2, the second term is inherently noisy. Hence, using all the terms in ReMINE only in optimization and removing the second term in estimation can effectively reduce noise. We call this trick ReMINE-J.
To verify the quality of lower bounds, we compare ReMINE and ReMINE-J to InfoNCE (Oord et al., 2018), JS (Hjelm et al., 2019; Poole et al., 2019), MINE (Belghazi et al., 2018), NWJ (Nguyen et al., 2010), SMILE (Song & Ermon, 2020), SMILE+JS (which estimates with SMILE, and optimizes with JS), TUBA (Barber & Agakov, 2004; Poole et al., 2019) and Iα (Poole et al., 2019). To make a fair comparison, ReMINE also uses the macro-averaging strategy, the same as the other methods. Our methods show comparable or better estimation performance with less variance than others, as shown in Fig. 6. Exact values for bias, variance, and mean square error to the true MI for each estimator are shown in Appendix.

heading: 5.3 REMINE IN THE IMAGE DOMAIN
text: Experiment settings. Song & Ermon (2020) introduced the self-consistency tests on image datasets to verify whether the MI estimates follow the basic properties of MI. For our experiments, we used the same network and optimizer, types of tests, and the number of epochs of Song & Ermon (2020) with batch size 16. We did not use our micro-averaging strategy for a fair comparison.
Comparison on self-consistency tests. To compare the stability of DV-based estimators, we conducted self-consistency tests on ReMINE, MINE, SMILE, and SMILE+JS. For type 1, every estimator successfully returns values between the theoretical bound with an increasing trend. For type 2, only ReMINE and SMILE+JS estimates are close to the ideal value. For type 3, none of the estimators worked well. However, ReMINE shows smaller variance compared to MINE and has similar stability to SMILE+JS.

heading: 6 CONCLUSION
text: In this paper, we studied how the neural network inside MINE handles the MI estimation problem. We delved into the drifting problem, where two terms of DV continue to fluctuate together even after the MI estimate converges, and the explosion problem, where the network outputs become unstable due to properties of the second term in DV when batch size is small. Based on the analysis, we penalized the objective function for obtaining a unique solution by using L2 regularization. Despite the simplicity, the proposed loss and the micro-averaging strategy mitigate drifting, exploding, and batch size limitation problems. Further, ReMINE enables us to directly interpret the network output values as the log-likelihood ratio of joint and marginal distribution probability and performs favorably against state-of-the-art methods. However, further investigation needs to be done on the impact of optimizers on the batch size limitation, and why DV-based estimators fail in some of the self consistency tests.

heading: 7 APPENDIX: PROOFS

heading: 7.1 PROOF OF FAMILY OF OPTIMAL FUNCTIONS
text: Theorem. For any constant C ∈ R, T = log dPdQ +C satisfiesDKL(P||Q) = EP[T ]− log(EQ[e T ]).
Proof. Suppose that T = log dPdQ +C. We can write the function T = (T ∗ −C∗) +C by Lemma 3 in the manuscript. Therefore,
EP[T ] = EP[T ∗ − C∗ + C] = EP[T ∗]− C∗ + C,
and
log(EQ[eT ]) = log(EQ[eT ∗−C∗+C ])
= log(eC−C ∗ EQ[eT ∗ ]) = (C − C∗) + log(EQ[eT ∗ ]).
Since EP[T ]− log(EQ[eT ]) = EP[T ∗]− log(EQ[eT ∗ ]), the function T also optimal.

heading: 7.2 PROOF OF REMINE LOSS FUNCTION
text: Theorem. Let d be a distance function on R. For any constant C ′ ∈ R and function T : Ω→ R,
DKL(P||Q) = sup T :Ω→R EP[T ]− log(EQ[eT ])− d(log(EQ[eT ]), C ′),
Proof. i) For any T ,
EP[T ]− log(EQ[eT ])− d(log(EQ[eT ]), C ′) ≤ EP[T ]− log(EQ[eT ]).
Therefore, supT :Ω→R EP[T ]− log(EQ[eT ])− d(log(EQ[eT ]), C ′) ≤ DKL(P||Q).
ii) By the theorem above, there exists T ∗ = log dPdQ + C ′ such that
DKL(P||Q) = EP[T ∗]− log(EQ[eT ∗ ])
and
log(EQ[eT ∗ ]) = log(EQ[eC ′ dP dQ ]) = log( ∫ eC ′ dP dQ dQ) = C ′.
Therefore,
sup T :Ω→R
EP[T ]− log(EQ[eT ])− d(log(EQ[eT ]), C ′) ≥ EP[T ∗]− log(EQ[eT ∗ ])− d(log(EQ[eT ∗ ]), C ′)
= DKL(P||Q)
Combining i) and ii) finishes the proof.

heading: 7.3 PROOF OF ESTIMATION BIAS CAUSED BY DRIFTING
text: Theorem. (Estimation Bias caused by Drifting) The two averaging strategies below produce a biased MI estimate when the drifting problem occurs.
1. Macro-averaging (similar to that of Poole et al. (2019)): Establish a single estimate through the average of estimated MI from each batch.
2. Micro-averaging (our method): Calculate DV representation using the average of the each individual network outputs.
Proof. Let the outputs of ith batch, jth sample inside the batch as T (J)ij , T (M) ij , joint and marginal case respectively, and the output without drifting as T ∗ij , and drifting constant for each batch Ci. Then, Tij = T ∗ij + Ci.
When the number of batch is B and each batch size is N ,
1. Macro averaging:
1 B Σi[ 1 N ΣjT (J) ij − log( 1 N Σj expT (M) ij )] (8)
= 1
B Σi[
1
N Σj(T
(J∗) ij + Ci)− log(
1
N Σj exp(T
(M∗) ij + Ci))] (9)
= 1
B Σi[
1
N Σj(T
(J∗) ij + Ci)− log(
1
N expCiΣj expT
(M∗) ij )] (10)
= 1
B Σi[
1
N ΣjT
(J∗) ij − log(exp(−Ci)
1
N expCiΣj expT
(M∗) ij )] (11)
= 1
B Σi[
1
N ΣjT
(J∗) ij − log(
1
N Σj expT
(M∗) ij )] (12)
= 1
NB ΣijT
(J∗) ij −
1 B Σi[log( 1 N Σj expT (M∗) ij )] (13)
6= 1 NB ΣijT (J∗) ij − log( 1 NB Σij expT (M∗) ij ) (14)
2. Micro averaging:
1
NB ΣijT
(J) ij − log(
1
NB Σij expT
(M) ij ) (15)
= 1
NB Σij(T
(J∗) ij + Ci)− log(
1
NB Σij exp (T
(M∗) ij + Ci))) (16)
= 1
NB ΣijT
(J∗) ij − log[(
1
NB Σij exp (T
(M∗) ij + Ci)) 1 BΣiCi ] (17)
6= 1 NB ΣijT (J∗) ij − log( 1 NB Σij expT (M∗) ij ) (18)

heading: 8 APPENDIX: ADDITIONAL EXPLANATIONS
text: Additional explanations for Fig. 4b As the N -dimensional one-hot discrete dataset is uniform, we can easily calculate the likelihood ratio of joint and non-joint case samples. For all the possible samples, PX ⊗ PY = 1/N2, as they are total of N2. Also, for joint case samples, PXY = 1/N , and PXY = 0 for non-joint case samples. Hence, the likelihood ratio for the joint cases is N , and non-joint cases is 0. These are consistent with the experimental results, where j converges to logN , and j keeps decreasing. Nonetheless, as exp(j) gets closer to zero, the second term of ReMINE loss has lesser influence; hence the decreasing speed of j gets slowed down to a halt as it reaches − 12λ = −5.
We can explain the same result from the perspective of j and j. As we observed in Section 3.2, the network output values of joint and non-joint cases converge to j and j, respectively. Since the dataset is uniform, the probability p of joint cases appearing from the marginal samples is 1N . Therefore, we can analyze the value of j and j after convergence as follows: as iteration i→∞,
EPXY [T (i)] = j → I(X;Y ) + C = logN + C (19)
log(EPX⊗PY [eT (i)]) = log(pej + (1− p)ej) = log( 1
N ej + N − 1 N ej)→ C (20)
where T (i) is the statistics network at iteration i. We combine Eq. (19) and Eq. (20) to
1
N elogN+C + N − 1 N ej → eC , and (21)
ej → 0. (22)
In summary, j will converge to logN +C, and ej to 0, as shown in Fig. 4b. Note that j and j serves as a back-of-the-envelope calculation for us to estimate network outputs easily on discrete settings.
What happens if the batch size is small? When the batch size is 1 and C = 0, the loss function of ReMINE changes its characteristics as follows.
• Joint case occurs. As the samples are indistinguishable,
L = j − log ej − d(j, 0) = −λj2, (23)
which is maximized when j = 0. • Non-joint case occurs.
L = j − log ej + d(log ej , 0) = j − j − λj2. (24)
The latter quadratic term of j is maximized when j = − 12λ .
If the statistics network succeeds to converge on both cases for our one-hot discrete dataset,
Î(X;Y ) = E(n) P̂XY [Tθ]− logE(n)̂PX⊗PY [e Tθ ] = 0− log(pe0 + (1− p)e− 12λ )→ − log p (25)
when λ→ +0. As p = 1N , Î(X;Y )→ logN . Intuitively, on smaller batch sizes, joint cases cannot occur in marginal samples, as mentioned in Section 3.2. Hence, EPXY [T ] and log(EPX⊗PY [eT ]) behave differently compared to the larger batch size. The regularizer term penalizes both terms in different ways. Joint cases in marginal samples can contribute only with Eq. (23), so EPXY [T ]→ 0. Moreover, as λ gets smaller, log(EPX⊗PY [eT ]) gets regularized less so that it can converge to −Î(X;Y ). In contrast, since MINE has no regularization term, namely λ = 0, there is no way for the joint case in marginal samples to influence T , hence failing to estimate MI as shown in Fig. 8b.
Impact of λwith batch size. We inspect the relationship between batch size and λ in detail. Fig. 8 shows that imposing regularization reduces noise on a large batch size domain. However, on a small batch size domain, log(EPX⊗PY [eT ]) cannot have nonzero value, hence failing to estimate MI value. The effect of the ReMINE loss in two different domains gets mixed in between.
Visualizing network outputs on 1-D Gaussian. The dataset forbids us to label joint and non-joint samples explicitly, so we visualized the network outputs on 2-D plane. We used the same experiment settings as Section 3.1, only changing the input dimension to 2. We can see in Fig. 9 that the network outputs of the overlapping region remain near 0, which indicates that the likelihood is equal between joint and marginal distribution. Other regions are separated by the sign of their outputs. Positive network output means that joint distribution is more probable than marginal distribution to sample that data point, and vice versa.

heading: 9 APPENDIX: ADDITIONAL EXPERIMENTAL DETAILS
text: Additional experiments outputs on 20-D Gaussian. For quantitative comparison with other approaches on the 20-D correlated Gaussian dataset, we show the bias, variance and mean squared error (MSE) of the neural network-based and nearest neighbor-based MI estimators in Fig. 10, Fig. 11 and Table 1. We omitted values which are more than 100 in Table 1. We additionally show results from KL (Jiao et al., 2018), Mixed KSG (Gao et al., 2017), CCMI (Mukherjee et al., 2020), TNCE (Oord et al., 2018; Poole et al., 2019) and ReMINE-L1 (our method with L1 regularization).
Both ReMINE and ReMINE-J shows comparable or better performance compared to other methods. Note that L1 regularizer also suffers from the explosion problem, as the gradient is not adaptively adjusted by the magnitude of network outputs, as discussed in Section 5.1.
Additional experiments on the self consistency test. We report the performance of other variational bound methods (JS, Iα, InfoNCE) in Fig. 12. Iα and JS often result in unstable MI estimates, as shown in the type 2 experiment. On the other hand, InfoNCE estimates MI quite reliably but also fails for type 3.
To observe on a different dataset setting, we also used MNIST (LeCun et al., 1998). As shown in Fig. 13, tests yielded similar results.
To observe on a different statistics network, we used modified ResNet18 (He et al., 2016) that outputs a single scalar. As shown in Fig. 14, SMILE has become more unstable, but there are no significant differences in other variational bounds. This experiment shows that the network size has a small impact on the validity of this test on CIFAR-10.
Effectiveness on the Conditional Mutual Information Estimation Task We compare the performance of various estimators on the conditional MI (CMI) estimation task. To set the baseline, we chose CCMI (Mukherjee et al., 2020), MINE (Belghazi et al., 2018), and KSG estimator (Kraskov et al., 2004). The Experiment is performed under the Model 1 setting in Mukherjee et al. (2020). We refer to the supplementary of Mukherjee et al. (2020) for hyper-parameter settings such as network structures and optimizer parameters. We only changed the objective function of MINE to test our method. As shown in Fig. 15, ReMINE can reach comparable performance without changing the form to classification loss. Also, ReMINE produces stable estimates compared to MINE.

Review summaries:
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.