Tight Mutual Information Estimation With Contrastive Fenchel-Legendre Optimization
Abstract
Successful applications of InfoNCE (Information Noise-Contrastive Estimation) and its variants have popularized the use of contrastive variational mutual information (MI) estimators in machine learning.
While featuring superior stability, these estimators crucially depend on costly large-batch training, and they sacriﬁce bound tightness for variance reduction.
To overcome these limitations, we revisit the mathematics of popular variational MI bounds from the lens of unnormalized statistical modeling and convex optimization.
Our investigation yields a new uniﬁed theoretical framework encompassing popular variational MI bounds, and leads to a new simple and powerful contrastive MI estimator we name FLO .
Theoretically, we show that the FLO estimator is tight, and it converges under stochastic gradient descent.
Empirically, the FLO estimator overcomes the limitations of its predecessors and learns more efﬁciently.
The utility of FLO is veriﬁed using extensive benchmarks, and we further inspire the community with novel applications in meta-learning.
Our presentation underscores the foundational importance of variational MI estimation in data-efﬁcient learning.
Introduction
Assessing the dependence between pairs of variables is integral to many scientiﬁc and engineering endeavors [61, 62].
Mutual information (MI) is a popular metric to quantify generic associations [45] and its empirical estimators have been widely used in applications such as independent component analysis [3], fair learning [28], neuroscience [53], Bayesian optimization [39], among others.
Notably, the recent advances in deep self-supervised learning (SSL) heavily rely on nonparametric
MI optimization [69, 52, 31, 13, 25].
In this study we investigate the likelihood-free variational approximation of MI using only paired samples, and improve the data-efﬁciency of current machine learning practices.
MI estimation has been extensively studied [6, 46, 45, 54, 56, 70, 10].
While most classical estimators work reasonably well for low-dimensional cases, they scale poorly to big datasets: naïve density-based estimator(s) and k -nearest neighbor estimators [41, 55, 20] struggle with high-dimensional inputs, while kernel estimators are slow, memory demanding and sensitive to hyperparameters [24, 23].
Moreover, these estimators are usually either non-differentiable or need to hold all data in memory.
Consequently, they are not well suited for emerging applications where the data representation needs to be differentiably optimized based on small-batch estimation of MI [35].
Alternatively, one can approach MI estimation through an estimated likelihood ratio [66, 35], but the associated numerical instability has raised concerns [2].
To scale MI estimation to the growing size and complexity of modern datasets, and to accommodate the need for representation optimization [8], variational objectives have been widely utilized recently [52].
Instead of directly estimating data likelihoods, density ratios, or the corresponding gradients [74], variational approaches appeal to mathematical inequalities to construct tractable lower or upper bounds of the mutual information [57], facilitated by the use of auxiliary critic functions 1 .
This practice turns MI estimation into an optimization problem.
Prominent examples include the
Barber-Agakov (BA) estimator [4], the Donsker-Varadhan (DV) estimator [15], and the NguyenWainwright-Jordan (NWJ) estimator [50].
These variational estimators are closely connected to the variational objectives for likelihood inference [1].
Despite reported successes, these variational estimators have a major limitation: their estimation variance grows exponentially to the ground-truth MI [47].
This is especially harmful to applications involving deep neural nets, as it largely destabilizes training [64].
An effective ﬁx is to leverage multi-sample contrastive estimators, pioneered by the work of InfoNCE [52].
However, the massive reduction in the variance comes at a price: the performance of the InfoNCE estimator is upper bounded by log K , where K is the number of samples used for estimation ( i.e. , batch-size) [57].
For a large MI, K needs to be sufﬁciently large to allow for an adequate estimate, consequently placing a signiﬁcant burden on computation and memory.
While variants of InfoNCE have been motivated to achieve more controllable bias and variance tradeoffs [57, 64], little research has been conducted on the cost-beneﬁt aspect of contrastive learning.
[Figure]
FLO provides a novel uniﬁed framework to analyze contrastive MI bounds.
A critical insight enabled by InfoNCE is that mutual information closely connects to contrastive learning [29, 52].
Paralleled by the empirical successes of instance discriminationbased self-supervision [48, 76, 13, 31] and multi-view supervision [68, 58], InfoNCE offers an InfoMax explanation to why the ability to discriminate naturally paired positive instances from the randomly paired negative instances leads to universal performance gains in these applications [43, 63, 57].
Despite these encouraging developments, the big picture of MI optimization and contrastive learning is not yet complete: ( i ) There is an ongoing debate about to what extent MI optimization helps to learn useful representation [72]; ( ii ) how does the contrastive view reconcile with those non-contrastive MI estimators; and crucial for practical applications, ( iii ) are the empirical tradeoffs made by estimators such as InfoNCE absolutely necessary?
Also theoretically, ( iv ) formal guarantees on the statistical convergence of popular variational non-parametric MI estimation are missing currently.
In this work we seek to bridge the above gaps by approaching the MI estimation from the novel perspective of energy modeling.
While this subject has recently been studied extensively using information-theoretic and variational inequalities, we embrace a new view from the lens of unnormalized statistical modeling.
Our main contributions include:
• Unifying popular variational MI bounds under unnormalized statistical modeling; • Deriving a simple but powerful novel contrastive variational bound called FLO ; • Providing theoretical justiﬁcation of the FLO bound (tightness and convergence); • Demonstrating strong empirical evidence of the superiority of FLO over its predecessors.
• Highlighting the importance of MI in data-efﬁcient learning with novel applications
We contribute in-depth discussion to bridge the gaps between contrastive learning and MI estimation, along with principled practical guidelines informed by theoretical insights.
Fenchel-Legendre Optimization for Mutual Information Estimation 2.1 Preliminaries
This section brieﬂy reviews the mathematical background needed for our subsequent developments.
Unnormalized statistical modeling deﬁnes a rich class of models of general interest.
Speciﬁcally, we are interested in problems for which the system is characterized by an energy function ˜
p θ ( x ) = exp( − ψ θ ( x )) , where θ is the system parameters and ψ θ ( x ) is known as the potential function .
The goal is to ﬁnd a solution that is deﬁned by a normalized version of ˜ p
θ ( x ) , i.e. ,
[Table]
, where L ( · ) is the loss function, µ is the base measure on X and
Z ( θ ) (cid:44) (cid:82) ˜
p θ ( x (cid:48) )
d µ ( x (cid:48) ) is called the partition function for ˜ p
θ ( x ) .
Problems in the above form arise naturally in statistical physics [60], Bayesian analysis [9], and maximal likelihood estimation [67].
A major difﬁculty with unnormalized statistical modeling is that the partition function Z ( θ ) is generally intractable for complex energy functions 2 , and in many applications Z ( θ ) is further composed by log Z ( θ ) , whose concavity implies any ﬁnite sample estimate Monte-Carlo of Z ( θ ) will render the loss function biased [59, 80].
Bypassing the difﬁculties caused by the intractable partition function is central to unnormalized statistical modeling [22, 49, 33, 36, 29].
Mutual information and unnormalized statistical models.
As a generic score assessing the dependency between two random variables ( X, Y ) , mutual information is formally deﬁned as the
Kullback-Leibler divergence (KL) between the joint distribution p ( x, y ) and product of the respective marginals p ( x ) p ( y ) [62], i.e. , I ( X ; Y ) (cid:44) E p ( x,y ) (cid:104) log p ( x,y ) p ( x ) p ( y ) (cid:105) .
The integrand log p ( x,y ) p ( x ) p
( y ) is often known as the point-wise mutual information (PMI) in the literature.
Mutual information has a few appealing properties: ( i ) it is invariant wrt invertible transformations of x and y , and ( ii ) it has the intuitive interpretation of the reduced uncertainty of one variable given another variable 3 .
To connect MI to unnormalized statistical modeling, we consider the classical Barber-Agakov ( BA ) estimator of MI [5].
To lower bound MI, BA introduces a variational approximation q ( y | x ) for the posterior p ( y | x ) , and by rearranging the terms we obtain an inequality
[Equation]
Here we have used notation I BA
( X ; Y | q ) to highlight the dependence on q ( y | x ) , and when q ( y | x ) =
p ( y | x ) this bound is sharp.
Unfortunately, this naïve BA bound is not useful for sample-based
MI estimation, as we do not know the ground-truth p ( y ) .
But we can bypass this difﬁculty by setting q θ ( y | x ) =
p ( y )
Z θ ( x ) e g θ ( x,y ) , where we call e g θ ( x,y ) the tilting function and recognize Z θ ( x ) =
E p ( y )
[ e g θ ( x,y ) ] as the associated partition function.
Substituting this q θ ( x | y ) into (1) gives the following unnormalized BA bound ( UBA ) that pertains to unnormalized statistical modeling [57]
[Equation]
[Equation]
While this UBA bound remains intractable, now with Z θ ( x ) instead of p ( y ) we can apply different techniques for empirical estimates of Z θ ( x ) to render a tractable surrogate target.
This has led to various popular MI bounds listed in Table 1 (see Appendix A for derivations).
InfoNCE and noise contrastive estimation.
InfoNCE is a multi-sample mutual information estimator proposed in [52], built on the idea of noise contrastive estimation (NCE) [29].
NCE learns statistical properties of a target distribution by comparing the positive samples from the target distribution to the “ negative ” samples from a carefully crafted noise distribution, and this technique is also known as negative sampling in some contexts [48, 26].
The InfoNCE estimator implements this contrastive estimation idea via using the naïve empirical estimate of Z θ ( x ) in UBA 4 , i.e.
[Equation]
where g θ is known as the critic in the nomenclature of contrastive learning, and we have used p K ( x, y ) to denote K independent draws from the joint density p ( x, y ) , and { ( x k , y k ) } Kk =1 for each pair of samples.
Here the positive and negative samples are respectively drawn from the joint p ( x, y ) and product of marginals p ( x ) p ( y ) .
Intuitively, InfoNCE tries to accurately classify the positive samples when they are mixed with negative samples, and the Proposition below formally characterizes
InfoNCE ’s statistical properties as a MI estimator.
[Equation]
Proposition 2.1 ([57]) .
InfoNCE is an asymptotically tight mutual information lower bound, i.e.
I K InfoNCE ( X ; Y | g θ ) ≤
I ( X ; Y ) , lim K →∞
I K InfoNCE ( X ; Y ) →
I ( X ; Y ) .
Fenchel-Legendre duality.
Our key idea is to exploit the convex duality for MI estimation.
Let f ( t ) be a proper convex, lower-semicontinuous function; then its convex conjugate function is deﬁned as f ∗ ( v ) (cid:44) sup t ∈D ( f ) { tv − f ( t ) } , where D ( f ) is the domain of function f
We call f ∗ ( v ) the Fenchel conjugate of f ( t ) , which is also known as the
Legendre transform in physics.
The Fenchel conjugate pair ( f, f ∗ ) are dual to each other, in the sense that f ∗∗ = f , i.e. , f ( t )
= sup v ∈D ( f ∗ ) { vt − f ∗ ( v ) } .
For f ( t ) = − log( t ) and its
Fenchel conjugate f ∗ ( v ) = − 1 − log( − v ) , we have inequality − log( t ) ≥ − u − e − u t + 1 , for u ∈ R (4) with the equality holds when u = log( t ) .
[Figure]
2.2 Fenchel-Legendre Optimization for tight mutual information estimation
With the above mathematical tools, we are ready to present the main result of this paper: a tight, data-efﬁcient variational MI lower bound that can be efﬁciently implemented.
Lower bounding MI with Fenchel-Legendre Optimization.
Our key insight is that MI estimation is essentially an unnormalized statistical model, which can be efﬁciently handled by the FenchelLegendre transform technique.
Take the integrand from UBA in (2) and we can rewrite it as
[Equation]
where p ( y (cid:48) ) is the same probability density as p ( y ) ( i.e. , Y (cid:48) is an independent copy of Y ).
Now let us use the Fenchel inequality of − log( t ) from (4), plugging it into the above equation
and then we have
[Equation]
for all u ∈ R .
This implies for any function u φ ( x, y ) :
X × Y → R , the following inequality holds
[Equation]
By putting (7) back to (2), we obtain our new Fenchel-Legendre Optimization ( FLO ) MI lower bound
I
[Equation]
and concludes the proof for the following Proposition.
Proposition
[Equation]
In practice, FLO can be estimated with the following naïve empirical K -sample estimator
[Equation]
Since the summation in ˆ I K FLO is not encapsulated by a convex log transformation, I K FLO (cid:44) E p K
[ˆ I K FLO ] is an unbiased estimator for I FLO
( X ; Y | g θ , u φ ) independent of the batch size K (see Figure 2).
Why is the FLO bound more appealing?
At ﬁrst sight, it may appear counter-intuitive that I FLO is a better MI bound compared to prior arts such as NWJ or InfoNCE : it seems to be more complicated as an extra variational function u φ ( x, y ) has been introduced.
To answer this question, we next explain the statistical meaning of the newly introduced u φ ( x, y ) , and establish some important statistical properties of FLO that makes it more favorable: that I FLO is tight, meaning the ground-truth MI can be recovered for some speciﬁc choice of g θ ( x, y ) and u φ ( x, y ) ; and that I K FLO for any batch size K is effectively optimizing InfoNCE with an inﬁnite batch size.
And in Sec 2.4, we further justify FLO ’s advantages from optimization perspectives.
Given the close connection between FLO and UBA , we ﬁrst recall UBA ’s optimal critic that gives the tight MI estimate is g ∗ ( x, y )
= log p ( x | y )
+ c ( x ) , where this c ( x ) can be any function of x [44].
This g ∗ ( x, y ) is not directly meaningful in a statistical sense, however, by integrating out y (cid:48) , we have
[Equation]
which is the likelihood ratio between the marginals and joint.
On the other hand, based on the Fenchel-Legendre inequality (4), we know for ﬁxed g ( x, y ) our FLO bound in (8) can be maximized with u g ( x, y )
=
log E p ( y (cid:48) )
(cid:104)
e g ( x,y
[Equation]
) − g ( x,y ) (cid:105) .
Putting these all together we have u g ∗ ( x, y )
= − log p ( x,y ) p ( x ) p ( y ) .
This shows the u φ ( x, y ) introduced in FLO actually tries to recover the negative PMI.
Comparing to the competing MI bounds that only optimizes for g θ , eliminating the drift term c ( x ) reveals FLO enjoys the appealing self-normalizing property [29] that helps stabilize training.
Plugging ( g ∗ , u g ∗ ) into (8), we readily see I FLO
( X ; Y | u g ∗ , g ∗ )
=
I ( X ; Y ) , proving FLO is a tight MI bound.
Proposition 2.3.
The FLO estimator is tight, the eqaulity holds when g ( x, y )
= log p ( x | y )
+ c ( x ) for arbitrary function c ( x ) and u ( x, y )
= − log p ( x,y ) p ( x ) p ( y ) .
Corollary 2.4.
Let ( g ∗ , u g ∗ ) be the maximizers for (8), then I ( X ; Y ) =
E p ( x,y ) [ − u g ∗ ( x, y )] .
Finally, we give a simple asymptotic argument showing FLO essentially optimizes InfoNCE with an inﬁnite batch size.
In virtue of the law of large numbers, we have the denominator in InfoNCE converging to lim K →∞ 1 K (cid:80) Kj
=
e g θ ( x i ,y j ) → E p ( y (cid:48) )
[ e g θ ( x i ,y (cid:48) ) ]
= Z θ ( x i ) , and consequently it recovers the UBA bound.
Since FLO is derived from UBA , we can view FLO as using the optimization of u φ ( x, y ) to amortize the difﬁculty of evaluating inﬁnite number of e g θ ( x i ,y j ) with InfoNCE .
Efﬁcient implementations of FLO .
A lingering concern is that the newly introduced u φ ( x, y ) can incur extra computation overhead.
This is not true, as we can maximally encourage parameter sharing by jointly model u φ ( x, y ) and g θ ( x, y ) with a single neural network f Ψ ( x, y ) :
X ×Y → R 2 with two output heads, i.e. , [ u i , g i ] =
f Ψ ( x i , y i ) .
Consequently, while FLO adopts a dual critics design, it does not actually invoke extra modeling cost compared to its single-critic counterparts ( e.g. , InfoNCE ).
Experiments show this shared parameterization in fact promotes synergies and speeds up learning (see our ablation studies in Appendix F.6).
To further enhance the computation efﬁciency, we consider a massively parallelized bi-linear critic design that uses all in-batch samples as negatives.
Let g θ ( x, y )
=
τ ·(cid:104) h θ ( x ) , ˜ h
( y ) (cid:105) , where h : X → S p and ˜
h :
Y → S p are respectively encoders that map data to unit sphere S p embedded in R p
+1 , (cid:104) a, b (cid:105) =
a T b is the inner product operation, and τ > 0 is the inverse temperature parameter.
Thus the evaluation of the Gram matrix G = τ · h ( X ) T ˜ h
( Y ) , where [ X , Y ] ∈ R K × ( d x + d y ) is a mini-batch of K -paired samples and g θ
( x i , y j )
=
G ij , can be parallelized via matrix multiplication.
In this setup, the diagonal terms of G are the positive scores while the off-diagonal terms negative scores.
A similar strategy has been widely employed in the contrastive representation learning literature ( e.g. , [13]) 5 .
We can simply model the PMI critic as u ( x, y ) =
MLP ( h ( x ) , ˜ h ( y )) , whose computation cost is almost neglectable in practice, where feature encoders h, ˜
h dominate computing.
2.3 Connections to the existing MI bounds
Due to space limitations, we elaborate the connections to the existing MI bounds here, and have relegated an extended related work discussion in a broader context to Appendix A, H.
From log -partition approximation to MI bounds.
To embrace a more holistic understanding, we list popular variational MI bounds together with our FLO in Table 1, and visualize their connections in
Figure 1.
With the exception of JSD , these bounds can be viewed from the perspective of unnormalized statistical modeling, as they differ in how the log partition function log
Z ( x ) is estimated.
We broadly categorize these estimators into two families: the log -family ( DV , MINE , InfoNCE ) and the exponentialfamily ( NWJ , TUBA , FLO ).
In the log -family, DV and InfoNCE are multi-sample estimators that leverage
Here g ( x, y ) , u ( x, y ) and u ( x ) are variational functions to be optimized, σ ( u )
= 1 1+exp( − u ) is the Sigmoid function, E [ f ( u ) , η ] denotes exponential average of function f ( u ) with decay parameter η ∈ (0 , 1) , and α ∈ [0 , 1] is the balancing parameter used by α -InfoNCE trading off bias and variance between InfoNCE and TUBA .
we use ( x i , y i ) to denote positive samples from the joint density p ( x, y ) , and ( x i , y j ) or ( x (cid:48) k , y (cid:48) k ) to denote negative samples drawn from the product of marginal p ( x ) p ( y ) .
In context, y ⊕ and y (cid:9) have the intuitive interpretation of positive and negative samples.
We exclude variational upper bounds here because their computations typically involve the explicit knowledge of conditional likelihoods.
[Table]
direct Monte-Carlo estimates ˆ Z for log Z ( x ) , and these two differ in whether to include the positive sample in the denominator or not.
To avoid the excessive in-batch computation of the normalizer and the associated memory drain, MINE further employed an exponential moving average (EMA) to aggregate the normalizer across batches.
Note for the log -family estimators, their variational gaps are partly caused by the log -transformation on ﬁnite-sample average due to Jensen’s inequality ( i.e. , log Z = log E
[ ˆ Z ] ≥ E [log ˆ Z ] ).
In contrast, the objective of exponential-family estimators do not involve such log -transformation, since they can all be derived from the Fenchel-Legendre inequality:
NWJ directly applies the Fenchel dual of f -divergence for MI [51], while TUBA exploits this inequality to compute the log partition log
Z ( x ) = log E p ( y (cid:48) )
[exp( g ( x, y (cid:48) ))] .
Motivated from a contrastive view, our FLO applies the Fenchel-Legendre inequality to the log -partition of contrast scores.
A contrastive view for MI estimation.
The MI estimators can also be categorized based on how they contrast the samples.
For instance, NWJ and TUBA are generally considered to be non-contrastive estimators, as their objectives do not compare positive samples against negative samples on the same scale ( i.e. , log versus exp ), and this might explain their lack of effectiveness in representation learning applications.
For JSD , it depends on a two-stage estimation procedure similar to that in adversarial training to assess the MI, by explicitly contrasting positive and negative samples to estimate the likelihood ratio.
This strategy has been reported to be unstable in many empirical settings.
The log -family estimators can be considered as a multi-sample, single-stage generalization of JSD .
However, the DV objective can go unbounded thus resulting in a large variance, and the contrastive signal is decoupled by the EMA operation in MINE .
Designed from contrastive perspectives, InfoNCE trades bound tightness for a lower estimation variance, which is found to be crucial in representation learning applications.
Our FLO formalizes the contrastive view for exponential-family MI estimation, and bridges existing bounds: the PMI normalizer exp( − u ( x, y )) is a more principled treatment than the EMA in MINE , and compared to DV the positive and negative samples are explicitly contrasted and adaptively normalized.
Important FLO variants.
We now demonstrate that FLO is a ﬂexible framework that not only recovers existing bounds, but also derives novel bounds such as
[Equation]
Recall the optimal g ∗ ( x, y )
= log p ( x | y )
+ c ( x ) and u ∗ ( x, y )
= − log p ( x,y ) p ( x ) p ( y ) , which motivates us to parameterize u ( x, y ) in the form of − g θ ( x, y )
+ s ψ ( x ) , where s ψ ( x ) models the arbitrary drift c ( x ) , and this recovers the TUBA bound.
Additionally, we note that ( i ) ﬁxing either of u and g , and optimizing the other also gives a valid lower bound to MI; and ( ii ) a carefully chosen multi-input u ( { ( x i , y i ) } ) can be computationally appealing.
As a concrete example, if we set u φ to u θ ( { ( x i , y i ) } ) ← log (cid:16) 1 K (cid:80) j e c ( x i ,y i ,y j ; g θ ) (cid:17) and update u θ ( x, y ) while artiﬁcially keeping the critic g θ ( x, y ) ﬁxed 6 , then FLO falls back to DV .
Alternatively, we can consider the Fenchel dual version of it: using the same multi-input u θ ( { ( x i , y i ) } ) above, treat u φ as ﬁxed and only update g θ , and this gives us the novel MI objective in (11), we call it Fenchel-Donsker-Varadhan (FDV) estimator.
2.4 Gradient and convergence analysis of FLO
In this section, we will establish that FLO better optimizes the MI because its gradient is more accurate than competing variational bounds such as NWJ and TUBA ; also, we provide the ﬁrst convergence analysis for variational MI estimation by showing FLO converges under SGD.
First, recall most tractable variational MI bounds are derived from and upper bounded by the intractable UBA bound [57].
For instance, with the same critic g θ we have I NWJ ≤ I TUBA ≤
I UBA .
So if we can show ∇ θ I FLO ≈ ∇ θ
I
UBA then FLO is better optimized.
To simplify notations, we denote c θ ( x, y, y (cid:48) )
(cid:44) g
θ
( x, y (cid:48) )
− g θ ( x, y ) and E θ ( x, y ) (cid:44) 1 / E p
( y (cid:48) )
[ e c θ ( x,y,y (cid:48) ) ] , and we can easily verify
E p ( y (cid:48) )
(cid:104) ∇ θ (cid:110)
e c θ ( x,y,y (cid:48)
[Equation]
Since for ﬁxed g θ ( x, y ) the corresponding optimal u ∗ θ ( x, y ) maximizing I FLO ( u φ , g θ ) (cid:44) 1 − (cid:110) u φ ( x, y ) + E p
y (cid:48) )
[ e − u φ ( x,y )
+ c ( x,y,y (cid:48) ; g θ ) ] (cid:111) is given by u ∗ θ ( x, y ) =
log E p ( y (cid:48) )
[ e c θ ( x,y,y (cid:48) ) ]
= − log E θ ( x, y ) (using (4)), we see that the term e − u φ ( x,y ) is essentially optimized to approximate
E θ ( x, y ) .
To emphasize this point, we now write ˆ E θ ( x, y ) (cid:44) e − u φ ( x,y ) .
When this approximation is sufﬁciently accurate ( i.e. , E θ ≈
ˆ E θ ), we can see that ∇ I FLO approximates ∇ I UBA as follows
[Equation]
We can prove FLO will converge under much weaker conditions, even when this approximation ˆ u ( x, y ) is rough.
The intuition is as follows: in (13), the term ˆ E θt E θt only rescales the gradient, so the optimizer is still proceeding in the same direction as UBA in SGD.
The informal version of our result is summarized in the Proposition below (see Appendix E for the formal version and proof).
Proposition 2.5 (Convergence of FLO , informal version) .
Let { η t } ∞ t =1 be the stochastic Robbins-Monro sequence of learning rates: (cid:80) t E [˜ η t ] = ∞ and (cid:80) t
E [˜ η 2 t ] < ∞ .
If ˆ E θt E θt is bounded between [ a, b ] ( 0 < a < b < ∞ ), then under the stochastic gradient descent scheme described in Algorithm 1, θ t converges to a stationary point of I UBA ( g θ ) with probability 1 , i.e. , lim t →∞ (cid:107)∇ I UBA ( g θ t ) (cid:107) = 0 .
Additionally assume I UBA is convex with respect to θ , then FLO converges with probability 1 to the global optimum θ ∗ of I UBA from any initial point θ 0 .
Algorithm
Empirical data ˆ p d = { ( x i ,y i ) } ni =
Model parameters Ψ
= ( θ,φ ) for t = 1 , 2 , ··· do
Sample i,j iid ∼ [ n ] u ii =
u φ ( x i ,y i ) ,g ii =
[Equation]
( x i ,y i ) , g
[Equation]
=
[Equation]
θ (
[Equation]
[Equation] [Equation]
+ exp( − u
[Equation]
Ψ t =
Ψ t − η t ∇ Ψ F end for
Importantly, this marks the ﬁrst SGD convergence analyses for variational MI estimators.
The convergence analyses for MI estimation is non-trivial and scarce even for those standard statistical estimators [54, 20, 59].
For variational MI bounds, existing convergence analyses only apply to oracle estimators but not practical estimators discussed here [50, 65].
The lack of convergence guarantees has led to a proliferation of unstable MI-estimators used in practice (in particular, DV , JSD , and
MINE ) that critically rely on various empirical hacks to work well (see discussions in [64]).
Our work establishes a family of variational MI estimators that provably converges, a contribution we consider signiﬁcant as it ﬁlls an important gap in current literature on both theoretical and practical notes.
While our convergence result is local, in practice we have observed that FLO .
[Figure] [Figure]
inefﬁciency, but its smaller variance helps in the more challenging dynamic case.
Experiments
We consider an extensive range of tasks to validate FLO and benchmark it against state-of-theart solutions.
To underscore the practical signiﬁcance of MI in efﬁcient machine learning, we demonstrate example applications from data collection (in statistical parlance, experimental design), self-supervised pre-training, to meta/transfer-learning.
Limited by space, we present only the key results in the main text, and defer ablation studies and details of our experimental setups to the
Appendix.
Our code is available from https://github.com/qingguo666/FLO .
All experiments are implemented with PyTorch .
Comparison to baseline MI bounds.
We start by comparing FLO to the following popular competing variational estimators: NWJ ,
TUBA , and InfoNCE .
We use the bilinear critic implementation for all models which maximally encourages both sample efﬁciency and code simplicity, and this strategy does perform best based on our observations.
We consider the synthetic benchmark from [57], where ( X ∈ R d , Y ∈ R d ) is jointly standard Gaussian with diagonal cross-correlation parameterized by ρ ∈ [0 , 1) .
We report d = 10 and ρ ∈ [0 , 0 . 99] here (other studies only report ρ up to 0 .
9 , which is less challenging.), providing a reasonable coverage of the range of MI one may encounter in empirical settings.
[Figure]
To focus on the bias-variance trade-off, we plot the decimal quantiles in addition to the estimated MI in Figure 3, where FLO signiﬁcantly outperformed its variational counterparts in the more challenging high-MI regime.
In Figure 5, we show FLO also beats classical MI estimators [41, 73, 21].
In the Appendix H, we further discuss recent works on parametric estimators [14, 10] and alternative information metrics [78].
Bayesian optimal experiment design (BOED).
We next direct our attention to BOED, a topic of signiﬁcant interest shared by the statistical and machine learning communities [11, 75, 32, 19].
The performance of machine learning models crucially relies on the quality of data supplied for training, and BOED is a principled framework that optimizes the data collection procedure (in statistical parlance, conducting experiments ) [18].
Mathematically, let x be the data to be collected, θ be the parameters to be inferred, and d be the experiment parameters the investigator can manipulate ( a.k.a ,
[Figure]
S
I R P
o p u
l a t i o n s
SIR Epidemic Trajectory
HealthyInfectedRecoveredFLO
[Figure]
FLO Ground
[Figure]
We consider three carefully-selected models from recent literature for their progressive practical signiﬁcance and the challenges involved [17, 37, 38]: static designs of ( i ) a simple linear regression model and ( ii ) a complex nonlinear pharmacokinetic model for drug development; and the dynamic policy design for ( iii ) epidemic disease surveillance and intervention ( e.g. , for Covid-19 modeling).
Designs with higher MI are more favorable, because it implies the data carries more information.
In
Figure 4 we compare design optimization curves using different MI optimization strategies, where
FLO consistently leads.
Popular NWJ and InfoNCE reports different tradeoffs that are less susceptible to FLO .
We also examine the FLO predicted posteriors and conﬁrm they are consistent with the ground-truth parameters (Figure 6 right).
For the dynamic policy optimization, we also manually inspect the design strategies reported by different models (Figure 6 left,middle).
Consistent with human judgement, FLO policy better assigns budgeted surveillance resources at different stages of pandemic progression.
A novel meta-learning framework.
A second application of our work is to meta-learning, an area attracting substantial recent interest.
In meta-learning, we are concerned with scenarios that at training time, there are abundant different labelled tasks, while upon deployment, only a handful of labeled instances are available to adapt the learner to a new task.
Brieﬂy, for an arbitrary loss (cid:96) t
(ˆ y, y ) , where t is the task identiﬁer and ˆ y = f
( x ) is the prediction made by the model, we denote the risk by R t ( f ) =
E p t ( x,y )
[ (cid:96) t ( f ( x ) , y )] .
Denote R ( f ) (cid:44) E t ∼ p ( t )
[ R t ( f )] as the expected risk for all tasks and ˆ R ( f ) for the mean of empirical risks computed from all training tasks.
Inspired by recent information-theoretic generalization theories [77], we derived a novel, principled objective
[Equation]
where λ is known given the data size and loss function, ( ˆ D t , ˆ E t ) are respectively data and task embeddings for training data, which for the ﬁrst time lifts contrastive learning to the task and data distribution level.
Our reasoning is that L Meta-FLO ( f ) theoretically bounds R ( f ) from above, and it is relatively sharp for being data-dependent.
We give more information on this in the Appendix and defer a full exposition to a dedicated paper due to independent interest and space limits here.
Note other MI bounds are not suitable for this task due to resource and vari-
[Table]
[Figure]
a Note SpecNCE does not explicitly target mutual information ance concerns.
In Figure 7 we show Meta-FLO wins big over the state-of-the-art model agnostic meta-learning (MAML) model on the regression benchmark from [16]. [Figure]
[Figure] [Figure]
Self-supervised learning (SSL) .
Finally, we wrap our experiments with one of the prime applications of contrastive MI estimation in machine learning: SSL for model pre-training.
Here we focus on how
FLO -inspired objectives can improve the current practice of SSL, and given this topic’s independent interest, we defer in-depth discussions in our dedicated work [12] where SSL-speciﬁc problems such as training diagnosis and low-precision numerical overﬂow are explored in detail.
In this experiment, we follow the SSL setup described in the SimCLR paper [13]: in the pre-training phase, we optimize the mutual information between difference augmentations of the same image ( i.e. , scaling, rotation, color jitting, etc. ); and use linear probing accuracy as our perfromance criteria.
We compare the effectiveness of the InfoNCE -based SimCLR framework [13] to our FLO -based alternatives.
To ensure fair comparison, we have used the FDV variant deﬁned in Eq.
(11) as our training objective, so that we are not introducing extra parameters to model u ( x, y ) .
We call our new model FlatCLR because, perhaps counter-intuitively, the second term in Eq.
(11) contributing all the learning signal is constant one in value ( i.e. , being ﬂat).
In Figure 8 and 9, we show our new model
FlatCLR shows superior sample efﬁciency compared to the SOTA SimCLR (a 8 × boost for the same performance, FlatCLR - 32 versus SimCLR - 256 ).
This result is signiﬁcant because SimCLR ’s crucial reliance on large-batch training is a well-known limitation [30, 79, 42].
Figure 10 shows typical training curves with the respective models.
Note that while the empirical estimates of MI are tied between the two methods, FDV optimized representation enjoys a better ground-truth MI 7 , which can be explained by its robustness to the numerical overﬂow issue (see [12] for details).
Further comparisons on the ground-truth MI estimation with different estimators can be found in Table 2.
Conclusion
We have described a new framework for the contrastive estimation of mutual information from energy modeling perspectives.
Our work not only encapsulates popular variational MI bounds but also inspires novel objectives such as FLO and FDV , which comes with strong theoretical guarantees.
In future work, we will leverage our theoretical insights to improve practical applications involving MI estimation, such as representation learning, fairness, and in particular, data efﬁcient learning.
Acknowledgements
The authors would like to thank the anonymous reviewers for their insightful comments.
Q Guo gratefully appreciate the support of Amazon Fellowship.
X Deng would like to thank the Advanced
Research Computing program at Virginia Tech and Virginia’s Commonwealth Cyber Initiative (CCI)
AI testbed for providing computational resources, also appreciate the CCI and CCI-Coastal grants to
Virginia Tech.
Part of this work is done before C Tao joined Amazon, and he was funded by National
Science Foundation Grant No. 1934964.
This work used the Extreme Science and Engineering
Discovery Environment (XSEDE), which is supported by National Science Foundation grant number
ACI-1548562 [71] and used the Extreme Science and Engineering Discovery Environment (XSEDE)
PSC Bridges-2 and SDSC Expanse at the service-provider through allocation TG-ELE200002 and
TG-CIS210044.
[-We want to solve for ∇ f ( θ ) = 
-]
[-Estimated ̂ g ( x , y ) − −   
-]