Year: 2019

Abstract:
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 efficient 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.

text: 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 efficient 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.

heading: 1 INTRODUCTION
text: 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 flow 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 scientific fields, 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 defined 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). Specifically 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 finding 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. (2018). 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 in Chen et al. (2016). On real data,
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.

heading: 2 LOWER BOUNDS ON KL DIVERGENCES AND MUTUAL INFORMATION
text: Consider two probability distributions P and Q, where P is absolutely continuous w.r.t. Q. Let p and q be their respective densities defined on X ⊂ Rd. Their KL divergence is defined as
KL(P,Q) = Ex∼P log ( p(x)
q(x)
) = ∫ p(x) log ( p(x)
q(x)
) dx.
We are interested in the MI between two random variables X , Y where X is defined on X ⊂ Rdx , and Y on Y ⊂ Rdy . Let pxy be their joint densities and px, py the marginals of X and Y respectively. The MI is defined as follows: I(X;Y ) = KL(pxy, pxpy), (1) 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 first variational characterization of the KL divergence goes back to Donsker & Varadhan (1976):
KL(P,Q) ≥ DHDV(P,Q) = sup{Ex∼Pf(x)− log(Ex∼Qef(x)) : f ∈H }, (2) 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. (2008); Nowozin et al. (2016) and derived through convex duality to be finally stated as follows:
KL(P,Q) ≥ DHf (P,Q) = 1 + sup{Ex∼Pf(x)− Ex∼Qef(x) : f ∈H }, (3) 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 ( Ex∼Qef(x) ) ≤ Ex∼Qef(x) − 1 which means that for any function space H we have:
KL(P,Q) ≥ DHDV(P,Q) ≥ DHf (P,Q), (4) from which we conclude that the DV bound is tighter than the f -div bound.
Now, given samples {xi, i= 1 . . . N, xi ∼ P}, {yi, i= 1 . . . N, yi ∼ Q}, estimating the KL divergence can be done by computing the variational bound from Monte-Carlo simulation. Specifically,
for the DV bound we have the following estimator:
D̂HDV(P,Q) = sup
{ 1
N
N∑
i=1
f(xi)− log ( 1
N
N∑
i=1
ef(yi) ) : f ∈H } . (5)
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:
D̂Hf (P,Q) = 1 + sup
{ 1
N
N∑
i=1
f(xi)− 1
N
N∑
i=1
ef(yi) : f ∈H } , (6)
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 fit 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. (2018). 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: log(x) = minη e−ηx+ η − 1.
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:
KL(P,Q) ≥ DHη-DV(P,Q) = − inf{L(f, η) : f ∈H , η ∈ R} (7)
L(f, η) = e−ηEx∼Qef(x) − Ex∼Pf(x) + η − 1, (8) We refer to this bound as η-DV bound.
Note that for η=0 we recover the f -div bound. Using Lemma 2, we can now rewrite the estimator for the DV bound in Eq. 5 as follows:
D̂HDV(P,Q) = − inf f,η
L̂(f, η), where L̂(f, η) = e−η 1
N
N∑
i=1
ef(yi) − N∑
i=1
f(xi) + η − 1, (9)
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. (2008); Bach et al. (2011).
3 η-DV BOUND, CONVEXITY, MMD RATIO ESTIMATES AND LAGRANGIANS
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. (2008). 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 finite dimensional feature map, i.e., H = {f |f(x) = 〈w,Φ(x)〉 ,Φ : X → Rm, w ∈ Rm}. Now for f ∈H , the loss given in Eq. 8 for η-DV can be rewritten as follows:
L(f, η) =∆ L(w, η) = e−ηEx∼Qe〈w,Φ(x)〉 − 〈w,Ex∼PΦ(x)〉+ η − 1. (10) 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 ‖w‖ 2 2 . The η-DV primal problem can now be defined as:
η-DV–P : min w,η L(w, η) + Ω(w). (11)
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 first establish the convexity of the η-DV loss function. Lemma 3. L and L̂ are jointly convex in w and η.
L̂ε(u, η) = L̂(u, η)+Ω(u)+εe−η.Our KL estimate is therefore given as D̂HDV(P,Q) = −L̂(w∗, η∗), 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 EQr̂ 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 = − (minw,η L(w, η) + Ω(w)), that has the following dual form:
D = min r>0 max η
∫
X r(y) log (r(y)) q(y)dy + (η − 1)
(∫
X r(y)q(y)dy − 1
) + Ω∗ (∆(r)) ,
where ∆(r) = ∫ X Φ(x)p(x)dx− ∫ 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:
C = min r>0
∫
X r(y) log (r(y)) q(y)dy + Ω∗
(∫
X Φ(x)p(x)dx−
∫
X r(y)q(y)Φ(y)dy
)
subject to: ∫
X r(y)q(y)dy = 1.
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: DHDV (P,Q) = ∫ X r ∗(y) log (r∗(y)) q(y)dy = L(w∗, η∗).
The regularizer Ω(w) = λ2 ‖w‖ 2 can now be given the following interpretation based on the results of Theorem 1. Recall the definition of the MMD (Gretton et al., 2012) between distributions using an RKHS: MMDΦ(P,Q) = ‖Ex∼PΦ(x)− Ey∼QΦ(y)‖. Replacing the Fenchel conjugate Ω∗(.) by its expression in Theorem 1, we see that η-DV is equivalent to the following dual (η-DV–D)
η-DV–D : min r>0 max η
∫
X r(y) log (r(y)) q(y)dy+(η−1)
(∫
X r(y)q(y)dy − 1
) + 1
2λ MMD2Φ(rq, p)
(12) which can be written as a constrained ratio estimation problem (η-DV–C)
η-DV–C : min r>0
∫
X r(y) log (r(y)) q(y)dy +
1
2λ MMD2Φ(rq, p), subject to:
∫
X r(y)q(y)dy = 1.
(13)
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) = 0 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:
DHDV(P,Q) = ∫
X r∗(y) log (r∗(y)) q(y)dy = −L(w∗, η∗). (14)
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):
f -div : min r>0
∫
X r(y) log (r(y)) q(y)dy + 1−
∫
X r(y)q(y)dy +
1
2λ MMD2Φ(rq, p). (15)
Let r∗ be the optimum, then the KL divergence can be estimated as follows Nguyen et al. (2008):
DHf (P,Q) = ∫
X r∗(y) log (r∗(y)) q(y)dy +
( 1− ∫
X r∗(y)q(y)dy
) . (16)
Comparing DHDV(P,Q) and DHf (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:
min ri>0
1
N
N∑
i=1
ri log (ri) + Ω ∗ ( 1
N
N∑
i=1
Φ(xi)− 1
N
N∑
i=1
riΦ(yi)
) , subject to: 1
N
N∑
i=1
ri = 1,
and the KL estimate is given by: D̂HDV(P,Q) = 1N ∑N i=1 r ∗ i log (r ∗ i ) = −L̂(w∗, η∗).
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 xi from P, yi from Q, the KL estimation problem becomes:
D̂NNDV(P,Q) = − (
min Φ min η,w
e−η 1
N
N∑
i=1
e〈w,Φ(yi)〉 − N∑
i=1
〈w,Φ(xi)〉+ η − 1 ) , (17)
which can be solved using BCD on (w, η,Φ). Note that if Φ(·) is known and fixed, 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:
1. Ratio estimation via Feature Matching. KL divergence estimation with variational bounds boils down to a ratio estimation using a form of MMD matching.
2. 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. (2019).
3. Ratio Normalization. η-DV bound introduces a normalization constraint on the ratio, that ensures a better consistent estimate.
4. 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.
4 η-DV MUTUAL INFORMATION ESTIMATION
We now specify the KL estimators given in Section 3 for the MI estimation problem. Given a function space H defined on X × Y ,
I(X;Y ) ≥ IHDV (X;Y ) = sup f∈H
Epxyf(x, y)− log ( EpxEpyef(x,y) ) ≥ IHf -div(X;Y ), (18)
where IHf -div(X;Y )=supf∈H Epxyf(x, y)− EpxEpyef(x,y). Equivalently, with the η-trick we have:
IHDV (X;Y ) = − inf f∈H ,η e−ηEpxpyef(x,y) − Epxyf(x, y) + η − 1. (19)
Now, given iid samples from marginals xi ∼ px and ỹi ∼ py, i = 1 . . . N , and samples from the joint (xi, yi) ∼ pxy, i = 1 . . . N , we can estimate the MI as follows:
ÎHDV (X,Y ) = D̂ H DV(px,y, pxpy) = − inf f∈H ,η e−η
1
N
N∑
i=1
ef(xi,ỹi) − 1 N
N∑
i=1
f(xi, yi) + η − 1, (20)
when H is an RKHS, η can be seen as a Lagrangian ensuring that the ratio estimation r(x, y) of pxy pxpy
is normalized when integrated on pxpy, i.e., η is a Lagrangian associated with the constraint∫ X×Y r(x, y)px(x)py(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).
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:
KL(px,y, pxpy) + Epx log(Epyg(x, y)) ≥ Epx,y log(g(x, y))
Therefore: I(X;Y ) = supf∈G (X×Y) Epx,yf(x, y) − Epy log(Epxef(x,y)) where G (X × Y) is the entire space of functions defined on X × Y . For any function space H map X × Y to R we have:
I(X;Y ) ≥ IHa-DV(X;Y ) := sup f∈H Epx,yf(x, y)− Epy log(Epxef(x,y)).
Using Jensen’s inequality, we can easily derive the following hierarchy of lower bounds:
I(X;Y ) ≥ IHa-DV(X;Y ) ≥ IHη-DV(X;Y ) ≥ IHf -div(X;Y ). (21)
In Theorem 2, we apply the η-trick of Lemma 1 to IHa-DV(X;Y ), where we replace each log ( Epxef(x,y) ) 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 (Lemma 1 in Dai et al. (2016)).
Theorem 2 (η-trick for Improved DV Bound). Let H be a space of bounded functions defined on X ×Y we have: I(X;Y ) ≥ − inff∈H ,η∈G(Y) Epx.pyef(x,y)−η(y) +Epyη(y)− 1−Epx,y (f(x, y)), where G (Y) is the entire space of functions defined on Y . Let Hη be a function space subset of G (Y) we obtain: I(X;Y ) ≥ IHa-DV(X;Y ) = sup
f∈H ,η(.)∈Hη Epx,yf(x, y)− Epx.pyef(x,y)−η(y) − Epyη(y) + 1. (22)
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 IHa-DV(X;Y ):
1. The IHa-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.
2. 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
∫ X r(x, y)p(x)dx = 1 for each y, instead of imposing a single constraint∫
X ∫ Y r(x, y)p(x)p(y)dxdy=1 as in the I H η-DV(X;Y ) bound.
3. No Free Lunch: while IHa-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.
4. We show in Appendix C how IHa-DV(X;Y ) can be extended to conditional MI estimation. 5. Generality: η-DV bound is applicable for general probability distributions P and Q: KL(P,Q) ≥ DHη-DV(P,Q), while IHa-DV(X;Y ) only works for P = pxy and Q = pxpy .
Finally, in Algorithm 1 we outline the steps of η-MINE for MI estimation.
Algorithm 1 η-MINE (Stochastic BCD )
Inputs: X,Y dataset X ∈ RN×dx , Y ∈ RN×dy , such that (xi = Xi,., yi = Yi,.) ∼ pxy Hyperparameters: αη, αθ (learning rates), nc (number of critic updates) Initialize η, θ parameter of the neural network fθ for i = 1 . . .Maxiter do
for j = 1 . . . nc do Fetch a minibatch of size N (xi, yi) ∼ pxy Fetch a minibatch of size N (xi, ỹi) ∼ pxpy {ỹi obtained by permuting rows of Y } Evaluate L̂(fθ, η) = e−η 1N ∑N i=1 e fθ(xi,ỹi) − 1N ∑N i=1 fθ(xi, yi) + η − 1
Stochastic Gradient step on θ: θ ← θ − α∂L̂(fθ,η)∂θ {We use ADAM}
end for Update η: η ← η − αη ∂L̂(fθ,η)∂η
end for Output: fθ, η, ÎHη-DV(X,Y ) = 1 N ∑N i=1 fθ(xi, yi)− e−η 1N ∑N i=1 e fθ(xi,ỹi) − η + 1

heading: 5 EXPERIMENTS
text: 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,
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 defined as f(·) = 〈w,Φ(·)〉 using pre-trained fixed 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. (2018). 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).
The encoder is trained by maximizing MI I(x′, y′) between features from one of the shallow layers l (x′ = El(x)) and a deeper layer k (y′ = Ek(y)). We examined different layer combinations and found that the encoder composed of only the first 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 classifier f(x′, y′), discriminating cases when x′ are y′ are coming from the same image and cases when x′ are y′ are unrelated. Finally, we attach a linear layer to the pre-trained and now fixed 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 benefits 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= {x1 . . . x9} 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 classifier CJ 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 classification 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. (2019). 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 classification task, E from CE and MI Jigsaw trainings are frozen (at 200 epochs) and followed by linear classifier C (Fig. 5.c); an adequate setup for comparing encoders as argued by Kolesnikov et al. (2019).
All models are compared on target task in Tab. 3. Best accuracy numbers are reported for Cs 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. (2019). 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.

