I MPROVED M UTUAL I NFORMATION E STIMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose a new variational lower bound on the KL divergence and show that the Mutual Information (MI) can be estimated by maximizing this bound using a witness function on a hypothesis function class and an auxiliary scalar variable.
If the function class is in a Reproducing Kernel Hilbert Space (RKHS), this leads to a jointly convex problem.
We analyze the bound by deriving its dual formulation and show its connection to a likelihood ratio estimation problem.
We show that the auxiliary variable introduced in our variational form plays the role of a Lagrange multiplier that enforces a normalization constraint on the likelihood ratio.
By extending the function space to neural networks, we propose an efﬁcient neural MI estimator, and validate its performance on synthetic examples, showing advantage over the existing baselines.
We then demonstrate the strength of our estimator in large-scale self-supervised representation learning through MI maximization.
I NTRODUCTION
Mutual information (MI) is an ubiquitous measure of dependency between a pair of random variables, and is one of the corner stones of information theory.
In machine learning, the information maximization principle for learning representation from unlabeled data through self-supervision (Bell &
Sejnowski, 1995) motivated the development of many MI estimators and applications (Hjelm et al., 2019; Noroozi & Favaro, 2016; Kolesnikov et al., 2019; Doersch et al., 2015; van den Oord et al., 2018b;
Hu et al., 2017).
The information bottleneck (Tishby et al., 1999; Kolchinsky et al., 2017) is another principle that triggered recent interest in mutual information estimation.
MI is also used to understand the information ﬂow in neural networks, in learning clusters (Krause et al., 2010) and in regularizing the training of Generative Adversarial Networks (GANs) (Chen et al., 2016).
In many of those machine learning applications and other scientiﬁc ﬁelds, one has to estimate MI given samples from the joint distribution of high dimensional random variables.
This is a challenging problem and many methods have been devised to address it.
Since MI is deﬁned as the KullbackLeibler (KL) divergence between the joint distribution and the product of marginals, one can leverage non parametric estimators of f-divergences (Nguyen et al., 2008;
Nowozin et al., 2016; Sriperumbudur et al., 2009).
Speciﬁcally of interest to us is the Donsker-Varadhan (DV) representation of the KL divergence (Donsker & Varadhan, 1976) that was used recently with neural networks estimators (Belghazi et al., 2018; Poole et al., 2019).
Other approaches to estimating the MI are through ﬁnding lower bounds using variational Bayesian methods (Alemi et al., 2016; 2017; Barber & Agakov, 2003;
Blei et al., 2017), as well as through geometric methods like binning (Kraskov et al., 2004).
In this paper we propose a new estimator of MI that can be used in direct MI maximization or as a regularizer, thanks to its unbiased gradients.
Our starting point is the DV lower bound of the KL divergence that we represent equivalently via a joint optimization that we call η -DV on a witness function f and an auxiliary variable η in Section 2.
In Section 3, we show that when the witness function f is learned in a Reproducing Kernel Hilbert Space (RKHS) the η -DV problem is jointly convex in both f and η .
The dual of this problem sheds the light on this estimator as a constrained ratio estimation where η plays the role of a Lagrange multiplier that ensures proper normalization of the likelihood ratio.
We also show how the witness function can be estimated as a neural network akin to Belghazi et al.
We specify our estimator for MI in Section 4, and show how it compares to alternatives in the literature (Nguyen et al., 2008; Belghazi et al., 2018; Poole et al., 2019).
The experiments are presented in Section 5.
On synthetic data, we validate our estimators by estimating
MI on Gaussian variables and by regularizing GAN training as
Chen et al.
(2016) real data,
[Table]
/Z+YuLMDppZhRcncxmczGfh89n83YvR0atuanvoMXqCxihEh+gIvUXHaIGY98H77H31vvlf/O/+D/9na/W9LucR6oX/6x/Q6ftB</latexit>
[ours] [ours] we explore our estimator in deep MI maximization for learning representation from unlabeled data.
Figure 1 shows an overview of all the bounds and related MI estimators discussed in this paper.
L OWER B OUNDS ON KL D IVERGENCES AND M UTUAL I NFORMATION
Consider two probability distributions P and Q , where P is absolutely continuous w.r.t.
Q .
Let p and q be their respective densities deﬁned on X ⊂ R d .
Their KL divergence is deﬁned as
[Equation]
We are interested in the MI between two random variables X , Y where X is deﬁned on X ⊂ R d x , and Y on Y ⊂ R d y .
Let p xy be their joint densities and p x , p y the marginals of X and Y respectively.
The MI is deﬁned as follows:
[Equation]
which is the KL divergence between the joint density and the product of marginals.
Non-parametric estimation of MI from samples is an important problem in science and machine learning.
In what follows, we review variational lower bounds on KL to enable such estimation.
Variational Characterization of KL divergence .
Let H be any function space mapping X to R .
The ﬁrst variational characterization of the KL divergence goes back to Donsker & Varadhan (1976):
[Equation]
where the equality holds if and only if (iif) f ∗ = log( p/q ) ∈ H .
We refer to this bound as the DV bound .
The second variational representation was introduced in Nguyen et al.
; Nowozin et al.
and derived through convex duality to be ﬁnally stated as follows:
[Equation]
with equality iif f ∗ = log( p/q ) ∈ H .
We call this bound the f -div bound (as in f -divergence).
From Eq.
2 and Eq.
3 we see that the variational bounds are attempting to estimate the log-likelihood ratio f ∗ = log( p/q ) , and the tightness of the bound depends on the representation power of H .
In order to compare these two lower bounds, observe that log( t ) ≤ t − 1 , t > 0 .
Therefore log
[Equation]
≤ E x ∼ Q e f ( x ) − 1 which means that for any function space H we have:
[Equation]
from which we conclude that the DV bound is tighter than the f -div bound.
Now, given samples { x i , i = 1 . . .
N, x i ∼ P } , { y i , i = 1 . . .
N, y i ∼ Q } , estimating the KL diver- gence can be done by computing the variational bound from Monte-Carlo simulation.
Speciﬁcally, for the DV bound we have the following estimator:
[Equation]
Note that the Mutual Information Neural Estimator (MINE) (Belghazi et al., 2018) considered the above formulation with the hypothesis class H being a neural network.
For the f -div bound we have similarly the following estimator:
[Equation]
for which Nguyen et al.
(2008) introduced and studied a convex estimator with H being an RKHS.
While DV bound is tighter than the f -div bound, in order to learn the function f using stochastic gradient optimization, f -div is a better ﬁt because the cost function is linear in the expectation, whereas in the DV bound we have a log non-linearity being applied to the expectation.
This nonlinearity introduces biases in the mini-batch estimation of the cost function as noted in Belghazi et al.
In the following, we show how to alleviate this problem and remove the non-linearity at the price of an additional auxiliary variable that will enable better optimization for the DV bound.
An η -trick for the DV Bound.
We start with the following elementary Lemma, that gives a variational characterization of the log .
All proofs are given in the Appendix A.
Lemma 1.
Let x > 0 , we have:
[Equation]
Using Lemma 1 we can now linearize the log in the DV bound.
Lemma 2 ( η − Donsker-Varadhan) .
Let H be any function space mapping X to R :
[Equation]
We refer to this bound as η -DV bound.
Note that for η =
we recover the f -div bound.
Using Lemma 2, we can now rewrite the estimator for the DV bound in Eq.
5 as follows:
[Equation]
where ˆ L ( f,
[Equation]
N
[Equation]
enabling unbiased stochastic gradient optimization of the function f .
We note that similar variational tricks of non-linearities have been devised for g ( η )= √ η in Argyriou et al.
; Bach et al.
η -DV B OUND , C ONVEXITY , MMD R ATIO E STIMATES AND L AGRANGIANS
We established that the DV bound is tighter than the f -div and derived η -DV as an alternative to
DV , enabling unbiased stochastic optimization.
However, the role of η in improving the lower bound remains unclear.
To highlight its role, we consider the η -DV bound from Eq.
7, and restrict the hypothesis function class to an RKHS, similarly to Nguyen et al.
Working in the dual, we can compare η -DV estimator to the f -div estimator and further understand this hierarchy of lower bounds.
See Figure 2 for an illustration.
For simplicity, and without loss of generality, we let RKHS be a ﬁnite dimensional feature map,
i.e., H = { f | f ( x ) = (cid:104) w, Φ( x ) (cid:105) , Φ : X → R m , w ∈ R m } .
Now for f ∈ H , the loss given in Eq.
8 for η -DV can be rewritten as follows:
[Equation]
Following Nguyen et al.
(2008), we consider the following regularized loss: L ( w, η ) =
L ( w, η ) + Ω( w ) , and the corresponding sample-based formulation ˆ L ( w, η ) =
ˆ L ( w, η ) + Ω( w ) , where Ω( w ) is a convex regularizer, e.g., Ω( w )
=
λ 2 (cid:107) w (cid:107) 22 .
The η -DV primal problem can now be deﬁned as:
[Equation] [Equation]
[Table]
+v+E46W18zIPycXN9rm3EF/zRg0W3zMndd0gaLFaVDSKYkUXZdOJNCBQzQPhwshwKxUzbrjAUPnGlqnh9UyKx42fuLz07XboKn3ezEtyc9xLT3rHV6fd/o91ay3ymRySLyQl56RPLsmADIkgf6I42o324lbci8/i81U0jtYzn8gG4v4TWyPIjg==</latexit>
In the following, we show that L ( w, η ) is jointly convex in ( w, η ) and derive its dual, which will shed light on the nature of the likelihood ratio estimate as well as the role of η .
Convex Estimate in RKHS.
In Lemma 3 we ﬁrst establish the convexity of the η -DV loss function.
Lemma 3.
L and ˆ L are jointly convex in w and η .
ˆ
[Equation]
=
ˆ L ( u, η )+Ω( u )+ εe − η .
Our KL estimate is therefore given as
[Equation]
where ( w ∗ , η ∗ ) =
arg min ( u,η ) ˆ L ε ( u, η ) .
η -DV
Dual is a Constrained Likelihood Ratio Estimation.
Another angle for estimating the KL divergence is through likelihood ratio estimation (Mohamed & Lakshminarayanan, 2016).
Given a ratio estimate ˆ r of p/q , the KL divergence can be computed as E Q ˆ r log(ˆ r ) , which can be easily estimated using samples from Q .
In Theorem 1, proven in Appendix A, we show that the dual problem (denoted D), corresponding to the primal minimization problem (denoted P) of η -DV , reduces to a constrained likelihood ratio estimation problem (denoted C).
Theorem 1.
Let Ω ∗ ( . ) be the Fenchel conjugate of Ω( . ) .
The η -DV bound restricted to an RKHS, amounts to the following regularized convex minimization:
P = − (min w,η L ( w, η ) + Ω( w )) , that has the following dual form:
[Equation]
where ∆( r ) =
(cid:82) X Φ( x ) p ( x ) dx − (cid:82)
X r ( y ) q
( y )
Φ( y ) dy .
Noticing that η − 1 plays the role of a Lagrangian multiplier, this is equivalent to the following likelihood ratio estimation problem:
[Equation]
subject to:
[Equation]
Therefore, we have P = D = C .
Let ( w ∗ , η ∗ ) be an optimizer of P .
Let r ∗ be an optimizer of D , then the KL estimate is given by: D H DV
[Equation]
The regularizer Ω( w ) = λ 2 (cid:107) w (cid:107) 2 can now be given the following interpretation based on the results of Theorem 1.
Recall the deﬁnition of the MMD (Gretton et al., 2012) between distributions using an
RKHS: MMD Φ ( P , Q ) = (cid:107) E x ∼ P
Φ
( x ) − E y ∼ Q Φ( y ) (cid:107) .
Replacing the Fenchel conjugate Ω ∗ ( . ) by its expression in Theorem 1, we see that η -DV is equivalent to the following dual ( η -
[Equation]
which can be written as a constrained ratio estimation problem ( η -DV –C)
[Equation]
, subject to:
[Equation]
Hence, it is clear now that the η -DV optimization problem is equivalent to the constrained likelihood ratio estimation problem r given in Eq.
13, where the ratio is estimated using the MMD distance in the
RKHS between rq and p .
It is also easy to see that p/q is a feasible point and for the universal feature map MMD Φ ( rq, p ) =
iif r = p/q , therefore, for a universal kernel, p/q is optimal and we recover the KL divergence for r = p/q .
When comparing η -DV –D from Eq.
12 and η -DV –C from Eq.
13, we see that η − 1 plays the role of a Lagrangian that ensures that rq is indeed a normalized distribution.
In practice, we solve the primal problem (P), while the dual problem (D) and its equivalent constrained form (C) explain why this formulation estimates the KL divergence and the role of η as a Lagrangian multiplier that enforces a normalization constraint.
Let r ∗ be the solution, then we have:
[Equation]
For comparison, the f -div bound, restricted to an RKHS, is equivalent to the following ratio estimation (follows from the proof of Theorem 1 by eliminating max on η , and setting η = 0 ), and is consistent with the results of Nguyen et al.
(2008) (see Eq.
51 therein):
[Equation] [Equation]
Let r ∗ be the optimum, then the KL divergence can be estimated as follows Nguyen et al.
[Equation]
Comparing D H DV ( P , Q ) and
D H
f
( P , Q )
we see that they both have a relative entropy term but the f -div bound does not impose the normalization constraint on the ratio, which biases the estimate.
We end with the following remark:
In Theorem 1, if we replace the loss L by its empirical counterpart ˆ L from Eq.
9, the equivalent Dual takes the following form:
[Equation]
and the KL estimate is given by:
[Equation]
From RKHS to Neural Estimation.
One shortcoming of the RKHS approach is that the method depends on the choice of feature map Φ .
We propose to learn Φ as a deep neural network as in MINE (Belghazi et al., 2018).
Given samples x
i from P , y i from Q , the KL estimation problem becomes:
[Equation]
which can be solved using BCD on ( w, η, Φ) .
Note that if Φ( · ) is known and ﬁxed, then optimization problem in Eq.
17 becomes convex in η and w .
We refer to this bound as η -DV –convex.
What Do Neural Estimators of KL or MI Learn?
We conclude with the following observations:
Ratio estimation via Feature Matching.
KL divergence estimation with variational bounds boils down to a ratio estimation using a form of MMD matching.
Choice of Feature/Architecture.
The choice of the feature space or the architecture of the neural network introduces a bias in the ratio estimation; also observed in Tschannen et al.
Ratio Normalization.
η -DV bound introduces a normalization constraint on the ratio, that ensures a better consistent estimate.
Impact of Regularizers.
The choice of the regularizer dictates the choice of the metric.
For instance, Ozair et al.
(2019) used gradient penalties as a regularizer, and this corresponds to a ratio matching in the Sobolev discrepancy sense.
Please see Appendix A.1 for more details.
η -DV M UTUAL I NFORMATION E STIMATION
We now specify the KL estimators given in Section 3 for the MI estimation problem.
Given a function space H deﬁned on X × Y ,
[Equation]
where I H f -div
( X ; Y )=
sup f
[Equation]
H E
[Equation]
x,
[Equation]
Equivalently, with the η -trick we have:
[Equation]
Now, given iid samples from marginals x i ∼ p x and ˜ y
i ∼ p
y , i = 1 . . .
N , and samples from the joint ( x i , y i ) ∼ p xy , i = 1 . . .
N , we can estimate the MI as follows:
[Equation]
when H is an RKHS, η can be seen as a Lagrangian ensuring that the ratio estimation r ( x, y ) of p xy p x p y is normalized when integrated on p x p y , i.e., η is a Lagrangian associated with the constraint (
cid:82) X
× Y r ( x, y )
p x ( x ) p y ( y ) dxdy
= 1 .
In Table 1 we review other variational bounds for MI based on f -div, DV , and η -DV bounds (a-MINE from Poole et al.
(2019) is discussed next).
[Table]
[Equation]
[Table]
is
[Table]
Improved Donsker-Varadhan for Mutual Information and a -DV Bound Discussion .
In the following Lemma we show that for the particular case of MI estimation, the DV bound can be made tighter:
Lemma 4.
For any function g that maps X × Y to R + we have:
[Equation]
Therefore: I ( X ; Y ) = sup
[Equation]
G ( X
[Equation]
Y )
E p x,y f ( x, y ) − E p
y log( E p x e f ( x,y ) )
where G ( X × Y ) is the entire space of functions deﬁned on X × Y .
For any function space H map X × Y to R we have:
[Equation]
Using Jensen’s inequality, we can easily derive the following hierarchy of lower bounds:
[Equation]
In Theorem 2, we apply the η -trick of Lemma 1 to I H
a -DV
( X ; Y ) , where we replace each log
[Equation]
y ) (cid:1) using the variational form of log which leads to η ( y ) .
The interchangeability between minimization and expectation and going from scalar η for each y to a function η ( . ) needs some care and we prove it in the Appendix B.1, appealing to the interchangeability principle (Lemmain Dai et al.
Theorem 2 ( η -trick for Improved DV Bound) .
Let H be a space of bounded functions deﬁned on
X × Y we have: I ( X ; Y ) ≥ − inf f ∈ H ,η ∈ G ( Y )
E p x .p
y e f ( x,y ) − η
( y )
+ E p y η
( y
[Equation]
−
[Equation]
−
[Equation]
,y (
[Equation]
x, y )) , where G ( Y ) is the entire space of functions deﬁned on Y .
Let H η be a function space subset of G ( Y ) we obtain:
[Equation]
Hence, we recover the lower bound given in Poole et al.
(2019) (given in Table 1), that was derived through other variational principles.
In the following, we make a few remarks about I H a -DV
( X ; Y ) : 1.
The I H a
-DV
( X ; Y ) bound, as given in Poole et al.
(2019), can be also obtained by setting η ( y ) =
log( a ( y )) , a ( y ) > 0 .
While this seems equivalent, the restriction of a to R + ( η is in R ) changes the landscape of the optimization problem and the interchangeability principle does not apply in a straightforward way.
Poole et al.
(2019) refers to a ( y ) as a “critic” (as in the reinforcement learning literature).
As shown in our analysis in Section 3, when we restrict H to an RKHS, η ( y ) plays a role of a
Lagrangian multiplier or a dual potential in a ratio estimation problem.
In this case it is enforcing multiple constraints
[Equation]
for each y , instead of imposing a single constraint
[Equation]
x ) p ( y ) dxdy =1 as in the
I H η -DV ( X ; Y ) bound.
No Free Lunch: while I H a -DV ( X ; Y ) is tighter than our I H
η -DV ( X ; Y ) bound, it comes at the price of estimating a function η
( y ) together with the witness function f .
This triggers a higher computational cost and larger estimation errors.
We show in Appendix C how I H a -DV
( X ; Y ) can be extended to conditional MI estimation.
Generality: η -DV bound is applicable for general probability distributions P and Q : KL ( P , Q ) ≥
D H
η -DV ( P , Q ) , while I H a -DV
( X ; Y ) only works for P = p xy and Q = p x p y .
Finally, in Algorithm 1 we outline the steps of η -MINE for MI estimation.
Algorithm 1 η -MINE (Stochastic BCD )
Inputs: X, Y dataset X ∈ R N × d x , Y ∈ R N × d y , such that ( x i
= X i,.
, y i
=
Y i,. )
∼ p xy Hyperparameters: α η , α θ (learning rates), n c (number of critic updates)
Initialize η , θ parameter of the neural network f θ for i = 1 . . .
Maxiter do for j = 1 . . .
n c do
Fetch a minibatch of size N ( x i , y i ) ∼ p xy Fetch a minibatch of size N ( x i , ˜ y i ) ∼ p x p y { ˜ y
i obtained by permuting rows of Y }
Evaluate ˆ L ( f θ , η ) =
e − η 1 N (cid:80)
Ni =
e f θ ( x i , ˜ y i ) − 1 N (cid:80)
Ni =1 f θ ( x i , y i ) + η − 1 Stochastic Gradient step on θ : θ ← θ − α ∂ ˆ L ( f θ ,η ) ∂θ { We use ADAM } end for
Update η : η ← η − α η ∂ ˆ L ( f θ ,η ) ∂η end for
Output
[Equation]
E XPERIMENTS
In this Section, we conduct a set of experiments using synthetic and real data on a number of applications to compare the proposed η -MINE MI estimator to the existing baselines.
MI estimation.
We compared different MI estimators on three synthetically generated Gaussian datasets [5K training and 1K testing samples].
Each evaluated MI estimator was run 10 times and its average performance (together with the standard deviation) is shown in Fig.
3.a.
As can be seen,
[Figure]
[Figure] [Figure]
[Table]
[Figure] [Figure]
[Figure]
[Table]
[Figure] [Figure]
[Figure]
[Table] [Table]
[Figure] [Figure]
[Figure]
[Table]
[Figure]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
/EIqkU8s5ZvvB4niteTVvKkENyRE6IQy5IkdySEikTRvrkhbySN+PZeDc+jM/J6pIxvTkgMzC+fgHEbJy1</latexit>
[Figure]
U n c o nd .
G A N
[Figure]
MI estimation in high dimensions is a challenging task, where the estimation error for all methods increases as the data dimensionality grows [red line shows true value for MI].
Nevertheless, the proposed η -MINE is able to achieve on average more accurate results compared to the existing baselines.
We also see that the convex formulation of η -MINE has overall a better performance and fast convergence rate.
This estimator has a linear witness function that is deﬁned as f ( · ) = (cid:104) w, Φ( · ) (cid:105) using pre-trained ﬁxed feature map Φ( · ) from the regular η -MINE.
In the experiments that follow we compare the proposed estimator η -MINE to DV-MINE, as a main baseline approach.
MI-regularized GAN.
We investigate GAN training improvements with MI, and especially diminishing mode collapse, as addressed in Belghazi et al.
(2018) in Section 5.1.
Belghazi et al.
(2018) uses a 25-Gaussian dataset to show improvements on GAN Clustering by using MI objective for regularization.
As in InfoGAN (Chen et al., 2016), the conditional generator G is supplied with random codes c along with noise z , we maximize the mutual information between I ( G ( z, c ) , c ) using η -MINE estimators.
In Fig.
3.b, we repeat this task and establish that η -MINE can recover all the modes within a fewer steps than DV-MINE (50K vs. 140K) and with a stable training.
Self-supervised: Deep InfoMax (Hjelm et al., 2018).
In unsupervised or self-supervised learning, the objective is to build a model without relying on labeled data but using an auxiliary task to learn informative features that can be useful for various downstream tasks.
In this experiment we evaluate the effectiveness of the proposed η -MINE estimator for unsupervised feature learning using the recently proposed Deep InfoMax method from Hjelm et al.
For feature representation we used an encoder similar to DCGAN (Radford et al., 2015), shown in Fig. 4.a
and evaluated results on
CIFAR10 and STL10 datasets (STL10 images were scaled down to match CIFAR10 resolution).
[Figure]
f i x e d t r a i n e d
Test
C
S
Train
DV η
Sup.
DV η
Sup.
C
[Table] [Table]
[Table]
S 67.
The encoder is trained by maximizing MI
I ( x (cid:48) , y (cid:48) ) between features from one of the shallow layers l ( x (cid:48)
=
E l ( x ) ) and a deeper layer k
( y (cid:48)
=
E k ( y ) ).
We examined different layer combinations and found that the encoder composed of only the ﬁrst two convolutional layers give the best performance on the downstream tasks.
As shown in Fig.
4.b
the encoder features are passed through additional trainable neural network layers, whose job is to build a classiﬁer f ( x (cid:48) , y (cid:48) ) , discriminating cases when x (cid:48) are y (cid:48) are coming from the same image and cases when x (cid:48) are y (cid:48) are unrelated.
Finally, we attach a linear layer to the pre-trained and now ﬁxed encoder (see Fig.
4.c)
to perform supervised training.
In Tab. 2 we present the results for two MI estimators: η -MINE and DV-MINE, whose loss functions are listed in Tab. 1.
As can be seen, η -MINE-based pre-training performs competitively with DV-MINE, achieving overall better results on both datasets, showing practical beneﬁts of the proposed approach.
Self-Supervised: Jigsaws with MI .
The self-supervision Jigsaw pretext task (Noroozi & Favaro, 2016; Kolesnikov et al., 2019) aims at solving an image jigsaw puzzle by predicting the scrambling permutation π .
From image X with x = { x 1 . . .
x 9 } 3 × 3 jigsaw patches, and permutation y
=
π i : [1 , 9] → [1 , 9] , scrambled patches x π i = { x π i (1) . . .
x π i (9) } generate the puzzle.
Each patch is passed through an encoder E , which needs to learn meaningful representations such that a classiﬁer C J can solve the puzzle by predicting the scrambling order π i (Noroozi & Favaro, 2016)
(see Fig. 5.a).
While this standard formulation relies on a CE-based classiﬁcation of the permutation, we propose here to use MI Jigsaw, where we train the encoder E to maximize I ( E ( x π i ); π i ) by using MI estimators
DVand η -MINE, as seen in Fig.
5.b.
A patch preprocessing similar to Kolesnikov et al.
(2019) avoids shortcuts based on color aberration, patch edge matching, etc.
(Noroozi & Favaro, 2016); for details of our implementation, see Appendix D.
All our models are built on a 10% subset of ImageNet (128K train.
, 5K val.
, 1K classes) as proposed by Kolesnikov et al.
This is a larger set than Tiny
ImageNet (200 classes) used in many publications.
E is a ResNet50 for all our experiments.
In our
ImageNet target classiﬁcation task, E from CE and MI Jigsaw trainings are frozen (at 200 epochs) and followed by linear classiﬁer C (Fig. 5.c); an adequate setup for comparing encoders as argued by
Kolesnikov et al.
All models are compared on target task in Tab. 3.
Best accuracy numbers are reported for C s trained for exactly 200 epochs.
For all results, E is trained from Jigsaw task (CE or MI) and frozen, with only C trained as in Kolesnikov et al.
DV- and η -MINE share the same f architecture.
η -MINE gives better accuracy performance compared to DV-MINE.
CE model trained from a Jigsaw- supervised encoder E provides an upper-bound for supervised performance for E from the Jigsaw task.
Despite CE being better than η -MINE and DV-MINE, η -MINE does a respectable job at learning a useful representation space for the target task, better than DV-MINE.
More details about the Jigsaw experiments setup can be found in Appendix D.
[Table] [Table]
+fwCZnpb6</latexit> f ( E ( x ⇡ ) ,y ) <latexit sha1_base64="6ZoobvEuxK05CaSONFMRRbkofmA=">AAAB+HicbVBNS8NAEJ34WetHox69BIvQgpSkCnosiuCxgv2ANoTNdtMu3WzC7kaMob/EiwdFvPpTvPlv3LY5aOuDgcd7M8zM82NGpbLtb2NldW19Y7OwVdze2d0rmfsHbRklApMWjlgkuj6ShFFOWooqRrqxICj0Gen44+up33kgQtKI36s0Jm6IhpwGFCOlJc8sBZWbyqOX9WM6qZ6mVc8s2zV7BmuZODkpQ46mZ371BxFOQsIVZkjKnmPHys2QUBQzMin2E0lihMdoSHqachQS6WazwyfWiVYGVhAJXVxZM/X3RIZCKdPQ150hUiO56E3F/7xeooJLN6M8ThTheL4oSJilImuagjWggmDFUk0QFlTfauEREggrnVVRh+AsvrxM2vWac1ar352XG1d5HAU4gmOogAMX0IBbaEILMCTwDK/wZjwZL8a78TFvXTHymUP4A+PzB1bYkjs
[Figure]
</latexit>
C ONCLUSION
In this paper we introduced a new lower bound on the KL divergence and showed how it can be used to improve the estimation of mutual information using neural networks.
Theoretically we proved that the dual of our η -DV formulation reduces to a constrained likelihood ratio estimation.
In practice, the stability of η -MINE is due to its unbiased gradient.
We tested our estimator η -MINE on synthetic data and applied it on various real-world tasks where MI can be used as a regularizer, or as an objective in self-supervised learning.
We used η -MINE in unsupervised learning of representations through MI maximization and by solving Jigsaw puzzles.
We leave for a future work the analysis and various possible applications of conditional mutual information neural estimators.
Supplementary Material for Improved Mutual Information Estimation
The supplementary material is organized as follows:
• In Appendix A we layout proofs of Lemmas 1, 2 and 3.
Those lemmas provide the η -DV formulation.
We also give in this Appendix the proof of the main Theorem 1.
This theorem gives the dual formulation of this variational bound in RKHS with Tikhonov Regularization and sheds the light on the ratio estimation interpretation of η -DV.
• We give in Appendix A.1 the impact of other regularizers on the estimation of KL divergence in RKHS.
• In Appendix B, we show that further improved bounds can be obtained for the particular case of mutual information leading to an improved Donsker-Varadhan.
We show that by using the interchangeability principle, we can obtain an equivalent functional η -DV bound that is related to a -MINE introduced in Poole et al.
• In Appendix C, we show how those improved variational bounds can be extended to conditional mutual information.
• In Appendix D, we give more details on the experimental setup of self-supervised learning with the jigsaw task.
A T ECHNICAL P ROOFS
Proof of Lemma:1.
Let g ( η ) =
e − η x + η − 1 .
We have g (cid:48)
( η ) =
− e − η x +
and g (cid:48)(cid:48)
( η ) =
e − η x > 0 .
Hence g is convex and admits a unique minimum.
First order optimality gives us: g (cid:48) ( η ∗ ) =
− e − η ∗ x + 1 = 0 , hence e − η ∗ x = 1 and η ∗ = log( x ) .
Finally, we get g ( η ∗ ) =
e − log( x ) x + log( x ) − 1 =
e log( 1 x ) x + log( x ) − 1 = 1 + log( x ) − 1 = log( x ) .
Proof of Lemma 2.
By Donsker-Varadhan representation we have:
[Equation]
Proof of Lemma 3.
L is a sum of a linear function in w , η and Ω( w ) , that are both convex.
We need to prove that f b ( w, η ) =
e − η e (cid:104) w,b (cid:105) is jointly convex, since L can be written as a sum of f b and other convex functions:
[Equation]
The Hessian of f b ( w, η ) , Hf b , has the following form:
[Equation]
(cid:21) Let us prove that for all ( w, η ) : ( w (cid:48) , η (cid:48) )
(cid:62) Hf b ( w, η )( w (cid:48) , η (cid:48) ) ≥ 0 , ∀ ( w (cid:48) , η (cid:48) ) .
We have:
[Equation]
Hence, f b is jointly convex in η and w , and so is L .
Now adding the perturbation εe − η to the loss, we get
[Equation]
for ﬁnite η
and hence we have a strict convexity.
Proof of Theorem 1.
Let’s consider:
[Equation]
Recall that the Fenchel conjugate of f is f ∗ ( p ) =
sup x (cid:104) p, x (cid:105) − f ( x ) .
It follows that ϕ ∗ ( t ) =
t log( t ) − t if t > 0 , 0 for t = 0 , and ∞ otherwise.
[Equation]
, Φ( y ) (cid:105) Introducing Lagrangian’s α we have = min w,η,z max
[Equation]
e
[Equation] η [Equation] using strong duality (convex function and lower semi-continuous) to swap min , max (Ekeland & Turnbull, 1983)
[Equation]
−
[Equation]
Hence [Equation] we have for α >
[Equation]
and ﬁnite η :
[Equation]
X
[Equation]
Deﬁne
[Equation]
then we have ( A ) :
[Equation]
and ( B ) :
[Equation]
Hence
[Equation]
by replacing (A) and (B) by their expressions
[Equation]
we get: min w,
[Equation]
Now consider again the η -DV Bound:
[Equation]
We see that η -DV amounts to a likelihood ratio r estimation, where we match with the dual norm Ω ∗ the difference of kernel mean embedding of p
[Equation]
and rq .
And η plays a role of a Lagrange multiplier that enforces the constraint that rq is a normalized distribution.
In particular, consider Ω( u ) =
λ
2 || u
|| 22 , so that Ω ∗ ( v ) =
12 λ
|| v
||
[Equation]
and hence we have the following equivalent problem of the regularized η -DV bound:
[Equation]
[Table]
[Equation]
[Table]
Another way of coming informally to the regularized η -DV objective is to start from expression of
KL using r :
[Equation]
We see that r in the dual function in estimating the ratio, η is the Lagrangian of the inequality constraint, and the MMD is added to the cost as a penality.
A. I MPACT OF REGULARIZERS ON KL E STIMATION
As discussed before, the choice of the regularizer Ω translates to a particular type of metric used in estimating the likelihood ratio: Ω( . )
= (cid:107) .
(cid:107) 22 translates in a ratio estimation in the MMD sense.
Here, we consider other choices of regularization used in Fisher discriminant analysis (Harchaoui et al., 2008; Mroueh & Sercu, 2017), speciﬁcally:
[Equation]
or gradient penalty regularizer (Gulrajani et al.
[Equation]
al
[Equation] [Equation]
[Equation]
It is easy to see that when we restrict f to H , η -DV –Fisher has the same objective as in Eq.with Ω( w ) =
λ 2 (cid:104) w, ( C ( P , Q )) w (cid:105) , where C ( P , Q ) =
E P + Q 2 Φ( x ) ⊗ Φ( x ) is a covariance matrix.
Similarly, η -DV –Sobolev corresponds to Eq.
11 with Ω( w )=
λ 2 (cid:104) w, ( D ( P , Q )) w (cid:105) , where D ( P , Q )=
E ( P + Q ) /(cid:80)
d j =1 ∂ j Φ( x ) ⊗ ∂ j Φ( x ) is a Gramian of derivatives.
The convex conjugate Ω ∗ ( w ) is therefore 12 λ (cid:10) w, ( C † ( P , Q ))
w (cid:11) and 12 λ (cid:10) w, ( D † ( P , Q ))
w (cid:11) , for η -DV –Fisher and η -DV –Sobolev respectively (where † denotes the pseudo-inverse).
Let δ P , Q = E P
Φ( x ) − E Q Φ( y ) .
Recall from Harchaoui et al.
that the Fisher discrepancy is F 2Φ ( P , Q )
= (cid:10) δ P , Q , C ( P , Q ) † δ P , Q (cid:11) , and from Mroueh et al.
(2018; 2019) the Sobolev discrepancy is
S 2Φ ( P , Q )= (cid:10) δ P , Q , D ( P , Q ) † δ P , Q (cid:11) .
Now replacing Ω ∗ by its expression in Theorem 1 we see that the choice of the regularizer dictates the discrepancy used for the ratio estimation:
[Equation]
, subject to: (cid:90) X
[Equation]
y
[Equation]
y
[Equation]
, subject to: (cid:90) X r ( y ) q
( y ) dy = 1 .
B F URTHER B OUNDING M UTUAL INFORMATION : a -MINE F ORMULATION FROM I NTERCHANGEABILITY P RINCIPLE
We start by the classical result of Donsker and Varadhan:
Lemma 5 (Donsker-Varadhan) .
For any Distribution P and any distribution Q and any non-negative function g for which E x ∼ P log g ( x ) is ﬁnite, we have:
[Equation]
This gives us a representation of KL divergence:
[Equation]
The following Lemma shows that a ﬁner application of the Donsker-Varadhan Lemma gives an improved variational representation of MI :
Lemma 6.
Let g >
[Equation]
Hence, we get
[Equation]
where G ( X × Y ) is the entire space of functions deﬁned on X × Y .
Proof.
[Equation]
Now applying Donsker-Varadhan Lemma
[Equation]
we have:
[Equation]
We conclude therefore that
[Equation]
B. η - TRICK AND INTERCHANGEABILITY PRINCIPLE FOR IMPROVED DV
We just showed that an improved Donsker-Varadhan representation of Mutual Information can be be derived as follows:
[Equation]
We show in this section that, by using interchangeability principle, we can replace the ” log ” non linearity preceding the expectation with an ”functional” η -trick.
We start ﬁrst by stating the interchangeability principle between expectation and minimum operations:
Lemma 7 (Interchangeability Principle see Lemma 1 in (Dai et al., 2016)) .
Let ξ be a random variable on Ξ and assume for any ξ ∈ Ξ , the function g ( ., ξ ) :
R → ( −∞ , + ∞ ) is a proper and lower semi-continuous convex function.
Then
[Equation]
Applying the interchangeability principle and the variational representation of log we obtain the following Theorem:
Theorem 3 (Theorem 2 restated: η -trick for improved DV Bound) .
Let H be a space of bounded functions deﬁned on X × Y , we have:
[Equation]
where G ( Y ) is the entire space of functions deﬁned on Y .
Let H η be a function space subset of G ( Y ) , we obtain:
[Equation]
Proof.
We know from Lemma 6 that for any function space H
[Equation]
Applying the η -trick for each log in the expectation:
[Equation] [Equation]
We assume here that the function space f is bounded.
Before applying the interchangeability Lemma 7, let us verify that all assumptions hold for
[Equation]
Convex.
g ( ., y ) is convex.
Considering derivatives wrt to η : g (cid:48)
( η, y )
=
− e − η E p x
( e f ( x,y ) )
+ 1 , and g (cid:48)(cid:48)
( η, y ) =
e − η E p x ( e f ( x,y ) ) > 0 .
Proper Function.
g ( η, . ) is a proper function, since f is bounded ( − M ≤ f ( x, y ) ≤ M ) :
[Equation]
and
[Equation]
since
[Equation]
Lower semi-continuity.
g ( η, . ) is lower semi-continuous since it is continuous (sum of continuous functions).
Now back to Eq.
26, we are ready to apply the interchangeability principle given in Lemma 7:
[Equation]
Using this equivalent form, we have:
[Equation]
It follows that
[Equation]
Consider now a function space H η from Y → R we have:
[Equation]
The MI lower bound is looser than regular η − MINE, since we have two lower bounds passing from all function spaces G ( X × Y ) on to H , then going from G ( Y ) to H η .
No free lunch: tighter bound
but we need to estimate a function η ( x ) which makes bigger estimation errors.
To get a good estimate of η ( x ) , to obtain its optimal value log( E p y e f ( x,y ) ) , we need for each x many samples from p y (which means many epochs of cycling through the same data in random order).
The η -trick with only a scalar variable is better computationally, since it enjoys convexity and has better optimization properties.
But as it is clear from this discussion, η function of a modality is good for conditional MI estimation which we will discuss next.
C C ONDITIONAL M UTUAL I NFORMATION
To illustrate the power of improved Donsker-Varadhan and its variational representation with the ηtrick and interchangeability principle, we show how to extend this to Conditional Mutual Information estimation.
We have a ( X, Y ) pair of random variables that may depend on another variable Z .
This conditional dependence is of high importance in all applications of joint embedding or in the information bottleneck principle, etc.
[Equation]
meaning that conditioned on Z , how much X and Y are dependent.
Therefore we have:
[Equation]
Now applying Donsker-Varadhan Lemma
[Equation]
we have for all z and for any function f
[Equation]
Now marginalizing on Z
[Equation]
we get
[Equation]
Hence
[Equation]
I
[Equation]
and
[Equation] [Equation]
Now using the η -trick for each z , followed by the interchangeability principle, we have the ﬁnal variational form:
[Equation]
D J IGSAW SELF - SUPERVISION AND M UTUAL I NFORMATION
Jigsaw Task Preprocessing For Jigsaw self-supervision task, the encoder needs to learn features from images randomly cropped and tiled into 3 × 3 jigsaw puzzles.
In order to avoid shortcuts based on matching edges of patches or picking up on color aberration as discussed in Noroozi & Favaro (2016), images are resized to 292 × 292, randomly cropped into 255 × 255 and jigsawed into 3 ×or 85 × 85 per patch wherein a 64 × 64 patch is randomly cropped, avoiding potentially easy edge matching.
× 255 crops are randomly assigned to gray scales with a probability of 2 ⁄ 3 prior to patch extraction.
Each patch RGB channel is normalized to zero mean, unit variance.
At test time, the 255 × 255 crops are obtained from images center crops.
ImageNet 10% For the Jigsaw experiments, we decided to subsample ImageNet to 10% within each class for both training and validation sets (128K training and 5K validation images).
The reason is that running on all of ImageNet 1.3 Million images was not practical for reasonable training turnaround on our infrastructure given what we had to train one set of encoder models from the pretext task and one set of classiﬁer models for the target task.
Even though we only train classiﬁers for the target task (as explained later in this section), we still need to forward input features through the Encoder which can still be computationally costly despite only backpropagating through the classiﬁer.
This 10% subset of ImageNet was proposed by (Kolesnikov et al., 2019) and can be used for faster training turnaround while still being
Tiny ImageNet dataset (
Jigsaw CE Training For self-supervision Jigsaw pretext task (Noroozi & Favaro, 2016; Kolesnikov et al., 2019)
we aim at solving an image jigsaw puzzle by predicting the scrambling permutation π .
Given an image X with x = { x 1 , . . .
, x 9 } 3 × 3 jigsaw patches, and a permutation y
=
π i : [1 , 9] → [1 , 9] , scrambled patches x π i = { x π i (1) , . . .
, x π i (9) } generates the puzzle.
The Encoder E is based on a ResNet50 and outputs 2048 features for each patch.
The classiﬁer C (or ”descrambler” in this case) is a simple Neural Network with a linear projection upfront from 2048 to 512 for each patch features that are then stacked (1 ×
and linearly projected to a 512 dimensional subspace, followed by batch normalization, ReLU, dropout, and another linear projection from 512 to 100 (number of our subset of permutations from Noroozi & Favaro (2016)).
E and C area trained with minibatch made from 16 permutations of 10 images (160 effective minibatch size).
The 10 images andpermutations are randomized for each minibatches.
For both E and C updates, we use ADAM with ( β 1 = 0 .
9 , β 2 = 0 .
with a learning rate of 10 − 4 that stays ﬁxed for all epochs and no weight decay.
We let the model build for a total of 200 epochs with a random seed ﬁxed (for reproducibility).
We perform 4 independent builds with 4 different random seeds.
We let the model build and just survey the Jigsaw solving accuracy assess the health of training only, not to select and encoder model as performance in pretext task is not a good predictor of performance in the target task (as observed in Kolesnikov et al.
We run for 200 epochs so to be able to turnaround results fast on our infrastructure, not for the goal of early stopping as no overﬁtting was observed.
[Table]
MI training For MI Jigsaw training, we train E so to maximize MI such that max E
I
[Equation]
, which means minimizing the following losses derived from Eq. 8:
[Equation]
where x are scrambled images patches, and y are scrambling permutations.
( x, y ) ∼ p x p y are sampled from marginals p x and p y ( x was not scrambled by permutation y )
while ( x, y ) ∼ P xy are sampled from the joint distribution ( y permutation scrambled x ).
For MI training, we need to deﬁne the witness function f ( x, y ) .
x is the set of 9 patches feature vectors from E and y is a one-hot embedding of our 100 permutations from Noroozi & Favaro (2016).
x and y are ﬁrst linearly projected in a common dimensionality of 512.
Then x goes through a shallow network made of 2 sequential stages each composed of CNN, batch normalization, ReLU, and max pooling.
After these 2 stages, we have about a combined 9216 features that we project to 512 dimensional space (same as dimension used for the projection of y ).
f ( x, y ) is then just the dot product of 512 feature vectors from x and y once normalized, i.e. a cosine distance that can give scores between -1 and 1.
We train using minibatches of 10 images and 16 permutations (160 effective minibatch size) as sampling from the joint distribution p xy , allowing for 16 unused permutations for sampling from marginal distribution p x p y .
Therefore expectation terms in the losses above for joint and marginal distribution samples are using the same number of samples (160) for each minibatch.
Again, we use ADAM for E and f ( x, y ) with ( β 1 = 0 .
9 , β 2 = 0 .
with a learning rate of 10 − 4 that stays ﬁxed for all epochs and no weight decay.
For η -MINE, we use ADAM for η with the same parameters as for E and f ( x, y )
but with a learning rate of 10 − 2 to allow for faster updates of η .
η is initialized to 0 for η -MINE.
We run models for 200 epochs of training.
A ﬁxed random seed is given to ensure reproducibility.
We use 4 seeds to ensure variability of our model builds (same seeds as used for Jigsaw CE).
Once again, we observe the training and validation loss as indicators of a training’s health and observe no overﬁtting.
ImageNet 10% Target Task
For our target classiﬁcation task, Encoders E from 200 epochs of
Jigsaw Task training are frozen and used to train a classiﬁer C for ImageNet 10% with 1000 output classes.
C is composed only of a linear projection from 2048 to 1000 classes.
For each input image, we resize the images to 232 × 232 and take a 192 × 192 center crop before turning it into a 3 × 3 cells grid where each individual patch (now 64 × 64) is forwarded through E to get a 2048-dimensional vector for each of the 9 patches.
Then, we take the mean of these 9 vectors and pass this 2048-dim mean vector to the classiﬁer C as done in Kolesnikov et al.
Only the classiﬁer C is trained using ADAM with ( β 1 = 0 .
9 , β 2 = 0 .
with a learning rate of 10 − 4 that stays ﬁxed for all epochs and no weight decay.
We let C build for a total of 200 epochs with a random seed ﬁxed (for reproducibility); we used 2 seeds for each classiﬁcation build, for a total of 8 ﬁnal model conﬁguration (4 seeds for CE and MI Jigsaw training and 2 seeds for Target Task task classiﬁcation).
The reason for 4 and 2 seeds is just so that our infrastructure can handle all the required builds (200 epochs of training each time).
We then report average and standard deviation over a set of 8 accuracies (top1, top5 and top10 in %) on the validation set (5K) in Tab. 3.
We used ADAM with these parameters for all classiﬁer builds.
Here again random seeds were ﬁxed.
Model CE We built a ”baseline” we call mode CE in Tab. 3.
The CE model is an attempt at providing a meaningful upper-bound on the performance on our MI trained Encoders by using a
Jigsaw CE trained encoder, freezing it and just learning C .
200 epochs of training is used, 2 random seeds as well.
Despite being better than η -MINE (and DV-MINE), the difference is only 1.9% absolute for top1 (for η − MINE ).
This means that our MI training, and particularly η -MINE, does a respectable job at learning a useful representation space for the target task.