Tight Mutual Information Estimation With Contrastive Fenchel-Legendre Optimization
Authors: Qing Guo
Junya Chen
Dong Wang
Yuewei Yang
Xinwei Deng
Jing Huang
Lawrence Carin
Fan Li
Chenyang Tao
Venue: NeurIPS 2022 Accept (2022)
Keywords: mutual information
variational inference
contrastive learning
few-shot learning
meta learning

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 sacrifice 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 unified theoretical framework encompassing popular variational MI bounds, and leads to a novel, 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 proposed FLO estimator overcomes the limitations of its predecessors and learns more efficiently. The utility of FLO is verified 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-efficient learning.

Tight Mutual Information Estimation With Contrastive Fenchel-Legendre Optimization
Qing Guo! Junya Chen', Dong Wang", Yuewei Wang", Xinwei Deng' Lawrence Carin Fan Li Jing Huang', Chenyang Tao
Virginia Tech Duke University "KAUST Amazon
gguo07010vt. edu, chenyang. tao@duke. edu
Abstract
Successful applications of InfoNCE (Information Noise-Contrastive Estimation) and its variants have popularized the use of contrastive variational mutual infor- mation (MI) estimators in machine learning. While featuring superior stability, these estimators crucially depend on costly large-batch training, and they sacrifice revisit bound tightness for variance reduction. To overcome these limitations, the mathematics of popular variational MI bounds from the lens of unnormalized statistical modeling and convex optimization. Our investigation yields new uni- fied theoretical framework encompassing popular variational MI bounds, and leads new simple and powerful contrastive MI estimator we name FLO. Theoreti- cally, 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 efficiently. The utility of FLO is verified using exten- sive benchmarks, and we further inspire the community with novel applications in meta-learning Our presentation underscores the foundational importance of variational MI estimation data-efficient learning.
Introduction
Assessing the dependence between pairs of variables is integral to many scientific and engineering endeavors [66, [67). Mutual information (MI) is popular metric to quantify generic associations [50]. and its empirical estimators have been widely used in applications such as independent compo- nent analysis (3]. fair learning [32] neuroscience [58), Bayesian optimization 43). among others.
Notably, the recent ad advances in deep self-supervised learning (SSL) heavily rely on nonparametric
15 29 In this study investigate the likelihood-free variational approximation of MI using only paired samples, and improve the data-efficiency of current machine learning practices.
MI optimization [76 557 35
MI estimation has been extensively studied 51 50 59 61 77 10). While most elassical 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 [45] 60, 24 struggle with high-dimensional inputs, while kernel estimators are slow, memory demanding and sensitive to hyperparameters 28 27 Moreover, these estimators are usually either non non-differemtiable or need to hold data in memory. Consequently, they not well suited for emerging applications where the data representation needs to be differentiably optimized based on small- estimation of MI 39]. Alternatively, one can approach MI estimation through estimated likelihood ratio 71 39 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 57 Instead of directly estimating data likelihoods, density ratios, or the corresponding gradients variational approaches appeal to mathematical inequalities to construct tractable lower or
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