heading: 6 CONCLUSION
text: 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.

heading: A TECHNICAL PROOFS
text: Proof of Lemma:1. Let g(η) = e−ηx+ η− 1. We have g′(η) = −e−ηx+ 1 and g′′(η) = e−ηx > 0. Hence g is convex and admits a unique minimum. First order optimality gives us: g′(η∗) = −e−η∗x+ 1 = 0, hence e−η∗x = 1 and η∗ = log(x). Finally, we get
g(η∗) = e− log(x)x+ log(x)− 1 = elog( 1x )x+ log(x)− 1 = 1 + log(x)− 1 = log(x).
Proof of Lemma 2. By Donsker-Varadhan representation we have:
DHDV(P,Q) = sup f∈H
Ex∼Pf(x)− log ( Ex∼Qef(x) )
= − inf f∈H
log ( Ex∼Qef(x) ) − Ex∼Pf(x)
= − inf f∈H {inf η e−ηEx∼Qef(x) + η − 1} − Ex∼Pf(x)
= − inf f∈H ,η∈R {e−ηEx∼Qef(x) − Ex∼Pf(x) + η − 1}.
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 fb(w, η) = e−ηe〈w,b〉 is jointly convex, since L can be written as a sum of fb and other convex functions: ∇wfb(w, η) = e−ηe〈w,b〉b, ∇2wfb(w, η) = e−ηe〈w,b〉b⊗ b PSD (Positive Semi-Definite)
∂2fb ∂w∂η = −e−ηe〈w,b〉b, ∂ 2fb ∂η2 = e−ηe〈w,b〉.
The Hessian of fb(w, η), Hfb, has the following form:
Hfb(w, η) =
[ ∂2L
∂w⊗∂w ∂2L ∂w∂η
∂2L ∂η∂w
∂2L ∂η2
] = [ e−ηe〈w,b〉b⊗ b −e−ηe〈w,b〉b −e−ηe〈w,b〉b> e−ηe〈w,b〉 ]
Let us prove that for all (w, η): (w′, η′)>Hfb(w, η)(w′, η′) ≥ 0,∀(w′, η′). We have: (w′, η′)>Hfb(w, η)(w ′, η′) = e−ηe〈w,b〉(〈w′, b〉)2 − 2η′e−ηe〈w,b〉 〈w′, b〉+ η′2e−ηe〈w,b〉
= e−ηe〈w,b〉 ( (〈w′, b〉)2 − 2η′ 〈w′, b〉+ η′2 )
= e−ηe〈w,b〉 (〈w′, b〉 − η′)2 ≥ 0.
Hence, fb is jointly convex in η and w, and so is L. Now adding the perturbation εe−η to the loss, we get
w′, η′)>Hfb(w, η)(w ′, η′) = e−ηe〈w,b〉 (〈w′, b〉 − η′)2 + εe−η > 0,
for finite η and hence we have a strict convexity.
Proof of Theorem 1. Let’s consider:
f(x) = 〈w,Φ(x)〉 , and ϕ(t) = et. Recall that the Fenchel conjugate of f is f∗(p) = supx 〈p, x〉 − f(x). It follows that ϕ∗(t) = t log(t)− t if t > 0, 0 for t = 0, and∞ otherwise.
min w,η
e−ηEy∼Qϕ (〈w,Φ(y)〉)− Ex∼P 〈w,Φ(x)〉+ η − 1 + Ω(w)
= min w,η
e−η ∫
X ϕ (〈w,Φ(y)〉) q(y)dy −
〈 w, ∫
X Φ(x)p(x)dx
〉 + η − 1 + Ω(w)
= min w,η,z
e−η ∫
X ϕ (z(y)) q(y)dy −
〈 w, ∫
X Φ(x)p(x)dx
〉 + η − 1 + Ω(w) s.t. z(y) = 〈w,Φ(y)〉
Introducing Lagrangian’s α we have
= min w,η,z max α
e−η ∫
X ϕ (z(y)) q(y)dy −
〈 w, ∫
X Φ(x)p(x)dx
〉 + η − 1 + Ω(w)− ∫
X α(y)(〈w,Φ(y)〉 − z(y))dy
= max α min w,η,z
e−η ∫
X ϕ (z(y)) q(y)dy −
〈 w, ∫
X Φ(x)p(x)dx
〉 + η − 1 + Ω(w)− ∫
X α(y)(〈w,Φ(y)〉 − z(y))dy
using strong duality (convex function and lower semi-continuous) to swap min, max (Ekeland & Turnbull, 1983)
= max α min η − ( max z − (∫ X (ϕ (z(y)) q(y)e−η + α(y)z(y))dy )) − (
max w
〈 w, ∫
X Φ(x)p(x)dx+
∫
X α(y)Φ(y)dy
〉 − Ω(w)− η + 1 )
= max α min η − (∫ X q(y)e−η max z −ϕ(z)− α(y) q(y)e−η z ) − Ω∗ (∫ X Φ(x)p(x)dx+ ∫ X α(y)Φ(y)dy ) + η − 1
= max α<0 min η − (∫ X q(y)e−ηϕ∗(− α(y) q(y)e−η ) ) − Ω∗ (∫ X Φ(x)p(x)dx+ ∫ X α(y)Φ(y)dy ) + η − 1
= max α>0 min η − (∫ X q(y)e−ηϕ∗( α(y) q(y)e−η ) ) − Ω∗ (∫ X Φ(x)p(x)dx− ∫ X α(y)Φ(y)dy ) + η − 1
= −min α>0 max η
(∫
X q(y)e−ηϕ∗(
α(y)
q(y)e−η )
) + Ω∗ (∫
X Φ(x)p(x)dx−
∫
X α(y)Φ(y)dy
) − η + 1
Hence, we have for α > 0, and finite η:
∫
X q(y)e−ηϕ∗
( α(y)
q(y)e−η
) dy = ∫
X q(y)e−η
( α(y)
q(y)e−η log(
α(y) q(y)e−η )− α(y) q(y)e−η
) dy
=
∫
X q(y)e−η
( α(y)
q(y)e−η log(
α(y) q(y)e−η )− α(y) q(y)e−η
) dy
=
∫
X [α(y) log
( α(y)
q(y)e−η
) − α(y)]dy
=
∫
X α(y) log
( α(y)
q(y)
) dy + (η − 1) ∫ α(y))dy.
Define
r(y) = α(y)
q(y) > 0
then we have (A) : ∫
X q(y)e−ηϕ∗
( α(y)
q(y)e−η
) dy = ∫
X r(y) log (r(y)) q(y)dy + (η − 1)
∫
X r(y)q(y)dy,
and (B) : Ω∗ (∫
X Φ(x)p(x)dx−
∫
X α(y)Φ(y)dy
) = Ω∗ (∫
X Φ(x)p(x)dx−
∫
X r(y)q(y)Φ(y)dy
) .
Hence, by replacing (A) and (B) by their expressions, we get:
min w,η L (w, η)
= −min α>0 max η
(∫
X q(y)e−ηϕ∗(
α(y)
q(y)e−η )
) + Ω∗ (∫
X Φ(x)p(x)dx−
∫
X α(y)Φ(y)dy
) − η + 1
= − min r(.)>0 max η
∫
X r(y) log (r(y)) q(y)dy + (η − 1)
(∫
X r(y)q(y)dy − 1
)
+ Ω∗ (∫
X Φ(x)p(x)dx−
∫
X r(y)q(y)Φ(y)dy
)
Now consider again the η-DV Bound:
−min w,η L (w, η)
= min r>0 max η
∫
X r(y) log (r(y)) q(y)dy + (η − 1)(
∫
X r(y)q(y)dy − 1)
+ Ω∗ (∫
X Φ(x)p(x)dx−
∫
X r(y)q(y)Φ(y)dy
)
= min r>0
∫
X r(y) log (r(y)) q(y)dy + Ω∗
(∫
X Φ(x)p(x)dx−
∫
X r(y)q(y)Φ(y)dy
)
s.t. ∫
X r(y)q(y)dy = 1.
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, 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||2, and hence we have the following equivalent problem of the regularized η-DV bound:
min r>0
∫
X r(y) log (r(y)) q(y)dy +
1
2λ
∥∥∥∥ ∫
X Φ(x)p(x)dx−
∫
X r(y)q(y)Φ(y)dy
∥∥∥∥ 2
2︸ ︷︷ ︸ MMD matching of rq and p
s,t. ∫
X r(y)q(y)dy = 1. ︸ ︷︷ ︸ Normalization Constraint
Another way of coming informally to the regularized η-DV objective is to start from expression of KL using r:
min r>0
∫
X r(y) log(r(y))q(y)dx
MMDΦ(rq, p) < ε∫
X r(x)q(x) = 1.
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.1 IMPACT OF REGULARIZERS ON KL ESTIMATION
As discussed before, the choice of the regularizer Ω translates to a particular type of metric used in estimating the likelihood ratio: Ω(.) = ‖.‖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), specifically:
η-DV–Fisher : min f,η
L(f, η) + λ
2 Ex∼ P+Q2 f 2(x), (23)
or gradient penalty regularizer (Gulrajani et al., 2017; Mroueh et al., 2018; 2019):
η-DV–Sobolev : min f,η
L(f, η) + λ
2 Ex∼ P+Q2 ‖∇xf(x)‖ 2 . (24)
It is easy to see that when we restrict f to H , η-DV–Fisher has the same objective as in Eq. 11 with Ω(w) = λ2 〈w, (C(P,Q))w〉, where C(P,Q) = E P+Q2 Φ(x) ⊗ Φ(x) is a covariance matrix. Similarly, η-DV–Sobolev corresponds to Eq. 11 with Ω(w)= λ2 〈w, (D(P,Q))w〉, where D(P,Q)= E(P+Q)/2 ∑d j=1 ∂jΦ(x) ⊗ ∂jΦ(x) is a Gramian of derivatives. The convex conjugate Ω∗(w) is
therefore 12λ 〈 w, (C†(P,Q))w 〉 and 12λ 〈 w, (D†(P,Q))w 〉 , for η-DV–Fisher and η-DV–Sobolev respectively (where † denotes the pseudo-inverse). Let δP,Q =EPΦ(x) − EQΦ(y). Recall from Harchaoui et al. (2008) that the Fisher discrepancy is F 2Φ(P,Q)= 〈 δP,Q, C(P,Q)†δP,Q 〉 , and from Mroueh et al. (2018; 2019) the Sobolev discrepancy is
S2Φ(P,Q)= 〈 δP,Q, D(P,Q)†δP,Q 〉 . Now replacing Ω∗ by its expression in Theorem 1 we see that the choice of the regularizer dictates the discrepancy used for the ratio estimation:
η-DV-Fisher : min r>0
∫
X r(y) log (r(y)) q(y)dy +
1
2λ F 2Φ(rq, p), subject to:
∫
X r(y)q(y)dy = 1,
η-DV-Sobolev : min r>0
∫
X r(y) log (r(y)) q(y)dy +
1
2λ S2Φ(rq, p), subject to:
∫
X r(y)q(y)dy = 1.
B FURTHER BOUNDING MUTUAL INFORMATION: a-MINE FORMULATION FROM INTERCHANGEABILITY PRINCIPLE
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 Ex∼P log g(x) is finite, we have:
KL(P,Q) + log(Ex∼Qg(x)) ≥ Ex∼P log(g(x)). This gives us a representation of KL divergence:
KL(P,Q) = sup g>0 Ex∼P log(g(x))− log(Ex∼Qg(x)) = sup f Ex∼Pf(x)− log(Ex∼Qef(x)).
The following Lemma shows that a finer application of the Donsker-Varadhan Lemma gives an improved variational representation of MI : Lemma 6. Let g > 0
KL(px,y, pxpy) + Epx log(Epyg(x, y)) ≥ Epx,y log(g(x, y)). Hence, we get
I(X;Y ) = sup g>0
Epx,y log(g(x, y))− Epx log(Epyg(x, y))
= sup f∈G (X×Y)
Epx,yf(x, y)− Epx log(Epyef(x,y))
= sup f∈G (X×Y)
Epx,yf(x, y)− Epy log(Epxef(x,y))
where G (X × Y) is the entire space of functions defined on X × Y .
Proof.
KL(px,y, pxpy) + Epx log(Epyg(x, y)) = Epxy log( pxy pxpy ) + Epx log(Epyg(x, y))
= Epx ( Epy|x log(
py|x py
) + log(Epyg(x, y)) )
= Epx ( KL(py|x, py) + log(Epyg(x, y)) ) .
Now applying Donsker-Varadhan Lemma, we have:
KL(py|x, py) + log(Epyg(x, y)) ≥ Epy|x log(g(x, y)). We conclude therefore that
KL(px,y, pxpy) + Epx log(Epyg(x, y)) ≥ Epx ( Epy|x log(g(x, y)) )
= Epx,y log(g(x, y)).
B.1 η-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:
I(X;Y ) = sup f∈G (X×Y)
Epx,yf(x, y)− Epy log(Epxef(x,y)).
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 first 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
Eξ min u∈R g(u, ξ) = min u(.)∈G (Ξ) Eξg(u(ξ), ξ).
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 defined on X × Y , we have:
I(X;Y ) ≥ − inf f∈H ,η∈G(Y) Epx.pyef(x,y)−η(y) + Epyη(y)− 1− Epx,y (f(x, y)),
where G (Y) is the entire space of functions defined on Y . Let Hη be a function space subset of G (Y), we obtain:
I(X;Y ) ≥ IHa-DV(X;Y ) = sup f∈H ,η(.)∈Hη Epx,yf(x, y)− Epx.pyef(x,y)−η(y) − Epyη(y) + 1.
Proof. We know from Lemma 6 that for any function space H
I(X;Y ) ≥ − inf f∈H Epy log(Epxef(x,y))− Epx,y (f(x, y)). (25)
Applying the η-trick for each log in the expectation:
Epy log(Epxef(x,y)) = Epy min η∈R e−ηEpxef(x,y) + η − 1. (26)
We assume here that the function space f is bounded. Before applying the interchangeability Lemma 7, let us verify that all assumptions hold for
g(η, y) = e−ηEpxef(x,y) + η − 1.
Convex. g(., y) is convex. Considering derivatives wrt to η: g′(η, y) = −e−ηEpx(ef(x,y)) + 1, and g′′(η, y) = e−ηEpx(ef(x,y)) > 0.
Proper Function. g(η, .) is a proper function, since f is bounded (−M ≤ f(x, y) ≤M):
g(0, y) = Epxef(x,y) − 1 ≤ eM − 1 <∞.
and g(η, y) > −∞,∀η,
since g(η, y) ≥ g(η∗, y) = log(Epxef(x,y)) ≥ −M,∀η.
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:
Epy log(Epxef(x,y)) = Epy min η∈R e−ηEpxef(x,y) + η − 1
= Epy min η∈R g(η, y)
= min η(.)∈G (Y) Epyg(η(y), y)
= min η(.)∈G (Y)
Epy ( e−η(y)Epxef(x,y) + η(y)− 1 ) .
Using this equivalent form, we have:
inf f∈H
Epy log(Epxef(x,y))− Epx,yf(x, y)
= inf f∈H min η(.)∈G (Y)
Epy ( e−η(y)Epxef(x,y) + η(y)− 1 ) − Epx,y (f(x, y)
= inf f∈H ,η∈G(Y)
Epx.pyef(x,y)−η(y) + Epyη(y)− 1− Epx,y (f(x, y).
It follows that
I(X;Y ) ≥ − inf f∈H ,η∈G(Y) Epx.pyef(x,y)−η(y) + Epyη(y)− 1− Epx,y (f(x, y)
= sup f∈H ,η∈G(Y)
Epx,yf(x, y)− Epx.pyef(x,y)−η(y) − Epyη(y) + 1.
Consider now a function space Hη from Y → R we have:
I(X;Y ) ≥ sup f∈H ,η∈Hη Epx,yf(x, y)− Epx.pyef(x,y)−η(y) − Epyη(y) + 1,
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(Epyef(x,y)), we need for each x many samples from py (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.

heading: C CONDITIONAL MUTUAL INFORMATION
text: 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.
I(X;Y |Z) = KL(pxyz, px|zpy|zp(z)), meaning that conditioned on Z, how much X and Y are dependent.
Therefore we have:
I(X;Y |Z) = Ez∼pzKL(px,y|z||px|zpy|z). Now applying Donsker-Varadhan Lemma, we have for all z and for any function f
KL(px,y|z||px|zpy|z) + logEpx|zpy|zf(x, y, z) ≥ Epx,y|z log(f(x, y, z)). Now marginalizing on Z, we get
Ez∼pzKL(px,y|z||px|zpy|z) + Ez∼pz logEpx|zpy|zf(x, y, z) ≥ Ez∼pzEpx,y|z log(f(x, y, z)).
Hence,
I(X;Y |Z) ≥ Ez∼pzEpx,y|z log(f(x, y, z))− Ez∼pz logEpx|zpy|zf(x, y, z), and
I(X;Y |Z) = sup f∈F ,f>0 Ez∼pzEpx,y|z log(f(x, y, z))− Ez∼pz logEpx|zpy|zf(x, y, z). (27)
Now using the η-trick for each z, followed by the interchangeability principle, we have the final variational form:
I(X;Y |Z) = sup f∈F ,η
Ez∼pzEpx,y|zf(x, y, z)− Ez∼pz ( e−η(z)Epx|zpy|ze f(x,y,z) + η(z)− 1 ) .
(28)

heading: D JIGSAW SELF-SUPERVISION AND MUTUAL INFORMATION
text: 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×3 or 85×85 per patch wherein a 64×64 patch is randomly cropped, avoiding potentially easy edge matching. 255×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 classifier models for the target task. Even though we only train classifiers 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 classifier. This 10% subset of ImageNet was proposed by (Kolesnikov et al., 2019) and can be used for faster training turnaround while still being larger than the Tiny ImageNet dataset (200 classes, 500 training, 50 vaildation, 50 test images per class) often used in publications.
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={x1, . . . , x9} 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 classifier 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×(9*512)) 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 and 16 permutations are randomized for each minibatches. For both E and C updates, we use ADAM with (β1 = 0.9, β2 = 0.999) with a learning rate of 10−4 that stays fixed for all epochs and no weight decay. We let the model build for a total of 200 epochs with a random seed fixed (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. (2019). 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 overfitting was observed.
Jigsaw MI training For MI Jigsaw training, we train E so to maximize MI such that max E I ( E(xπ(1)), . . . , (xπ(9));π ) , which means minimizing the following losses derived from Eq. 8:
LDV(f) = logEx,y∼PxPy [exp (f(x, y))]− Ex,y∼Pxy [f(x, y)] (DV-MINE algo.) Lη-DV(f, η) = e
−ηEx,y∼PxPy [exp (f(x, y))]− Ex,y∼Pxy [f(x, y)] + η − 1, (η-MINE algo.) where x are scrambled images patches, and y are scrambling permutations. (x, y)∼pxpy are sampled from marginals px and py (x was not scrambled by permutation y) while (x, y)∼Pxy are sampled from the joint distribution (y permutation scrambled x). For MI training, we need to define 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 first 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 pxy , allowing for 16 unused permutations for sampling from marginal distribution pxpy. 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.999) with a learning rate of 10−4 that stays fixed 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 fixed 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 overfitting.
ImageNet 10% Target Task For our target classification task, Encoders E from 200 epochs of Jigsaw Task training are frozen and used to train a classifier 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 classifier C as done in Kolesnikov et al. (2019). Only the classifier C is trained using ADAM with (β1 = 0.9, β2 = 0.999) with a learning rate of 10−4 that stays fixed for all epochs and no weight decay. We let C build for a total of 200 epochs with a random seed fixed (for reproducibility); we used 2 seeds for each classification build, for a total of 8 final model configuration (4 seeds for CE and MI Jigsaw training and 2 seeds for Target Task task classification). 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 classifier builds. Here again random seeds were fixed.
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.

Review summaries:
This paper develops variations on the Donsker-Varadhan (DV) lower bound on KL divergence which yields a lower bound on mutual information as a special case.  The paper is primarily theoretical with an emphasis on the case where the witness function f is drawn from an RKHS (a support vector machine).  They discuss, and present experimental results where the feature map of the RKHS is computed by a neural network, although the theoretical results largely do not apply to optimization of the neural network.