Combating the Instability of Mutual Information-based Losses via Regularization
Abstract
Notable progress has been made in numerous fields of machine learning based on neural networkdriven mutual information (MI) bounds.
However, utilizing the conventional MI-based losses is often challenging due to their practical and mathematical limitations.
In this work, we first identify the symptoms behind their instability: (1) the neural network not converging even after the loss seemed to converge, and (2) saturating neural network outputs causing the loss to diverge.
We mitigate both issues by adding a novel regularization term to the existing losses.
We theoretically and experimentally demonstrate that added regularization stabilizes training.
Finally, we present a novel benchmark that evaluates MI-based losses on both the MI estimation power and its capability on the downstream tasks, closely following the pre-existing supervised and contrastive learning settings.
We evaluate six different MI-based losses and their regularized counterparts on multiple benchmarks to show that our approach is simple yet effective.
INTRODUCTION
Identifying a relationship between two variables of interest is one of the key problems in mathematics, statistics, and machine learning [Goodfellow et al., 2014, Ren et al., 2015, He et al., 2016, Vaswani et al., 2017].
One of the fundamental approaches is information theory-based measurement, namely the measure of mutual information (MI).
Due to its mathematical soundness and the rise of deep learning, many have designed differentiable MI-based losses for neural networks.
Some utilize the MI-based losses to bridge the gap between latent variables and representations in generative adversarial networks [Nowozin et al., 2016, Chen et al., 2016, Belghazi et al., 2018, van den Oord et al., 2018, Hjelm et al., 2019], where others introduce MI-based methodologies identifying the relationship between input, output, and hidden variables [Tishby and Zaslavsky, 2015, Shwartz-Ziv and Tishby, 2017, Saxe et al., 2018].
Furthermore, recent self-supervised losses use contrastive losses, where its origin can be traced back to MI-based losses [Cheng et al., 2020, Hénaff, 2020, Chuang et al., 2020].
Although many have shown computational tractability and usefulness of MI-based losses, others still struggle with their instability during optimization.
Contrastive learning literature with MI-based losses such as Chen et al.
He et al.
[2020] use huge batch sizes to reduce the variance of losses.
Bardes et al.
[2021] adds a regularization term to the neural network embeddings to stabilize the training.
McAllester and Stratos [2020] and Song and Ermon [2020] further provide theoretical limitations of variational
MI estimators, arguing that the limited batch size induces
a
MI estimation variance too large to handle.
We argue that mitigating the variance of MI-based losses is critical for stabilizing training, where it is well known that more stable optimization of neural networks yields better predictive performance on the downstream tasks [Rothfuss et al., 2019,
Bear and Cushman, 2020, Chavdarova et al., 2019, Richter et al., 2020, Zeng et al., 2020, Colombo et al., 2021].
In this paper, we concentrate on identifying the cause behind the instability of MI-based losses and propose a simple yet effective regularization method that can be applied to various MI-based losses.
We start by analyzing the behaviors of two MI estimators; the MI Neural Estimator (MINE) loss [Belghazi et al., 2018] and Nguyen-Wainwright-Jordan loss (NWJ) loss [Nguyen et al., 2010].
We identify two distinctive behaviors that induce instability during training, drifting and exploding neural network outputs.
Based on these observations, we design two novel dual representations of the KL-divergence called Regularized Donsker-Varadhan representation (ReDV) and Regularized NWJ representation (ReNWJ)
We show theoretically and experimentally that adding our regularizer term suppresses two behaviors of drifting and exploding, avoiding instability during training.
Finally, we design a novel benchmark that bridges the gap between variational MI estimators and real-world tasks, whereas previous works either do not directly show the MI estimation performance or evaluate only on toy problems.
We reformulate both the supervised and the contrastive learning problem [Chen et al., 2020, He et al., 2020, Khosla et al., 2020] as MI estimation problems and show that our regularization yields better performance on both perspectives, downstream task and MI estimation performance.
BACKGROUND & RELATED WORKS
Definition of MI The mutual information between two random variables X and Y is defined as
[Equation]
where P XY and P X ⊗ P Y are the joint distribution and the product of the marginal distributions, respectively.
D KL is the Kullback-Leibler (KL) divergence.
Without loss of generality, we consider P XY and P X ⊗ P Y as being distributions on a compact domain Ω ⊂ R d .
MI through dual representation of D KL
We first introduce two dual representations of D KL , as MI is defined using
D KL .
The most widely known is
the Donsker-Varadhan representation D DV
[Donsker and Varadhan, 1975].
For given two distribution P and Q on some compact domain Ω ⊂ R d , D
[Equation]
where both the expectations E P ( T ) and E Q ( e T ) are finite.
If we substitute P and Q into P XY and P X ⊗ P Y , D DV yields the definition of MI.
The optimal T ∗ = log d P d Q + C , where
C ∈ R can be any constant.
In contrast to D DV , the Nguyen-Wainwright-Jordan representation D NWJ [Nguyen et al., 2010] is induced by the convex conjugate known as Fenchel’s inequality [HiriartUrruty and Lemaréchal, 2004]:
[Equation]
The optimal T ∗ = log d P d Q + 1 is unique unlike the optimal
T ∗ of D DV due to its self-normalizing property [Belghazi et al., 2018].
However, D DV guarantees tighter lower bounds than D NWJ [Ruderman et al., 2012, Polyanskiy and Wu, 2014].
These two representations provide the theoretical soundness for numerous variational MI bounds.
Variational MI estimation With the increasing success of neural networks, several neural network-driven variational bounds of MI are proposed.
They are widely employed, such as contrastive learning [van den Oord et al., 2018,
He et al., 2020, Chen et al., 2020] or generative adversarial training [Belghazi et al., 2018, Nowozin et al., 2016].
Variational bounds of MI commonly focus on estimating T ∗ via a neural network T θ : Ω → R , called the statistics network [Belghazi et al., 2018], which outputs a single real value given the input sample pairs.
I MINE [Belghazi et al., 2018] directly maximize D DV as the objective function by feeding the samples ( x, y ) of P XY and P X ⊗ P Y into T θ :
[Equation]
where P ( n ) is the empirical distribution associated to n i.i.d.
samples for given distribution P .
Belghazi et al.
[2018] also utilizes moving averages of mini-batches to reduce the MI estimation variance caused by the limited batch size.
I InfoNCE [van den Oord et al., 2018] is also commonly used due to its stability and decent performance:
[Equation]
where the N samples ( x i , y i )
Ni =1 are drawn from P XY , which becomes equivalent to using the Softmax function with the negative log loss.
I InfoNCE is also equivalent to I MINE up to a constant, but upper bounded by log N , hence not able to estimate large MI values [van den Oord et al., 2018].
Poole et al.
[2019] introduced I TUBA , a unified lower bound, by expanding D NWJ [Barber and Agakov, 2003, Nguyen et al., 2010].
[Equation]
where a ( y ) is the variational parameter.
However, unlike
I MINE or I InfoNCE , directly using the exponential term often causes numerical instability.
Even if T θ outputs a moderately sized value, e T θ can easily exceed the floating-point range.
To avoid this problem, Poole et al.
[2019] introduce D NWJbased lower bound I JS by using a softplus-activated neural network as T θ ,
[Equation]
Variance problem of MI estimators Despite the variety of bounds proposed, many still suffer from the bias-variance trade-off [Poole et al., 2019].
McAllester and Stratos [2020] and Song and Ermon [2020] prove that the I MINE estimator must have a batch size proportional to the exponential of true MI to control the variance of the estimation.
Many bounds try to mitigate this problem by reducing the variance of low-biased estimators, such as by interpolating with a low variance bound [Poole et al., 2019] or dropping the formal theoretical guarantees [McAllester and Stratos, 2020].
Song and Ermon [2020] proposed I SMILE to clip the range of T θ trained with I MINE , sacrificing the estimation quality to reduce the variance.
[Equation]
where clip ( v, l, u ) =
max(min
( v, u ) , l ) for v, u, l ∈ R .
Practical usages of MI MI-based losses are often applied in generative modeling, such as for better mode coverage [Belghazi et al., 2018] or learning disentangled representations without supervision [Chen et al., 2016, Ojha et al., 2020, Li et al., 2021b, Jeon et al., 2021].
Representation learning employs MI-based losses [Tian et al., 2020b, Hjelm et al., 2019, Tschannen et al., 2020, Cheng et al., 2020, Wu et al., 2020, Wen et al., 2020, Boudiaf et al., 2020, Tian et al., 2020a, Li et al., 2021a] to yield feature extractors that reflect its downstream tasks well.
We emphasize that these approaches can be further utilized to measure the performance of MI estimators.
Comparing between MI estimators Toy datasets such as correlated multivariate Gaussian distributions has been widely accepted for the evaluation of MI estimation [Belghazi et al., 2018, Poole et al., 2019, Song and Ermon, 2020,
Cheng et al., 2020, Lin et al., 2019].
However, we emphasize that using synthetic data as a definitive benchmark will end up in a disparity with real-world tasks.
There have been some approaches that compared different MI estimators on generative modeling [Belghazi et al., 2018, Hjelm et al., 2019] or representational learning [Tian et al., 2020b].
However, finding the ideal MI for each downstream task is not trivial, making it impossible to directly assess the MI estimation quality.
Moreover, Tschannen et al.
[2020] and Tian et al.
[2020b] showed the gap between MI estimation quality and downstream performance on specific tasks.
Hence, it is crucial to evaluate both perspectives.
The closest work to our benchmark is the consistency test of Song and Ermon [2020] using CIFAR-10 [Krizhevsky, 2009] and MNIST [LeCun et al., 1998].
However, the test only offered to assess the ratio of two separate MI estimations, making it difficult to separately measure the quality of each estimation. [Figure] [Figure]
100 for 3000 iterations.
{+We observe the statistics network outputs T θ ( x 1 , x 2 ) , where we split the outputs into two: ( x 1 , x 
+}
) ∈ Supp ( P XX ) and ( x ′ 1 , x ′ 2 ) ∈ Supp ( P X ⊗ P X ) \ Supp ( P XX ) .  INSTABILITY OF MI BOUNDS
To demonstrate and analyze the instability of variational
MI bounds, we design a synthetic problem with the Onehot dataset.
We then solve the task via I MINE and I NWJ , which are the losses derived from the two most commonly used representations of KL-divergence, D DV and D NWJ , respectively.
Both losses consist of two terms, each derived from the statistics of joint distribution E P XY and the product of marginal distributions E P X ⊗ P Y .
Hence, to observe the behavior of each loss during training, we plot the two terms separately.
Also, to observe how each distribution differ by the statistics network outputs T θ ( x, y ) , we plot each output from ( x, y ) ∼ Supp ( P XY ) and ( x, y ) ∼
Supp ( P X ⊗ P Y ) \ Supp ( P XY ) , where we denote the support of P as Supp ( P ) .
Support is the set of values that the random variable can take [Taboga, 2021].
One-hot Dataset We design a one-hot discrete problem with uniform distribution X ∼ U (1 , N ) to estimating
I ( X, X ) = log N for a given integer N .
This task is intentionally created to easily discern samples ( x, x ) ∼ P XX from ( x, x ) ∼ P X ⊗ P X , so that we can directly observe its network outputs T θ ( x, x ) . [Figure]
Seemingly Stable Case We first observe the behaviors of the statistics network T θ when the losses are seemingly stable, producing a successful MI estimate.
Fig.
1 shows the MI estimates and the two terms that construct each MI estimate per batch.
We observe that the first and the second term estimates of I MINE , unlike I NWJ , drifting in parallel even after the MI estimate converge.
This is due to the free constant term C in the optimal T ∗ of D DV , where the self-normalizing D NWJ avoids this problem.
This drifting phenomenon implies that T θ is not stable even after the loss seems to be converged, as shown in Fig.
Also, the plot demonstrates how T θ is trained; it isolates the samples ( x, y ) ∼ P XY from the samples ( x, y ) ∼ P X ⊗ P Y .
Unstable Case We also demonstrate the behaviors of
T θ when the losses get unstable in Fig.
We reduce the batch size to make the optimization unstable, where this behavior is often reported in multiple works [van den Oord et al., 2018, He et al., 2020, Chen et al., 2020].
However, even though the losses seem unstable, T θ successfully discerns the samples before the outputs explode.
We believe that this is because of how T θ is optimized during training.
The statistics network outputs T θ ( x 1 , x 2 ) of ( x 1 , x {+
+} ) ∈ Supp ( P XX ) gets increased by the first term but occasionally decreased by the second term.
However,
T θ ( x ′ 1 , x ′ 2 ) of ( x ′ 1 , x ′ 2 ) ∈ Supp ( P X ⊗ P X ) \ Supp ( P XX ) gets decreased whatsoever, as ( x ′ 1 , x ′ 2 ) is used only for the second term.
This makes the second term more unstable and motivates us to regularize it for better numerical stability during optimization.
To summarize, we suspect the instability of variational bounds comes from two reasons.
Firstly, the statistics net- [Figure] work did not converge even after the loss seemingly converged.
We argue that this is due to the unnormalized constant term in the optimal T ∗ of D DV , where D NWJ successfully avoids via self-normalization.
Secondly, the loss gets unstable as T θ ( x ′ 1 , x ′ 2 ) endlessly decrease due to the second term.
This observation is also consistent with the theoretical findings of Song and Ermon [2020], McAllester and Stratos [2020], where they show that large variance of the second term leads to failed MI estimation.
We claim that the outputs have to be regularized in some form to avoid the instability.
STABILIZING THE MI BOUNDS
In this section, we introduce two novel regularized representations and its corresponding losses to tackle the instability during optimization.
We show both theoretically and experimentally that adding regularization mitigates the unstable behavior of the statistics network T θ .
We also describe a simple windowing method that can sidestep the batch size limitation problem of the MI estimation problem.
We defer all the proofs to the Appendix.
Regularized representations We stabilize the two existing representations D DV and D NWJ by regularizing the second term.
We introduce two novel representations: Regularized DV ( D ReDV ) and Regularized NWJ ( D ReNWJ ),
[Equation]
where C ∗ ∈ R is any constant and d ( ∗ , ∗ ) is a distance function on R .
Theorem 1.
D ReDV and D ReNWJ is a dual representation for
D KL such that
[Equation]
We emphasize that both representations are not MI-specific but dual representations of D KL , which can be easily extended to numerous variational MI bounds based on D DV and D NWJ .
Especially, the newly added regularizer grants
D ReDV the normalizing property, effectively solving the drifting problem of D DV .
Regularizing I MINE and I NWJ
Based on D ReDV and
D ReNWJ , we propose a novel neural network-driven variational MI bound I ReMINE
and I ReNWJ by choosing the Euclidean distance d ( x, y ) =
( x − y ) 2 and the log-Euclidean distance d ( x, y ) =
(log x − log y ) 2 , respectively.
[Equation]
where C ∗ ∈ R is any constant and λ is a hyperparameter that controls the degree of regularization.
We can also easily regularize other losses such as I InfoNCE , I SMILE , I TUBA , and
I JS in a plug-and-play manner.
See Table 1 for more details on its regularized counterparts.
Solving the drifting problem Due to the self-regularizing nature of D NWJ , we must fix C ∗ = 1 for I ReNWJ .
We also set
C ∗ = 0 for I ReMINE on future experiments, but to demonstrate the ability of the regularizer term to stop the drifting, we experiment with various C ∗ in Fig. 4.
Comparing to
I MINE in Fig.
1, we can observe that I ReMINE successfully solves the drifting problem by regularizing the second term to have a single value.
Solving the explosion problem We previously observed the instability of I MINE and I NWJ when using a small batch in Fig.
We apply the same setting to I ReMINE
and I ReNWJ to observe if the regularizer mitigates the instability problem.
Both regularized losses successfully avoid the explosion problem and limit the statistics network outputs T θ ( x 1 , x 2 ) within a certain boundary.
As discussed in Section 3, the second term was the culprit of the variance in MI estimation.
The newly added term directly regularizes it to stabilize training, giving the statistics network T θ additional hints for the second term to converge to a specific value C ∗ successfully.
Furthermore, we empirically found that our regularization works well with I SMILE ’s strategy of clipping
T θ .
Gradient zeros out for the original I SMILE
if T θ ( x, y ) exceeds a certain threshold.
This behavior makes T θ act as if it were frozen, failing to further optimize during training.
However, with the regularizer term, we can clip T θ ( x, y ) [Figure] counterparts, I ReMINE and I ReNWJ , with the same small batch settings from
Fig. Regularization effectively mitigates both instability symptoms, shifting and exploding. only on first and second term, i.e., on the original loss. Now, clipping filters out the noisy gradients while the gradients calculated from the regularizer avoid freezing T θ entirely.
Mathematical properties of I ReMINE and I ReNWJ
Following Belghazi et al.
[2018], we show the soundness of
I ReMINE and I ReNWJ in two perspectives, strong consistency and sample complexity.
These properties relate to whether the trained T θ can be sufficiently similar to the optimal T ∗ .
Theorem 2.
I ReMINE
and I ReNWJ are strongly consistent.
For the two losses, we also provide the mathematical bound on the number of samples required for the empirical MI estimation at a given accuracy and with high confidence.
Similar to Belghazi et al.
[2018], let T θ satisfy L -Lipschitz with respect to the parameter θ such that | θ | < K and d is dimension of the parameter space of T θ .
Theorem 3.
Assume that T θ is bounded above by M .
Let k be the number of sample means.
Given any ϵ, δ of the desired accuracy and confidence parameters, we have
[Equation]
whenever the number n of samples satisfies
[Equation]
Theorem 4.
Assume that 1 ≤ | T θ | < M and d ( x, 1) ≤ | x − 1 | .
Let k be the number of sample means.
Given any ϵ, δ of the desired accuracy and confidence parameters, we have
[Equation]
whenever the number n of samples satisfies
[Equation]
Drifting may lead to noisy MI estimate We prove that the variance of the second term on the empirical distributions is affected by the constant term C ∗ .
Theorem 5.
Let Q ( n ) be the empirical distributions of n i.i.d.
samples from Q .
For the optimal T 1 = log dpdq + Cand T 2 = log dpdq + C 2 where C 1 ≥ C 2 ,
[Equation]
This implies that unregulated C ∗ may lead to worse MI estimation quality, as the source of the estimate variance are mainly due to the second term.
Increasing the effective sample size for MI estimation
The drifting problem caused by the unnormalized constant term C ∗ raises more issues when estimating MI.
Poole et al.
[2019] use a simple macro-averaging technique, i.e., averaging the estimated MI from each batch.
We can also consider a slight modification to the technique, where we call it the micro-averaging technique, by saving all the statistics network outputs T θ ( x, y ) for each batch and producing a single estimate based on all the outputs.
However, we proved that both averaging techniques yield wrong final estimates for biased estimators like I MINE [Belghazi et al., 2018], I SMILE
[Song and Ermon, 2020], I CLUB [Cheng et al., 2020], and
I InfoNCE [van den Oord et al., 2018].
Theorem 6.
(Estimation bias caused by drifting)
Both macro- and micro-averaging strategies produce a biased MI estimate when the drifting problem occurs.
To the contrary, self-normalizing or regularized MI estimators have the upper hand in this perspective.
By utilizing all the samples from multiple batches, they can effectively sidestep the batch size limitation problem [McAllester and
Stratos, 2020, Song and Ermon, 2020].
EXPERIMENTS 5.1 MI ESTIMATION VS. DOWNSTREAM TASK PERFORMANCE
Benchmark Design To measure the performance of the
MI estimators, one must design the target task to have the ground truth MI.
This constraint led previous works to evaluate the estimators only on artificial toy problems [Belghazi et al., 2018, Poole et al., 2019], where its connection to actual problems is fairly limited.
We design the two types of
MI estimation tasks with de facto image datasets to improve the existing benchmarks to reflect on the real-world tasks.
We defer all the proofs to the Appendix.
Theorem 7.
(Supervised learning)
Given a dataset D =
( X, Y ) where X is an sample, Y is the label for X , and
H ( Y ) is the entropy of the label set, I ( X, Y )
=
H ( Y ) .
Similarly, the true MI between images from the same class is also tractable based on the same assumption.
Theorem 8.
(Contrastive learning) Consider the dataset
D =
( X, Y ) .
Let X 1 be a sample drawn from the dataset and X 2 be another sample drawn from the subset with the same label Y to X 1 .
Then, I ( X 1 , X 2 )
=
I ( X 1 , Y ) = I ( X 2 , Y ) =
H ( Y ) .
Note that we assume statistical dependence between the image X and label Y from the point of view of information bottleneck [Tishby and Zaslavsky, 2015].
We derive the theorems above based on the assumption, where Y implicitly determines X .
Based on the above theorems, we use the two MI estimation problems as benchmarks that evaluate the performance of estimators.
We intentionally design Theorem 7 and Theorem 8 to mimic the existing tasks closely, namely, supervised and contrastive learning.
For Theorem 7, we can set the statistics network T θ ( X, Y ) =
f θ ( X ) · o ( Y ) where f θ ( X ) is the logits obtained from feeding the image X to the classification neural networks and o ( Y ) is the one-hot representation of the label Y .
If we use the InfoNCE estimator, this formulation becomes identical to solving the classification problem using negative log loss with the Softmax function, hence the name being the supervised learning benchmark (SLB).
Similarly, for Theorem 8, we can set
T θ ( X 1 , X 2 ) =
f θ ( X 1 ) · f θ ( X 2 ) and use the InfoNCE estimator to yield a commonly used contrastive loss [van den Oord et al., 2018, Chen et al., 2020].
Due to the strict assumption of statistical dependence, the theorems above cannot be used on standard datasets like
ImageNet dataset [Deng et al., 2009], as its samples often violate the single-label assumption.
However, we can still empirically compare the MI estimators by the relative size of their final MI estimation.
We conduct a demo experiment on ImageNet in the Appendix.
Evaluation To verify the performance of MI estimators, we perform our benchmark tasks on the CIFAR10 and CIFAR100 dataset [Krizhevsky, 2009].
As both CIFAR and CIFAR100 have a uniform label distribution, ideal MI is log 10 and log 100 , respectively.
In addition, to check whether this MI estimate task is actually helpful for downstream tasks, we evaluate each estimator on both dimensions: MI estimation and test set accuracy.
Similar to the existing settings in the contrastive learning literature [Chen et al., 2020, He et al., 2020], we design the test accuracy of
CLB by defining the label estimate ˆ y of each test set sample x Test to be the label of x = argmax x ∈ X Train f ( x ) · f ( x Test ) of the train dataset X Train .
Similarly for SLB, we chose Loss
Loss settings
[Table]
ˆ y = argmax y f ( x Test ) · o ( y ) where o ( y ) is the one-hot encoding of y . We ran the same experiment 5 times with different seeds to yield a 95% confidence interval.
5.2 COMPARISON WITH OUR BENCHMARK
To demonstrate the effectiveness of our novel regularization term, we regularize the two representations, D DV and D NWJ .
We test three realizations for each representation, I MINE [Belghazi et al., 2018], I SMILE
[Song and Ermon, 2020], and
I InfoNCE [van den Oord et al., 2018] for D DV , I NWJ [Nguyen et al., 2010], I TUBA [Poole et al., 2019], and I JS [Hjelm et al., 2019] for D NWJ .
We compare the original losses with its regularized counterparts, a total of 6 × 2 = 12 .
We do not apply averaging scheme on any of the losses, and choose the regularization weight λ ∈ { 0 . 1 , 0 . 01 , 0 . 001 } that shows the best MI estimation results.
See Table 1 for more details.
We observe in Table 2 that additional regularization generally induces better performance on both the MI estimation task and the downstream task (test accuracy).
Hence, adding the regularizer to a pre-existing supervised or contrastive learning loss seems to be a viable option to increase the performance further.
Even when the performance of the regularized loss slightly degrades, its negative impact is minimal.
This implies that even for the case where the regularizer is not greatly helpful, it does not greatly hinder optimization.
Especially, it is intriguing that many losses, I MINE , I ReMINE , and I ReTUBA , are better than I InfoNCE in SLB, which is used as the de facto standard in classification.
Also, I SMILE , I NWJ , and I TUBA fail to converge in CLB, where simply adding a regularization term solves the issue altogether, yielding a competitive or even better performance than all the other losses.
Given the fact that numerous contrastive learning literature suffers from instability [Caron et al., 2021, Bardes et al., 2021, Chen et al., 2020, He et al., 2020, Bardes et al., 2021], we emphasize that adding our regularization term can be a simple yet effective method to stabilize training.
Additionally, to observe the impact of regularization strength λ , we plot the benchmark performance for each λ in Fig.
We compare the losses on CLB as experimental results suggest that CLB is a more difficult task than SLB, showing significant performance differences between various losses.
On CIFAR-10, λ acts as a trade-off parameter [Figure] between test accuracy and MI estimation quality.
Performance trade-off has also been reported in other literature, where better MI estimation does not necessarily deliver better downstream performance [Tschannen et al., 2020,
Tian et al., 2020b].
However, compared to CIFAR-100, test accuracy differences are minimal, where MI differences are apparent.
I ReMINE
and I ReSMILE show excellent MI estimation quality in CIFAR-10 compared to other losses.
In contrast, test accuracy and MI estimation quality align well in the CIFAR-100 case.
I ReSMILE shows good overall performance, albeit its sensitivity towards regularization strength.
I ReInfoNCE , on the other hand, shows stable performance in the downstream task, sacrificing the MI estimation quality.
This result is further supported by the prominence of
I InfoNCE in the contrastive learning domain.
It is yet unclear where the difference between CIFAR-10 and CIFARcomes from, whether it is due to the difference in the level of difficulty of the dataset or the batch size used throughout the training.
We leave further analysis as future work.
5.3 COMPARISON WITH THE STANDARD TOY PROBLEM
We provide the quality of MI-based losses on the 20D Correlated Gaussian task [Belghazi et al., 2018, Poole et al., 2019] where the true MI is increased 5 times during optimization in Fig.
This experiment demonstrates how stable the MI-based losses estimate MI in a dynamically changing environment.
We apply the same settings from
[Table] Table 1, where we fix the regularization strength λ = 1 .for
all the losses.
With the exception of I InfoNCE , regularized losses show clear superiority over the original losses.
Regularization facilitates I MINE and I SMILE to avoid the instability which is mentioned in Section 3.
Also, regularization greatly enhances the MI estimation quality of I JS and lessens the variance of both
I NWJ and I TUBA .
CONCLUSION
In this paper, we identify the two symptoms behind the instability: The statistics network was not converging even after the loss seemed to converge, and its outputs from the product of marginal distribution explode during training.
We propose a novel regularization term to mitigate the instability during training by adding to various existing MI-based losses.
We theoretically and experimentally demonstrate that the added regularizer directly alleviates the two instability symptoms.
Finally, we present a benchmark that evaluates both the MI estimation power and its capability on the downstream tasks by imitating the supervised or contrastive learning settings.
We compare six different losses and their regularized counterparts on various benchmarks to show the method’s effectiveness and broad applicability. [Figure]
LIMITATIONS AND FUTURE WORKS
We suspect that the instability of MI estimators can also be related to the collapse problem [Bardes et al., 2021, Caron et al., 2021].
Further loss-based approaches to combat this problem by regularizing the network outputs may be helpful.
We expect that extending our methods to various contrastive learning losses may yield fruitful results for self-supervised learning, notably for other domains such as text or audio.
Also, our mathematical analysis is mainly focused on the drifting problem of I MINE , not the explosion problem of
I NWJ . For I NWJ , we suspect that the absence of the log function wrapping the exponential values makes the second term much more susceptible to output explosion due to its numerical instability. The added regularizer gives additional hints for the second term to converge to a specific value. However, we did not expand the discussion further in this paper.