upper bounds of the mutual information [62]. facilitated by the use of auxiliary critic functions This practice turns MI estimation into an optimizaaiinn problem. Prominent examples include the Barber-Agakov (BA) estimator [4], the Donsker- Varadhan (DV) estimator [19], and the Nguyen- Wainwright-Jorlan (NWJ) estimator (55). These variational estimators are closely connected the variational objectives for likelihood inference [
Despite reported successes, these variational estimators have major limitation: their estimation variance grows exponentially to the ground-truth MI [52]. This is especially harmful to applications involving deep neural nets, as it largely destabilizes training [9]]. An effective fix is to leverage multi-sample contrastive estimators, pioneered by the work of Inf oNCE [57]. However, the massive reduction in the variance comes bounded by log K, where K is the number of samples used for estimation (i.e., batch-size) [62]. For a large MI, K needs to be sufficiently large to allow for an adequate estimate, consequently placing a significant burden on computation and memory. While variants of Inf oNE have been motivated to achieve more controllable bias and variance tradeoffs 62 69],
price: the performance of the Inf estimator is upper
little research has been conducted on
the cost-benefit aspect of contrastive learning.
A critical insight enabled by InfoNCE is that mutual infor- mation closely connects to contrastive learning 57]. Par- alleled by the empirical successes of instance discrimination- based self supervisro [53] 83 15 35] and multi-view super- vision 77 633, InfoNCE offers an InfoMax explanation to why the ability to discriminate naturally paired posi- tive instances from the randomly paired negative instances leads to universal performance gains in these applications 47 68 62). Despite these encouraging developments, the big picture of MI optimization and contrastive learning is not yet complete: (i) There an ongoing debate about to what extent MI optimization helps to learn useful represen- tation 79]; (ii) how does the contrastive view reconcile with those non-contrastive MI estimators; and crucial for practical applications, (dit) are the empirical tradcoffs made by estimators such InfoNCE absolutely necessary? Also theoretically. (iv) formal guarantees on the statistical con- vergence of popular variational non-parametric MI estimation
Intractable
Mutual Information
((X:P)
Tractable
Curs
JSD
RA
l'enceel Family
TUDA
FLO
NWJ
HDV
MINE
InfoNCE
Boands
Figure 1:
Schematic of variational lower bounds of mutual information. FLO provides novel unified framework to analyze contrastive MI bounds. 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-theoreii and variational inequalities, we embrace new view from the lens of unnormal- ized 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 justification 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-efficient learning with novel applications
We contribute in-depth discussion bridge the gaps between contrastive learning and MI estimation, along with principled practical guidelines informed by thcoretical insights.
2
Fenchel-Legender Optimization for Mutual Information Estimation
2.1
Preliminaries
This section briefly reviews the mathematical background needed for our subsequent developments.
Unnormalized statistical modeling defines rich class of models of general interest. Specifi- cally, we are interested in problems for which the system is characterized by an energy function fo(x) exp(-1ha (z)), where 0 is the system parameters and (2) is known as the potential function. The goal is to find a solution that is defined by normalized version of fo(x), i.e.,
When estimates sharp, these critic functions usually recover some transformation of the likelihood ratio.

---

ming {c J Fo(x)) du(()) J. where [() is the loss function, f is the base measure on X' and 2(0) _ Po (x') du((²") ) is called the partition function for (z)). Problems in the above form arise naturally in statistical physics (65], Bayesian analysis [9] and maximal likelihood estimation [72). A major difficulty with unnormalized statistical modeling is that the partition function Z generally intractable for complex energy functions P] and in many applications Z(0) is further composed by log Z()), whose concavity implies any finite sample estimate Monte- Carto of Z will render the loss function biased [64] [7). Bypassing the difficulties caused by the intractable partition
function is central to unnormalized statistical modeling (26 54 37 40 33
Mutual information and unnormalized statistical models. As generic score assessing the dependency between two random variables (X. Y), mutual information is formally defined as the Kullback-Leibler divergence (KL) between the joint distribution p(æ;, y) and product of the respective
pla-y) marginals p(z)) (1)) 677. i.e., ((X;Y) & Ep(x,y) log p(a))tt)] The integrand log p(#)pty) is often
p(z.y)
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
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 ((ylz) for the posterior p(yl (z). and by rearranging the terms we obtain inequality
p(y[æ)]
(u1x)) p(v)
I(X;Y)
Ep(z.))
log
= Ep(z.))
log
+ Bp(n) (L(p(vl)) I| [((l)))]
p(v)
log 4(vlxr)
(X;:
p(v)
Here we have used notation (x: lg) to highlight the dependence on ((vlz), and when ((ylz) =
p(y(z) this bound is sharp. Unfortunately, this naive BA bound is not useful for sasple-based
MI estimation, as we do not know the ground-truth p(v). But we can bypass this difficulty by
setting qo (y((x) Za(a)
p(y) cBo(2.2), where we call (1.(v))) the tilting function and recognize Zo(x)
Ep(u) [e (s.ar) as the associated partition function. Substituting this (o(xl)) into (I) gives the
following unnormalized BA bound (UBA) that pertains to unnormalized statistical modeling [62]
egu(z.p))
LUBa (X; Y (lga) A E) (go y) - log Zo(x)] 1pp(x) (EpGlle)
log Zo(x)
(2)
While this UBA bound remains intractable, now with Zp (x:) instead of p(y) we can apply different techniques for empirical estimates of Zz(z)) to render tractable surrogate target. This has led to various popular MI bounds listed in Table (see Appendix A for derivations).
InfoNCE and noise contrastive estimation. InfoNCE is multi-sample mutual information esti- mator proposed in [57], built on the idea of noise contrastive estimation (NCE) 33 NCE learns statistical properties of a target distribution by comparing the positive samples from the target dis- tribution to the "negative" samples from a carefully crafted noise distribution, and this technique is also known as negative sampling in some contexts (53 (3). The Inf oNCE estimator implements this contrastive estimation idea via using the naive empirical estimate of Zz(x) in UBA i.e.
(0((2101))
Jk InfoUcE (X;Y Y (g)) A (
log
K In (X;Y) A max [TK F INSORCE (X; Y(go))
E, cg0(21.35) K
(3)
where go is known as the critic in the nomenclature of contrastive learning, and we have used pK (z, y) to denote K independent draws from the joint density p(z, y), and {(Ik, ) k=z 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(() Intuitively, ONCE tries to accurately classify the positive samples when they are mixed with negative samples, and the Proposition below formally characterizes InfoNCE's a statistical properties a MI estimator.
In the sense that they do not render closed from expressions.
è Formalll, YX; (X)) - H (X H(Y) - (Y))). where (X) (resp. (X())) denotes the Shannon entropy (resp. conditional Shannon entropy) of random variable.
This estimator is technically equivalent to the original definition due the symmetry of K samples.
3

---

Proposition 2.1 ((62). InfoNCE is an asymptotically tight mutual information lower bound, i.e. I Intouue go) S Y), k I (antone(Xx;;)) (X; (X;)).
Fenchel-Legendre duality. Our key idea to exploit the convex duality for MI estimation. Let [() be proper con- vex, lower-semicootiinous function; then its convex conju- gate function defined as (u) - suprep(r) (tv f()), where D()) the domain of function f [38]. We call (v) the Fenchel conjugate of f(t), which is also known the Legendre transform in physics. The Fenchel conjugate pair
Comparison of Sample Efficiency
M Information
True MI
InfoNCE K=2
1.5
InfoNCE K=5
InfoNCE
InfoNCE K-20
InfoNCE K-50
f*) are dual to each other, in the sense that f f. i.e.
f(t) = suPrED(s-) fut - (u)). For f(t) = - log(t) and its
0.0
0.2
0.4
0.6
0.8
1.0
Fenchel conjugate I (v) = -1 (og(-w), we have inequality
Figure 2: K-sample InfoNCE and single-sample FLO. Note FLO is tight regardless of sample-size
(tg((t) 2 -u - et + for u E R
(4)
with the equality holds when log(!).
2.2 Fenchel-Legendre Optimization for tight mutual information estimation
With the above mathematical tools, we are ready present the main result of this paper: tight, data-efficient variational MI lower bound that can be efficiently implemented.
Lower bounding MI with Benchel-Legendee Optimization. Our key insight is that MI estimation essentially an unnormalized statistical model, which can be efficiently handled by the Fenchel- Legendre transform technique. Take the integrand from UBA in (2) and we can rewrite
log (xp(go(x,)) - log (exp(ax, v²) g( v))]}
(5)
(z(x)
where p(y") is the same probability density as ply) (Le., Y' an independent copy of Y). Now let us use the Fenchel inequality of log(t) from plugging it into the above equation and then we have
exp(go (z. 3) 2 f-uu - e e "E (exp(g(x, v/) g(x, y))l} +
log
(6)
(z)
for all u E R. This implies for any function ( 4) X IR, the following inequality holds exp(ge(xx))) 2 -Ius
log
(us((; (z, y) (2.0) E p(y' [exp(e(<, v²) g(2, w))]} + 1.
(7)
Zo(x)
By putting 7 back to 2p we obtain our new Fenchel-Legendre Optimization (FLO) MI lower bound
(2, v) + -=)()))))p (c Bpoopledoooss + +1, (8)
Ino (x::XYYloo us) Ep(r.a))
and concludes the proof for the following Proposition.
Proposition 2.2. (X; Ylge, ub) S JuBa (X; Y Ylge) S (X;
In practice, FLO can be estimated with the following naïve empirical K-sample estimator
FLo (Xx (X; Y (go, us) A -
Σ go(m))))) go(m))s)) 371
1.
K-1
Since the summation in I is not encapsulated by a convex log transformation, I & Er, [iA is an unbiased estimator for VFLo(X ; Ylge, u&) independent of the batch size K (see Figure 2}.
Why is the FLO bound more appealing? At first sight, it may appear counter-intuitive that 4TLO is better MI bound compared to prior arts such as NWJ or InfoNCE: it seems to be more complicated as an extra variational function us (z. y) has been introduced. To answer this question, we next explain the statistical meaning of the newly introduced us(x, y), and establish some important statistical properties of FLO that makes more favorable: that TFLO is tight, meaning the ground-truth MI can be recovered for some specific choice of (r y) and wo( y); and that 7 for any batch size K is effectively optimizing InfE with an infinite batch size. And in See 2.4 we further justify FLO's advantages from optimization perspectives.

---

Given the close connection between FLO and UBA, we first recall UBA's s optimal critic that gives the tight MI estimate is g" (z, y) log p(zly) + c(x), where this c(z) can be any function of x [4]] This o" (x, y) is not directly meaningful in statistical sense, however, by integrating out y', we have
Ep(()) (00' = Ep(rr)
p(aclu)) p(xl))]
p(z) p(x)p(v) p(zly)
(10)
p(x,
which is the likelihood ratio between the marginals and joint. On the other hand, based on the Fenchel-Legender inequality 4 we know for fixed g(2, 3) our FLO bound in (8) can be
maximized with (1 v) = log Ep(c") (00(2,0)-0(xx0))
Putting these all together we have
p(z)p(a)) This shows the us(2z,) y) introduced in FLO actually tries to recover the negative PMI. Comparing to the competing MI bounds that only optimizes for ge, eliminating the drift term c(z) reveals FLO enjoys the appealing self nelf-ozmal property (33 that helps stabilize
"gg (z. y) log p(e.))
training. Plugging (g" ug+) into (8). we readily see [xLo(X; Y [Ug- g¹) = [X;; Y). proving FLO is a tight MI bound.
Proposition 2.3. The FLO estimator is tight, the eqaulity holds when g(r, 3) = logp(xly) - c(x) p(z.z)
for arbitrary function c(z) and u(x, y) = -loo p(s)p(o))
Corollary 2.4. Let (g", ((g)) be the maximizers for (b. then ((X;Y) = Bp(a.wl- uug (1r, v)].
Finally, we give a simple asymptotic argument showing FLO essentially optimizes InfoNCE with an infinite batch size. In virtue of the law of large numbers, we have the denominator in InfoNCE converging --mo- R
Ze (ax), and consequently recovers the UBA bound. Since FLO is derived from UBA, we can view FLO using the optimization of uz(2 y) to amortize the difficulty of evaluating infinite number of ( with InfoNCE.
1 2j=1 e9o(m..35) + Ep(r)]eso(xx."")
Efficient implementations of FLO. A lingering concern is that the newly introduced us y) can incur extra computation overhead. This is not true, as we can maximally encourage parameter sharing by jointly model (1, y) and go (z. y) with a single neural network fy (z. y) X' x x + R with two output heads, i... [Mg. 9il ( (r,, V:). Consequently, while FLO adopts a dual crities 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 Appendix F.6).
To further enhance the computation efficiency, we consider massively parallelized bi-linear critic design that uses all in-batch samples as negatives. Let (0 v) T- T-(ho(z), h(y)), where h : X + SP and h ) + S² are respectively encoders that map data to unit sphere S embedded in R (a, b) = aT b is the inner product operation, and T >> the inverse temperature parameter. Thus the evaluation of the Gram matrix G = T ((X)))())........ where [ Y) E (dx is a mini-batch of K K-paired samples and go (2t, Vy) = Gij. 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., 115
We can simply model the PMI critic u(x, y) MLP(h(>) M(y)), whose computation cost 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 broader context to Appendix A, H.
From log-partition approximation to MI bounds. To embrace more holistic understanding, we list popular variational MI bounds together with our FLO in Table I and visualize their connections in Figure I 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, Inf oNCE) and the exponential- family (NVJ, TUBA, FLO). In the log-family, DV and Inf oNEE are multi-sampi estimators that leverage
$ As an important note to the community, most open source implementations for the bilinear contrastive loss have mechanically implemented ((-)) following the practice from pioneering contrastive learning studies, which is numerically unstable compared to our parameterization -(:-)) proposed here.
5

---

Table 1: Comparison of popular variational MI estimators. Here g(z, y), u(x, 3) and u(x) are variational functions to be optimized, a(u) = [+exp(-u) is the Sigmoid function, (1). mi denotes exponential average of function f(() with decay parameter 1) (0, 1), and [0. 1] is the balancing parameter used by a-InfoNCE trading off bias and variance between InfoNCE and TUBA. we use (Z4) Me) denote positive samples from the joint density p(x, y), and (xt My) or (4r4, M/4 denote negative samples drawn from the product of marginal p(z)p(v). In context, Ve and ye 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.
Name
Objective
Bias
Var.
Converge
(Ft. Ub) Rd p(a, (2 v() ~ p(x)p(y), ma v1.)) _ #a
k= exp(o(<r, mv.)) (1 o)) exp(a(z))
DV 19
ofar. 91) log Mng(Ei (. g( u.) log (og([exp((((x..)))
high
high
no
MINE
low
high
no
NWJ
glas, ui) explo(es vs)
low
high
no
39
9" (1). explo" (1) -1)
low
high
no
May min (log (ogo(( will log ol ((21 8:))
TUBA 62
g( ud) + u(a,) + 1 M3) u(x,))
low
high
no
InfoNCE 57
g(a..u)) log(L;, expig(as, M; K)
high
low
no
a- InfoNICE [62
g( v.) g(uz, w/) - log Ma (1r, gr + log(ma, (og(ma..(44.. v())
no
InfoNCE interpolates between biw- high-var (a 1. NWJ) high-bias low-var (a + 0, InfoNCE)
FLO (ours)
u((;; exp( u( 1)) g( My) g( Mi]) low
moderate
yes
direct Monte- Carlo estimates 2 for log (xx), 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 finite-sample average due to Jensen's 's inequality (i.e., log = log [ 2 Ellog 2).. In contrast, the objective of pppooential-famill estimators do not inequality: NWJ directly applies the Fenchel dual of f-divergence for MI (56), while TUBA exploits this inequality to compute the log partition log Z(x) = log [p(r>)exp((( (x, }]]. Motivated from contrastive view, our FLO applies the Fenchel-Legent inequality to the log- pog-partition of contrast scores.
involve such log- trannfoommaiinn since they can all be derived from the Fenchel-
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 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-fan oog-family estimators can be considered as multi-sapple, single-stage generalization of JSD, However, the DV objective can go unbounded thus resulting in 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-famin MI estimation, and bridges existing bounds: the PMI normalizer exp(-u(ar; y)) is 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 flexible framework that not only recovers existing bounds, but also derives novel bounds such as
Σ, explec((e......))) TFOV A StopGrad (((x,, y1)))] Stopprad]E, exp (2....
1.
(11)
Recall the optimal 9¹ (x, 9) = logp(xly) + c(x) and ux(,,, us to parameterize ((x,y) in the form of -go(a, v) + (2)),, where so (xx) models the arbitrary drift c(x), and this recovers the TUBA bound. Additionally, note that (i) fixing either of
(x, y) log p(a.))
(()))(y)) which motivates

---

and g, and optimizing the other also gives valid lower bound to MI; and (ii) a carefully chosen multi-input u( (1, }1 can be computationally appealing. As a concrete example, if we set uø to uo(((a ys)1) log k E, 20 and update uo(x, y) while artificially keeping the critic go(2,,) fixed] then FLO falls back to DV. Alternatively, we can consider the Fenchel dual version of it: using the same multi-input uo( (Ia, (s)]) above, treat Ui as fixed and only update go. and this gives us the novel MI objective in (TT). we call it Fenchel-Donsker- -araad-ana (FDV) estimator.
2.4 Gradient and convergence analysis of FLO
In this section, we will establish that FL.0 better optimizes the MI because its gradient is more accurate than competing variational bounds such NWJ and TUBA; also, we provide the first 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 [62]. For instance, with the same critic go we have Ino I ITUBA S JuBa. So if we can show VolFLO = VeluBa then FLD is better optimized. To simplify notations, we denote c0(2, U. V.V') _ go v/) go (z. and (o v) - 1 1/Ep(w) (c )), and we can easily verify
Ep(rr) Vo { (ecoeesss) Vo (E((z,)) }
VEo(,, y) (Eg(z, = Vo Volog& Eo
(12)
Eo(1.
Since for fixed go the corresponding optimal uz(x, v) maximizing I 9o) - (uo(z,y) + Rep(rr) le (o(x.))) (901)] is given by u)(x,y) = log Dog (e -
loog (1r, v) (using (4)) we see that the term e-us(e.)) is essentially optimized to approximate (z, y)- To emphasize this point, we now write E0 y) & e-us(aw) When this approximation is
sufficiently accurate (i.e., Eo se Ee), we can see that VILO approximates Vluaa as follows
VoL!FLo(up: [o)} -Ery (--o(xxe0 E. [Voes((x
Euy
Ea(a,v) Vo log Ex w) (o
(13)
2 Eay Vo log & (s, v)] = Vo { (Rp( flog &a (1, [x], = Vo (904
We can prove FLO will converge under much weaker conditions, even when this approximation Enc ((2, y) is rough. The intuition is as follows: in (13). the term Eo, 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 (m be the stochastic Bobbins-Mome sequence of < co. If & Eor is bounded between [a, b] (0 < a < b < oo), then under the stochastic gradient descent scheme described in Algo- rithm 11 0r converges to a stationary point of LUAA (go) with probability 1, i.e., lim || V (go.) )l O. Additionally assume JUBA is convex with respect to 8, then FLO converges with probability 1 to the global optimum H* of TUNA from any initial point fy-
Algorithm FLO
Empirical data pa {(x,, y1 91) Model parameters 1 = (0, () for t 1,2, do
learning rates: Σ, Blmr] ao and Σ, E(
Sample i,j rid
~ [n]
us unl ug 46). go (1)
Py gol [T; go
F un + expl Uii + 9h) - 9a))
V, V. m Fg
end for
Importantly, this marks the first SGD convergence analyses for variational MI estimators. The convergence analyses for MI estimation is non-trivial and scarce even for those standard statistical estimators [59. 24 64). For variational MI bounds, existing convergence analyses only apply to oracle estimators but not practical estimators discussed here [35] 70). The lack of convergence guarantees has led to 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 (9)). Our work establishes family of variational MI estimators that provably converges, contribution we consider significant as it fills 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.
'hht is say go in up is independent copy of go.
7

---

TUB4
12.
10.0
5.0
0.6 0.8
1.0
0.0
0.2
0.4
0.6 0.8
1.0
0.0
0.2
0.4
Figure 3: Bias- vaas-vaieenee plot for popular variational MI bounds with the 10-D Gaussians. Estimators that are more concentrated around the dashed line is considered better (low-bias, low-variance). the more challenging high-MI regime, FLO shows clear advantage over competing alternatives, where FLO pays less price in variance to achieve even better accuracy when tight estimation is impossible.
Nalsy Linear Regression Model
Norllieer Phermecokinelic Model
Policy Design for Epidemics
t0us) Infonce
1.5
3-
Design ML
40-
10ums)
3.5
30
1.2
2003
60:0
1003
12099
1003
220c
2000
40.00
$3.00
2043
49:00
88:00
20000
26 530
756
20:00
1224
1500
1726
2008
Iterations
Ieee nteratioos
Iterations
Figure 4: Bayesian Optimal Experiment Design results. FLO consistently performs best, demonstrat- ing superior strength in learning efficiency and robustness. NWJ takes the runner-up, but it has larger variance and is sensitive to network initializations. InfoNCE is less competitive due to low sample inefficiency, but its smaller variance helps in the more challenging dynamic case.
3 Experiments
We consider an extensive range of tasks to validate FLO and benchmark against state-of-the- art solutions. To underscore the practical significance of MI in efficient machine learning, we demonstrate example applications from data collection (in statistical parlance, experimental design), self-supervised pre-training, meta/transfer-leamin results in the main text, and defer ablation studies and details of our experimental setups to the Appendix. Our code is available from https::/ //ggthub......t// cww qingguo666/FLO All experiments are implemented with PyTorch.
Limited by space, we present only the key
Comparison to baseline MI bounds. We start by comparing FLO to the following popular competing variational estimators: NWJ, TUBA, and Inf INN.. We use the bilinear critic implementation for all models which maximally encourages both sample efficiency and code simplicity, and this strategy does perform best based on our observations. We consider the synthetic benchmark from 162), where (X E [dd Y E Id)) is jointly standard Gaussian with diagonal cross-correlation parameterized by & E [0, 1). We report d 10 and p E [0, 0.99] here (other studies only report p up to 0.9, which is less challenging.), providing a reasonable coverage of the range of MI one may encounter in empirical settings.
True MI
kde
15-
Knnie
MI (10D)
Knnie Revised
10-
npeee
FLO (Ours)
0.0
0.8
Figure 5: FLO compares favor- ably to classical MI estimators.
To focus on the biri-variance trade-off, we plot the decimal quantiles in addition to the estimated MI in Figure 3 where FLO significantly outperformed its variational counterparts in the more challenging high-MI regime. In Figure [] show FLO also beats classical MI estimators
45
80 25 In
the Appendix H, we further discuss recent works on parametric estimators [7 10 and alternative information metrics (85).
Bayesian optimal experiment design (BOED). We next direct our attention BOED, a topic of significant interest shared by the statistical and machine learning communities [[1
82 36 23 The
performance of machine learning models crucially relies on the quality of data supplied for training, and BOED is principled framework that optimizes the data collection procedure (in statistical parlance, conducting experiments) [22]. Mathematically, let be the data to be collected, 0 be the parameters to be inferred, and d be the experiment parameters the investigator can manipulate (a.k.a,
8

---

SIR Epidemic Trajectory
Policy Design Epidemics
106
4.4
Ground
165-
SiR Populations
Design lime
TUBA INONCC
422
6o-
thy Irfocted Recoverel
148
4.0
100
FLO Design
3.5
3.6-
Design Sequence
Figure 6: Diagnosis of learned sequential designs. The disease surveillance windows designed by FLO makes more sense: measures more frequently as infection spikes. and more sparsely when the pandemic slowly fades. The estimated parameter posterior (right) is consistent with the ground truth.
the design parameters), BOED tries to find the optimal data collection procedure that is expected to generate data that is most informative about the underlying model parameters, i.e., solves for arg max, [xx; 0; d). In this study, we focus on the more generic scenario where explicit likelihoods are not available, but we can still sample from the data generating procedure [43 44).
We consider three carefully-selecte
models from recent literature for their progressive practical significance and the challenges involved [21 41, 42] static designs of (1) simple linear regression model and (ii) complex nonlinear pharmacokinetic model for drug development; and the dynamic policy design for (iéi) epidemic disease surveillance and intervention (e.g., for Covid-19 19 modeling). Designs with higher MI are more favorable, because it implies the data carries more information. In Figure we compare design optimization curves using different MI optimization strategies, where FLO consistently leads. Popular NWJ and InfoNCE reports different tradeoff's that are less susceptible to FLO. We also examine the FLO predicted posteriors and confirm 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 @leff. 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. Briefly, for arbitrary loss (0, y). where t is the task identifier and 9 = f(z) is the prediction made by the model, we denote the risk by (f( = Ep. (f v)]. Denote R(f) Ex-p(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-theorett generalization theories [84), we derived a novel, principled objective
Lxeta-FLO(f) = R(f) + X/ I (Dr Ex),
(14)
where À is known given the data size
Table Multi-view representation learning on Cifar
and loss function, (De, Ex) are respec-
Model
InfoNCE
SpecNCE (34]
FLO
FDV
tively data and task embeddings for training data, which for the first time lifts contrastive lcarning to the task and data distribution level. Our reason- ing is that FLo theoretically bounds R(f) from above, and it is rel- atively sharp for being dependeen... We give more information on this in the Appendix and defer a full exposi- tion to dedicated paper due inde- pendent interest and space limits here. Note other MI bounds are not suitable for this task due to resource and vari-
MI 5.73 + .07
4.76 + .08
5.83 + .08
5.93 + .08
Feeeshot Adaptation Tain Dynamics
Few-shet Adaptation
rogpes
Mesaruo
3.5-
no-
2.5
1.5-
Greand Truth Support Set
200
400
600
800
1000
2
Retaations bx100)
Figure 7: Few-shot adaptation with Meta-PLO.
"Note SpecNCE does explicitly larget mutual information
ance concerns. Figure we show Meta-FLO wins big over the state-of-the-aa model agnostic meta-learning (MAML) model on the regression benchmark from [20].

---

Salf supervised Leaminn
9.5
Topt Test Recurecy
26-
Citar 10,10, ResMer50
Kumaal Information
5.0
inpreserration
8.0
35
The bound
Gooin be
30
78-
25
PADCEE
76
20
InfoICE
30
100
150
200
350
350
400
160
240
400
Traniing Spacis
Mutual Information
Training Epocht
Figure 8: Sample efficiency Figure 9: Representation Figure 10: FlatNCE better optimizes comparison for SimCLR and MI strongly correlates the true MI with the same mini-batch FlatCLR on Cifar10.
with performance.
size. (Cifar10 SSL training)
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-specific problems such as training diagnosis and low- now-precision numerical overflow 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 our performance criteria.
We compare the effectiveness of the InfoNCE-based SimCLR framework [5] to our FLO-based alternatives. To ensure fair comparison, we have used the FDV variant defined in Eq. [ 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-intuitie intuitively, the second term in Eq. IIT contributing all the learning signal constant one in value (i.e., being flat). In Figure and 9 we show our new model FlatCLR shows superior sample efficiency compared to the SOTA SimCLR (a 8x boost for the same performance, FlatCLR-32 versus SimCLR-256) This result is significant because SimCLR's ' crucial reliance on large-batch training is a well-known limitation [34 86 46). Figure TO 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-trus -tuth MI which can be explained by its robustness to the numerical overflow issue (see [12] for details). Further comparisons on the ground-truth MI estimation with different estimators can be found in Table 2
4 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, causal inference [48], clustering (31], and in particular, data efficient learning. We are also interested in further exploring the deep connections
between contrastive MI estimation and generative modeling [16] 73 14
18
74
Acknowledgements
The authors would like 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 's Commonwealth Cyber Initiative (CCI) Al 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 OCI-1548562 (78) 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
1 "Code available httpt:://gtthuub...... com/ humm/Junaa-aaa Saanya-Chen/FFan Chen/FlatCii.
*Ground-trthh MI approximated by InfoNCt using very large negative sample pool (100x mini-batchi
10

---

References
[1] Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing broken ELBO. In ICML, pages 159-168, 2018.
[2] Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adver- sarial networks. In ICLR, 2017.
[3] Francis R Bach and Michael I Jordan. Kernel independent component analysis. Journal of Machine Learning Research, 3(Jul):1-48, 2002.
[4] David Barber and Felix Agakov. The IM algorithm: variational approach to information maximization. NIPS, 16:201, 2004.
[5] David Barber and Felix Agakov. Information maximization in noisy channels: A variational approach. NIPS, 16, 2003.
[6] Roberto Battiti. Using mutual information for selecting features in supervised neural net learning. IEEE transactions on Neural Networks, 5(4):537-550, 1994.
[7] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In ICML, 2018.
[8] Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on Pattern Analysis and Machine Intelligence, 35(8):1798- 1828, 2013.
[9] James O Berger. Statistical decision theory and Bayesian analysis. Springer Science & Business Media, 2013.
[10] Rob Brekelmans, Sicong Huang, Marzyeh Ghassemi, Greg Ver Steeg, Roger Baker Grosse, and Alireza Makhzani. Improving mutual information estimation with annealed and energy-based bounds. In ICLR, 2022.
[11] Kathryn Chaloner and Isabella Verdinelli. Bayesian experimental design: A review. Statistical Science, pages 273-304, 1995.
[12] Junya Chen, Zhe Gan, Xuan Li, Qing Guo, Liqun Chen, Shuyang Gao, Tagyoung Chung, Yi Xu, Belinda Zeng. Wenlian Lu, et al. Simpler, faster, stronger: Breaking the log-K K curse on contrastive learners with flatnce. arXiv preprint arXiv22107.01152, 2021.
[13] Junya Chen, Danni Lu, Zidi Xiu, Ke Bai, Lawrence Carin, and Chenyang Tao. Variational inference with Holder bounds. arXiv preprint arXiv:2111.02941 2021.
[14] Liqun Chen, Chenyang Tao, Ruiyi Zhang, Ricardo Henao, and Lawrence Carin Duke. Vari- ational inference and model selection with generalized evidence bounds. In ICML, pages 893-902 PMLR, 2018
[15] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In ICML, 2020.
[16] Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Info- GAN: Interpretable representation learning by information maximizing generative adversarial nets. In NIPS, 2016.
[17] Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, and Lawrence Carin. CLUB:
A contrastive log-ratio upper bound of mutual information. In ICML. 2020.
[18] Bo Dai, Hanjun Dai, Niao He, Weiyang Liu, Zhen Liu, Jianshu Chen, Lin Xiao, and Le Song. Coupled variational bayes via optimization embedding. NeurlPS, 31, 2018.
[19] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time. Communications on Pure and Applied Mathematics, 36(2):183-212. 1983.
11

---

[20] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adapta- tion of deep networks. In ICML, 2017.
[21] Adam Foster, Desi R Ivanova, Ilyas Malik, and Tom Rainforth. Deep adaptive design: Amortiz- ing sequential bayesian experimental design. In /CML, 2021.
[22] Adam Foster, Martin Jankowiak, Eli Bingham, Paul Horsfall, Yee Whye Teh, Tom Rainforth, and Noah Goodman. Variational bayesian optimal experimental design. In NeurlPS, 2019.
[23] Adam Foster, Martin Jankowiak, Matthew O' O'Meara, Yee Whye Teh, and Tom Rainforth. A unified stochastic gradient approach to designing bayesian-optimal experiments. In AISTATS, 2020.
[24] Shuyang Gao, Greg Ver Steeg, and Aram Galstyan. Efficient estimation of mutual information for strongly dependent variables. In A/STATS, 2015.
[25] Weihao Gao, Sewoong Oh, and Pramod Viswanath. Demystifying fixed k-ncarest neighbor information estimators. IEEE transactions on Information Theory, 64(8):5629-5661. 2018.
[26] Charles Geyer. On the convergence of monte carlo maximum likelihood calculations. Journal of the Royal Statistical Society: Series (Methodological), 56(1):261-274 1994.
[27] Arthur Gretton, Ralf Herbrich, Alexander Smola, Olivier Bousquet, Bernhard Schölkopf, et al. Kernel methods for measuring independence. Journal of Machine Learning Research, 2005.
[28] Arthur Gretton, Ralf Herbrich, and Alexander J Smola. The kernel mutual information. In ICASSP, 2003.
[29] Jean- Bant-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent: A new approach self-supervised learning. In NeurIPS, 2020.
[30] Aditya Grover and Jure Leskovec. node2vec: Scalable feature learning for networks. In SIGKDD, 2016.
[31] Qing Guo, Xinwei Deng, and Nalini Ravishanker. ietoccattio-bbssee optimal subpopulation selection for multivariate data. In Innovations in Multivariate Statistical Modeling, pages 3-21. Springer, 2022.
[32] Umang Gupta, Aaron Ferber, Bistra Dilkina, and Greg Ver Steeg. Controllable guarantees for fair outcomes via contrastive information estimation. arXiv preprint 1xXiv2201..01008 2021.
[33] Michael Gutmann and Aapo Hyvärinen. Noise-comtrastive estimation: A new estimation principle for unnormalized statistical models. In AISTATS, 2010.
[34] Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self- supervised deep learning with spectral contrastive loss. NeurlPS, 2021.
[35] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. CVPR, 2020.
[36] José Miguel Hernández-Lobato, Matthew W Hoffmman, and Zoubin Ghahramani. Predictive entropy search for efficient global optimization of black-box functions. In NIPS, 2014.
[37] Geoffrey Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation, 14(8): 2002.
[38] Jean-Baptiste Hiriart- and Claude Lemaréchal. Fundamentals of convex analysis. Springer Science & Business Media, 2012.
[39] R Devon Hjelm, Alex Fedorov, Samuel Lavoie- Marchildan Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In ICLR, 2019.
12

---

[40] Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(Apr):695-709, 2005.
[41] Desislava Ivanova, Adam Foster, Steven Kleinegesse, Michael U Gutmann, and Thomas Rainforth. Implicit deep adaptive design: Poliey-based experimental design without likelihoods. NeurIPS, 2021.
[42] Steven Kleinegesse, Christopher Drovandi, and Michael Gutmann. Sequential bayesian experimental design for implicit models via mutual information. Bayesian Analysis, 1(1):1-30, 2021.
[43] Steven Kleinegesse and Michael U Gutmann. Bayesian experimental design for implicit models by mutual information neural estimation. In ICML, 2020.
[44] Steven Kleinegesse and Michael U Gutmann. Gradient-based bayesian experimental design for implicit models using mutual information lower bounds. arXiv preprint arXiv:2105.04:3
2021.
[45] Alexander Kraskov, Harald Stögbauer, and Peter Grassberger. Estimating mutual information. Physical review E, 69(6):0661 068.38,66688800 2004.
[46] Kyungmin Lee and Jinwoo Shin. RényiCL: Contrastive representation learning with skew rényi divergence. In NeurIPS, 2022.
[47] Ralph Linsker. Self-organization in perceptual network. Computer, 21(3):005-117, 1988.
[48] Danni Lu, Chenyang Tao, Junya Chen, Fan Li, Feng Guo, and Lawrence Carin. Reconsidering generative objectives for counterfactual reasoning. NeurIPS, 33:21539-21553, 2020.
[49] Zhuang Ma and Michael Collins. Noise contrastive estimation and negative sampling for conditional models: Consistency and statistical efficiency. arXiv preprint 18Xiv:1809.01812. 2018.
[50] David JC MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
[51] Frederik Maes, Andre Collignon, Dirk Vandermeulen, Guy Marchal, and Paul Suetens. Mul- timodality image registration by maximization of mutual information. IEEE transactions Medical Imaging, 16(2):187-198, 1997.
[52] David McAllester and Karl Stratos. Formal limitations on the measurement of mutual informa- tion. arXiv preprint 1rX2v.181..0425,, 2018.
[53] Andriy Mnih and Koray Kavukcuoglu. Learning word embeddings efficiently with noise- contrastive estimation. In NIPS, 2013.
[54] Radford M Neal. Annealed importance sampling. Statistics and computing, 11(2):125-139, 2001.
[55] XuanI Long Nguyen, Martin Wainwright, and Michael Jordan. Estimating divergence func- tionals and the likelihood ratio by convex risk minimization. IEEE transactions on Information Theory, 56(11)584775861. 2010.
[56] Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-GAN: Training generative neural samplers using variational divergence minimization. In NIPS, 2016.
[57] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1807:0807.0377 2018
[58] Stephanie Palmer, Olivier Marre, Michael J Berry, and William Bialek. Predictive information in sensory population. Proceedings of the National Academy of Sciences, 112(22):6008-691) 2015.
[59] Liam Paninski. Estimation of entropy and mutual information. Neural computation, 15(6): 1253, 2003.
13

---

[60] Fernando Pérez-Cruz. Estimation of information theoretic measures for continuous random variables. In NIPS, 2008.
[61] Josien PW Pluim, JB Antoine Maintz, and Max Viergever. nnttual-infrrmation-- based registration of medical images: survey. IEEE transactions on Medical Imaging, 22(8):986- 1004, 2003.
[62] Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In ICML, 2019.
[63] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. arXiv preprint (0:::2203.000000 2021.
[64] Tom Rainforth, Rob Cornish, Hongseok Yang, Andrew Warrington, and Frank Wood. On nesting monte carlo estimators. In International Conference on Machine Learning, 2018.
[65]
Linda E Reichl. A modern course in statistical physics. John Wiley & Sons, 2016.
[66] David Reshef, Yakir A Reshef, Hilary K Finucane, Sharon R Grossman, Gilean Mc Peter J Turnbaugh, Eric S Lander, Michael Mitzenmacher, and Pardis C Sabeti. Detecting novel associations in large data sets. Science, 334(6062):1518-- 151--124, 2011.
[67] Claude E Shannon. A mathematical theory of communication. The Bell system technical journal, 27(3):379-422 1948.
[68] Ravid Shwartz- Ziiv and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv preprint arXiv:1703.008/00 2017.
[69] Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual informa- tion estimators. In ICLR, 2020.
[70] Sreejith Sreekumar and Ziv Goldfeld. Neural estimation of statistical divergences. Journal of Machine Learning Research, 23(126))1-75, 2022.
[71] Taiji Suzuki, Masashi Sugiyama, Jun Sese, and Takafumi Kanamori. Approximating mutual information by maximum likelihood density ratio estimation. In New challenges for feature selection in data mining and knowledge discovery, 2008.
[72] Chenyang Tao, Liqun Chen, Shuyang Dai, Junya Chen, Ke Bai, Dong Wang, Jianfeng Feng, Wenlian Lu, Georgiy Bobashev, and Lawrence Carin. On Fenchel mini-max learning. In NeurlPS, 2019.
[73] Chenyang Tao, Liqun Chen, Ricardo Henao, Jianfeng Feng, and Lawrence Carin Duke. Chi- square generative adversarial network. In ICML, pages 4887-4896. PMLR, 2018.
[74] Chenyang Tao, Shuyang Dai, Liqun Chen, Ke Bai, Junya Chen, Chang Liu, Ruiyi Zhang, Georgiy Bobashev, and Lawrence Carin Duke. Variational annealing of GANs: A langevin perspective. In ICML, pages 6176-6185. PMLR, 2019.
[75] Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. arXiv preprint arXiv:1906.055 058499 2019.
[76] Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Information Theory Workshop (ITW), pages 1-5. IEEE, 2015.
[77] Kari Torkkola. Feature extraction by non-parametric mutual information maximization. Journal of machine learning research, 2003.
[78] John Towns, Timothy Cockerill, Maytal Dahan, Ian Foster, Kelly Gaither, Andrew Grimshaw; Victor Hazlewood, Scott Lathrop, Dave Lifka, Gregory D Peterson, et al. Xsede: accelerating scientific discovery. Computing science & engineering, 16(5):62-74, 74, 2014.
[79] Michael Tschannen, Josip Djolonga, Paul Rubenstein, Sylvain Gelly, and Mario Lucie. On mutual information maximization for representation learning. ICLR, 2020.
14

---

[80] Greg Ver Steeg and Aram Galstyan. ceemmation-thooreee measures of influence based on content dynamics. In Proceedings of the sixth ACM international conference on Web search and data mining, pages 3-12, 2013.
[81] Liangjian Wen, Yiji Zhou, Lirong He. Mingyuan Zhou, and Zenglin Xu. Mutual information gradient estimation for representation learning. In ICLR, 2020.
[82] CF Jeff Wu and Michael S Hamada. Experiments: planning, analysis, and optimization, volume 552. John Wiley & Sons, 2011.
[83] Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non-parametric instance discrimination. In CVPR, 2018.
[84] Aolin Xu and Maxim Raginsky. Information-theorrticc analysis of generalization capability of learning algorithms. In NIPS, 2017.
[85] Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints. In ICLR, 2020.
[86] Zhuoning Yuan, Yuexin Wu, Zi-Hao Qiu, Xianzhi Du, Lijun Zhang, Denny Zhou, and Tianbao Yang. Provable stochastic optimization for global contrastive learning: Small batch does harm performance. In ICML, 2022.
[87] Sue Zheng, Jason Pacheco, and John Fisher. A robust approach to sequential information theoretic planning. In /CML, 2018.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [N/A]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] They are the Supplementary Material.
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] They are summarized in the Supplementary Material.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] Smaller variance is the highlight this paper.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
15

---

(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsoureing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "vV0737mBlrd", "reviewer": "Reviewer_NZCL", "summary": "The paper proposes a new bound of mutual information (MI) based on Fenchel-Legendre Optimization (FLO).  The bound has a number of beneficial properties such as tightness (it can approach the true MI), the empirical estimator is unbiased, and it can be estimated using samples from the model but does not require evaluation of model probabilities (i.e. it can be used for implicit likelihood models).  The paper discusses properties of the FLO estimator in relation to existing variational estimators and those based on the conjugate function.  Finally , the paper provides theoretical convergence guarantees along with empirical analysis (and ablation study) on several interesting models that are relevant in the literature on MI estimation and sequential experiment design. ", "questions": "* Proposition 2.2 states that FLO is a lower bound on UBA with the same critic.  If so, what is the benefit of FLO over UBA (the unbiased empirical estimator is obviously one, but are there others)?  Is this bound true for every $u_\\phi$ or only for an optimal one?\n\n* Several estimators (e.g. DV) are based on a similar conjugate duality construction to FLO.  Yet, the authors do not discuss the relationship between these estimators in great detail.  How do these estimators compare, e.g. for similar choices of the surrogate function (i.e. the critic in the FLO case)?\n\n", "limitations": "The authors clearly state limitations.  This reviewer does not see obvious negative societal impact of the work.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**STRENGTHS**\nThe paper is very well-written, the problem is well-motived, and the topic is timely and important.  The authors provide a solid foundation for the reader, along with all necessary context.  The paper is a good blend of theory and practical application.  The empirical analysis provides comparison to most relevant MI estimators and sufficiently supports the stated claims.  \n\nOne strength that receives less attention due to the focus on tightness (see critique below) is that the empirical estimator lacks finite sample bias.  Much of the literature on MI estimation focuses on properties of the analytic approximations / bounds, when in reality these bounds cannot be calculated exactly and require further estimation.  In most cases there do not exist unbiased estimators, for example in the UBA bound, and estimators are only asymptotically unbiased / consistent.  By contrast, FLO admits a straightforward empirical estimator with zero finite sample bias.  Discussion of this is provided in Zheng et al. (2018) and Rainforth et al. (2018) for the Barber and Agakov bound.\n\n**WEAKNESSES**\nThe empirical analysis in the main text is fairly superficial.  Results on the SIR model are particularly cursory, and are largely qualitative.  For example, the caption of Figure 6 states that FLO designs are somehow more useful because they are sampled more densely around the infection spike.  Yet the authors provide no comparison and the comment is debatable from visual inspection alone.  Additionally, the estimates provided in Figure 3 are mostly qualitative and do not seem to show a strong benefit of FLO compared to NWJ for example.\n\nThe paper leans heavily on tightness of the FLO estimator, yet this property does not hold in practice.  The discussion under Prop. 2.3 shows that tightness is achieved only under the optimal critic $g_{*}(x,y)$ which requires knowledge of the true posterior $p(x\\mid y)$.  In addition, the tightness property is not unique to FLO, despite what the authors suggest in Table 1.  Many of the provided estimators are also asymptotically tight or tight given knowledge of the true posterior.  For example, most variational MI bounds become tight when the variational distribution $q(x \\mid y) = p(x \\mid y)$ almost everywhere.  By Prop. 2.2 the UBA bound is also tight given the same optimal critic, and indeed FLO is in general a looser bound than UBA by this proposition.   Similar properties hold for DV and NWJ estimators as well as InfoNCE (i.e. Proposition 2.1).\n\nSome detailed comments below:\n* It is not clear what point the authors intend to convey with the section on unnormalized statistical modeling and its connection to MI estimation.  Essentially any posterior inference problem contends with the lack of an analytic partition function, and so it is not particular to this setting.\n* The statement in Prop. 2.3 is a bit terse.  At the very least authors should state that the estimator is tight at $(g_{*}, u_{g_{*}})$. \n* Line 266 : Change “direct out attention” to “direct our attention”\n\n**REFERENCES**\n\nRainforth, Tom, et al. \"On nesting monte carlo estimators.\" International Conference on Machine Learning. PMLR, 2018.\n\nZheng, Sue, Jason Pacheco, and John Fisher. \"A robust approach to sequential information theoretic planning.\" International Conference on Machine Learning. PMLR, 2018.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "uSoq1qRnVoa", "reviewer": "Reviewer_4a7X", "summary": "This paper introduces new type of contrastive mutual information estimator, which is tight and provides theoretical convergence. As the variance of previous estimator was big trouble for the accurate MI estimation, it provides new optimization technique with fenchel-legendre optimization. The relationship with the existing mutual information estimator has been nicely provided, which makes it easier to evaluate the introduced estimator with the fair manner.", "questions": "Q1. how this framework can be applied on the supervised contrastive learning? Additional works from supervised contrastive learning stated that it is connected with conditional mutual information maximization. I am curious that the mutual information estimation or maximization given the label can be conducted with FLO.\n", "limitations": "Stated in weakness session.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nTheoretical analyses : it provides the tightness and boundedness for the given estimator with nice derivation. It was easy to follow.\nExperiments : FLO significantly outperformed its variational counterparts in the more challenging 263 high-MI regime, which implies the robustness from the noisy environments.\n\nWeaknesses : Experiments for additional data-efficient learning is not provided. The author referred that efficient implementation for FLO can be provided, which makes it easier to apply FLO on data-efficient learning. However, it was not adequately managed in the manuscripts.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "JtU1azLAIhy", "reviewer": "Reviewer_FJTi", "summary": "The paper introduces FLO, a novel contrastive estimator for mutual information. FLO is inspired by a reexamined connection between mutual information and unnormalised statistical modelling, which yields a unified theoretical basis for several variational-based MI estimators. FLO is then proven to be tight, and to converge under SGD. Finally, several experiments confirm FLO's superior performance compared to previous estimators.", "questions": "My main suggestions for improvements would go for the experimental section:\n- running time comparisons\n- giving more context as to levels of performance for various tasks.", "limitations": "Limitations are addressed adequately", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is *very* well written, and as a consequence a pleasure to read, even for non-experts in the field of mutual estimation. Connections are explained clearly and the new insights and methods are well-motivated. The 'zoo' of estimators in Figure 1 is much appreciated.\n\nI can find precious little weaknesses in the paper (although I'm by no means an expert in this particular field). A few things:\n- some running time experiments contrasting the performance of the various estimators would be good, to verify the tractability claims\n- in the experiments, some explanation of how significant the gains are would be good, given that some of the tasks are not familiar to all readers\n- all told, given the amount of material in the paper (and consequently how much of it is left in the appendix), I wonder if a journal version would not do it better justice.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "J2XCnOznmF2", "reviewer": "Reviewer_gbJf", "summary": "The paper proposes a new mutual information estimator based on Fenchel-Legendre optimization (FLO). The paper describes existing related techniques, and then presents the derivation of the FLO estimator via convex duality. The FLO estimate is established as a lower bound on the true mutual information value, and it is shown that appropriate choices of the g and u functions in the variational bound retrieves the true value. The authors also provide a convergence to stationarity claim for SGD-based optimization of the FLO objective. Numerical experiments comparing the proposed approach to existing estimators is provided, showing good performance on the considered examples. Applications to Bayesian optimal experiment design and meta-learning are also presented.", "questions": "* The feasibility claims from Proposition 2.3 and Corollary 2.4 are trivial and have limited utility since the optimal functions depend on the population distribution, which is unknown in practice.  \n\n* The abstract and intro make it sound like a statistical performance analysis for FLO will be provided, but it is nowhere to be found. The short paragraph after Corollary 2.4 is quite laconic and surely doesn't qualify as such. If there is a rigorous statement the authors can provide, it should be clearly formulated, stated, and proven. The current presentation is misleading in this regard.\n\n* The convergence claim from Proposition 2.5 is local and should be framed as such from the get-go. The abstract and introduction make it sound like global guarantees would be provided, which is not the case and hence this presentation is misleading. The 2nd part of that statement that assumes convexity if I_{UBA} surely does not quality as a global guarantee since this convexity in the parameters is rarely available (e.g., when parametrizing the functions by a neural network, as often performed in practice).\n\n* I find the statement in Table 1 that FLO has low bias and moderate variance quite odd. What actual evidence (beyond the said heuristics) did the authors provide to support this bold claim?\n\n* In the `Efficient implementation of FLO' paragraph: jointly modeling u and g by a single neural network makes sense, but doesn't doing so typically require more parameters than approaches based on a single critic? If so, how can the authors say that this incurs no extra modeling cost? \n\n* The bi-linear critic design proposed in the next paragraph is indeed highly parallelizable, but how expressive is this model? Under what conditions can we expect that class to be able to approximate the optimal u and g functions (that depend on the joint population distribution)? Computational efficiency is important, but it is crucial to understand the tradeoffs involved.\n\n* I found the content of Section 2.3 full of jargon and quite confusing. I am not sure what the goal of this section is. The authors may consider omitting it to make space for a more comprehensive and rigorous treatment of the SGD convergence guarantees.\n\n* It is unclear how FLO helps mitigate challenges associated with large sample sets and high-dimensional data. The text suggests that this is the case on several occasions, but I have seen no results to that effect. In a similar vein, the authors should contrast their approach with the McAllester and Stratos formal limitations paper [44].\n\n* Equation (11) is unclear; in particular, StopGrad is never defined.", "limitations": "I don't think that limitations of the approach were sufficiently discussed. The paper advocates quite strongly that FLO is superior to existing approaches but provided evidence to that effect is at best partial. As mentioned above, these claims would have been easier to digest if supporting theory and formal guarantees was provided. Clarity on the limitations of the approach and/or when alternative methods are preferable would be appreciated as well.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed estimation approach is interesting and may very well have merits, but in my opinion, the paper does not provide sufficient concrete evidence to support that. I liked the overview of existing/past techniques, which was comprehensive and informative, and think that the FLO approach was clearly described. Beyond that, however, while there is a considerable amount of discussion and heuristic arguments for why FLO is favorable to alternative methods, these lack rigorous supporting evidence. In various cases in the past, an estimator of mutual information was proposed and argued to drastically improve upon past approaches only for the next paper to provide examples on which it fails. The best approach to avoid such pitfalls is to couple a proposed estimator with formal performance guarantees (e.g., consistency, error bounds, sample complexity, etc.), but the paper comes short in that regard. I also found the paper to be quite heavy in jargon and at times sloppy in its mathematical notation, which are aspects the authors should work to improve on. Specific concerns/comments/questions are provided in the next section. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept