Critics and variational bounds: f-divergences, InfoNCE and contrastive learning

Before this lesson: lesson 03 (Jensen's inequality, KL divergence), lesson 04 (cross-entropy, softmax cross-entropy), lesson 06 (differential entropy, Gaussians) and lesson 07 (mutual information as a KL divergence, pointwise mutual information). Logistic regression (ML fundamentals track, in preparation) helps; convex functions are in optimization lesson 02. This lesson defines the convex conjugate it needs. Time: about 90 minutes, plus the notebook. This is the hardest lesson of the track. Take a break after the variational bound (section 3, before "KL: the NWJ and DV bounds").

---

Why this matters #

• Contrastive learning. CPC, SimCLR, MoCo and CLIP train encoders with the InfoNCE loss. It is a softmax cross-entropy over a batch, and $\log N$ minus the loss is a lower bound on the mutual information between two views. The $\log N$ ceiling is one reason these methods use very large batches (CLIP used 32,768 pairs per batch).
• GANs. The discriminator of a GAN is a classifier between real and generated samples. At its optimum the GAN objective is $2\,\mathrm{JS} - \log 4$, a Jensen–Shannon divergence. f-GAN uses the same recipe for any f-divergence.
• Estimating mutual information. MINE, NWJ and InfoNCE estimate the mutual information between high-dimensional variables, where the tables of lesson 07 are hopeless.
• Density-ratio estimation. A classifier between two data sets gives importance weights for covariate shift, two-sample tests, and neural ratio estimation in simulation-based inference.

As in the whole track, $\log$ is the natural log, so every value in this lesson is in nats (divide by $\log 2 \approx 0.693$ for bits). PyTorch and NumPy losses are in nats too.

1. Predict first #

2. The picture #

Picture 1: a classifier learns the density ratio #

Take two distributions $P$ and $Q$ with densities $p$ and $q$. Draw samples of both. Give every sample of $P$ the label $1$ and every sample of $Q$ the label $0$, with equal numbers of each. Now ask: what is the best possible classifier? Bayes' rule answers this in three steps.

1. The two classes are equally likely, so $\Pr(\text{label } 1) = \Pr(\text{label } 0) = \tfrac12$.
2. Bayes' rule: $\Pr(1 \mid x) = \dfrac{\tfrac12 p(x)}{\tfrac12 p(x) + \tfrac12 q(x)} = \dfrac{p(x)}{p(x) + q(x)}$, and $\Pr(0 \mid x) = \dfrac{q(x)}{p(x) + q(x)}$.
3. The logit (log-odds) is the log of their ratio: $\log\dfrac{\Pr(1 \mid x)}{\Pr(0 \mid x)} = \log\dfrac{p(x)}{q(x)}$.

So the best classifier's logit is the log density ratio $\log r(x)$, with $r(x) = p(x)/q(x)$. This is the density-ratio trick. A logistic regression or a neural network trained with cross-entropy (lesson 04) approaches this classifier when it has enough data and capacity. It never needs the densities, only samples.

Example. $P = \mathcal{N}(1, 1)$ and $Q = \mathcal{N}(0, 1)$. The constants of the two Gaussian densities cancel:

The log ratio is a straight line: a logistic regression on $x$ can represent it exactly. If the class sizes differ, $n_P \ne n_Q$, the logit becomes $\log r(x) + \log(n_P/n_Q)$: a constant shift.

Interactive figure: loading…

A function of $x$ that we learn from samples, and whose best value is a function of the density ratio, is called a critic. The classifier's logit is our first critic. The rest of the lesson shows how to turn any critic into a number that bounds a divergence.

Picture 2: tangent lines give lower bounds #

The second picture is about convex functions. A convex function $f$ lies above each of its tangent lines. In fact $f$ is the upper envelope of its tangent lines: at every point $u$, $f(u)$ equals the highest tangent line at that point.

Name each line by its slope $t$. Write the line as $u \mapsto tu - f^*(t)$, so $-f^*(t)$ is the point where the line crosses the vertical axis. For the line to be a tangent line, $f^*(t)$ must be $\sup_u\,[\,tu - f(u)\,]$: the largest gap between the line $tu$ and the graph of $f$. Then:

• for every slope $t$ and every point $u$: $f(u) \ge tu - f^*(t)$ (the line is below $f$);
• for the best slope, $t = f'(u)$, the line touches $f$ at $u$ and the two sides are equal.

Now the key step. In a divergence, $u$ will be the density ratio $r(x)$ at a point $x$. We may choose a different slope at every point: a function $T(x)$. However badly we choose it, the line stays below $f$:

In section 3 we average this over $Q$, and the unknown ratio $r(x)$ disappears. What remains is a lower bound made of two averages over samples. The function $T$ is the critic.

Interactive figure: loading…

3. Definitions and main results #

Notation. $P$ and $Q$ are distributions on the same space, with densities (or probability mass functions) $p$ and $q$. Unless we say otherwise, we assume $p(x) > 0$ only where $q(x) > 0$. Then $r(x) = p(x)/q(x)$ is finite, and it is called the likelihood ratio or density ratio. A critic is a function $T(x)$ (or $T(x, y)$ for pairs). In this lesson the letter $f$ is always a convex generator, and critics are always called $T$.

f-divergences #

Definition. Let $f: [0, \infty) \to \mathbb{R}$ be convex with $f(1) = 0$. This $f$ is the generator. The f-divergence of $P$ from $Q$ is

The average is under the second argument $Q$, and the ratio has $p$ on top. The generator says how much a ratio $u$ is penalized; $u = 1$ (no difference) costs $f(1) = 0$.

Theorem (non-negativity). $D_f(P\,\|\,Q) \ge 0$, and $D_f(P\,\|\,P) = 0$.

Proof.

1. The ratio has average 1 under $Q$: $\mathbb{E}_Q[r] = \int q(x)\frac{p(x)}{q(x)}\,dx = \int p(x)\,dx = 1$.
2. Jensen's inequality for the convex $f$ (lesson 03): $\mathbb{E}_Q[f(r)] \ge f(\mathbb{E}_Q[r])$.
3. With step 1, the right side is $f(1) = 0$. If $P = Q$, then $r = 1$ everywhere and $D_f = \mathbb{E}_Q[f(1)] = 0$. $\blacksquare$

If $f$ is strictly convex at $u = 1$, then $D_f = 0$ only when $P = Q$. Every f-divergence also satisfies the data processing inequality of lesson 07.

Three facts that prevent confusion.

1. The generator is not unique. $f(u)$ and $f(u) + c(u - 1)$ give the same divergence for every constant $c$, because $\mathbb{E}_Q[r - 1] = 0$. For example, $u\log u - u + 1$ also gives KL, and it is $\ge 0$ for every $u$. Tables in papers differ by such terms and by constant factors.
2. Swapping the arguments. $D_f(Q\,\|\,P) = D_{\tilde f}(P\,\|\,Q)$ with $\tilde f(u) = u\,f(1/u)$. For $f(u) = u\log u$ this gives $\tilde f(u) = -\log u$: the reverse KL.
3. Points outside the support. A point with $q(x) = 0 < p(x)$ contributes $p(x)\lim_{u\to\infty} f(u)/u$. This is $+\infty$ for KL and $\chi^2$, but finite for TV, Hellinger and JS. A point with $p(x) = 0 < q(x)$ contributes $q(x)f(0)$, which is $+\infty$ for reverse KL.

The six standard generators. $M = \tfrac12(P + Q)$ is the mixture. The last two columns are explained below (conjugate and optimal critic).

Divergence	$D_f(P\,\Vert\,Q)$	Generator $f(u)$	Conjugate $f^*(t)$	Optimal critic $T^* = f'(r)$
KL	$\int p\log\frac{p}{q}$	$u\log u$	$e^{t-1}$	$1 + \log r$
reverse KL	$\int q\log\frac{q}{p}$	$-\log u$	$-1-\log(-t)$, for $t<0$	$-1/r$
total variation (TV)	$\frac12\int\lvert p-q\rvert$	$\frac12\lvert u-1\rvert$	$t$, for $\lvert t\rvert\le\frac12$	$\frac12\operatorname{sign}(r-1)$
Pearson $\chi^2$	$\int\frac{(p-q)^2}{q}$	$(u-1)^2$	$t + \frac{t^2}{4}$, for $t \ge -2$	$2(r-1)$
squared Hellinger $H^2$	$\int(\sqrt p-\sqrt q)^2$	$(\sqrt u-1)^2$	$\frac{t}{1-t}$, for $t<1$	$1 - 1/\sqrt r$
Jensen–Shannon (JS)	$\frac12\mathrm{KL}(P\Vert M)+\frac12\mathrm{KL}(Q\Vert M)$	$\frac12\big[u\log u-(u+1)\log\frac{u+1}2\big]$	$-\frac12\log(2-e^{2t})$, for $t<\frac12\log 2$	$\frac12\log\frac{2r}{1+r}$

Two inequalities link these divergences:

• $\mathrm{KL}(P\,\|\,Q) \le \log\big(1 + \chi^2(P\,\|\,Q)\big)$. Proof: $\mathrm{KL} = \mathbb{E}_P[\log r] \le \log\mathbb{E}_P[r]$ (Jensen, $\log$ is concave), and $\mathbb{E}_P[r] = \mathbb{E}_Q[r^2] = 1 + \chi^2$.
• Pinsker's inequality: $\mathrm{TV}(P, Q) \le \sqrt{\mathrm{KL}(P\,\|\,Q)/2}$, with KL in nats (no proof here).

The convex conjugate #

Definition. The convex conjugate (or Fenchel conjugate) of $f$ is

It is the largest gap between the line $tu$ and the graph of $f$. As a supremum of functions that are affine in $t$, $f^*$ is convex.

Example (KL), step by step. $f(u) = u\log u$.

1. Maximize $g(u) = tu - u\log u$. Its derivative is $t - \log u - 1$.
2. Set it to zero: $\log u = t - 1$, so $u = e^{t-1}$. ($g$ is concave, so this is the maximum.)
3. Insert this $u$: $f^*(t) = t e^{t-1} - e^{t-1}(t - 1) = e^{t-1}$.

Fenchel–Young inequality. For all $u \ge 0$ and all $t$: $f(u) \ge tu - f^*(t)$. Equality holds exactly when $t$ is a slope of $f$ at $u$ ($t = f'(u)$ when $f$ is differentiable). This is picture 2 in one line.

Fenchel–Moreau theorem. If $f$ is convex and lower semicontinuous (its epigraph is a closed set), then $f^{**} = f$, that is, $f(u) = \sup_t\,[tu - f^*(t)]$: $f$ is the envelope of its tangent lines. All generators in the table have this property.

The variational representation #

Theorem (variational lower bound). Let $f$ be a convex, lower semicontinuous generator with $f(1) = 0$, and let $p > 0$ only where $q > 0$. For every critic $T$ whose values lie where $f^*$ is finite (and with finite expectations):

Equality holds when $T(x)$ is a slope of $f$ at $r(x)$ for ($Q$-almost) every $x$, for example $T^*(x) = f'(r(x))$. Hence $D_f(P\,\|\,Q) = \sup_T\,\big\{\mathbb{E}_P[T] - \mathbb{E}_Q[f^*(T)]\big\}$.

Proof.

1. Fenchel–Young at $u = r(x)$ and $t = T(x)$: $f(r(x)) \ge T(x)\,r(x) - f^*(T(x))$ for every $x$.
2. Multiply by $q(x) \ge 0$ and integrate over $x$. The left side becomes $\int q\,f(r) = D_f(P\,\|\,Q)$.
3. On the right, $q(x)\,r(x) = p(x)$. So $\int q\,T\,r = \int p\,T = \mathbb{E}_P[T]$, and $\int q\,f^*(T) = \mathbb{E}_Q[f^*(T)]$. This gives the bound.
4. Step 1 is an equality at $x$ exactly when $T(x)$ is a slope of $f$ at $r(x)$. For $T^* = f'(r)$ this holds at every $x$, so the bound is an equality. $\blacksquare$

What the theorem gives you in practice.

• Only samples are needed. Replace $\mathbb{E}_P$ and $\mathbb{E}_Q$ by averages over samples of $P$ and of $Q$. No density is evaluated.
• Any critic is safe. With a restricted family of critics $\mathcal{T}$ (for example a neural network), $\sup_{T\in\mathcal{T}} \le D_f$. The gap is how well $\mathcal{T}$ can represent $T^* = f'(r)$.
• The best critic knows the density ratio: $r = (f')^{-1}(T^*)$. Training the critic is density-ratio estimation, as in picture 1.

KL: the NWJ and Donsker–Varadhan bounds #

NWJ bound. Put $f(u) = u\log u$ and $f^*(t) = e^{t-1}$ into the theorem. This gives the bound of Nguyen, Wainwright and Jordan (2010):

Check at $T^*$: $\mathbb{E}_P[1 + \log r] - \mathbb{E}_Q[e^{\log r}] = 1 + \mathrm{KL} - \mathbb{E}_Q[r] = 1 + \mathrm{KL} - 1 = \mathrm{KL}$.

Donsker–Varadhan (DV) bound.

Proof (DV is NWJ with the best constant shift).

1. Apply NWJ to the shifted critic $T + c$, with a constant $c$: $\mathrm{KL} \ge \mathbb{E}_P[T] + c - e^{c-1}\,\mathbb{E}_Q[e^T]$.
2. Maximize the right side over $c$. Its derivative is $1 - e^{c-1}\,\mathbb{E}_Q[e^T]$, which is zero at $c = 1 - \log\mathbb{E}_Q[e^T]$.
3. Insert this $c$: $\mathbb{E}_P[T] + 1 - \log\mathbb{E}_Q[e^T] - 1 = \mathbb{E}_P[T] - \log\mathbb{E}_Q[e^T]$. $\blacksquare$

So for the same critic, DV $\ge$ NWJ. DV does not change when you add a constant to $T$, and its optimal critics are $\log r + c$ for any constant $c$. MINE (Belghazi et al., 2018) maximizes DV with a neural critic.

Bounds on mutual information #

Mutual information is a KL divergence (lesson 07):

So put $P = p(x, y)$ (the joint) and $Q = p(x)p(y)$ (the product of the marginals) into NWJ and DV. The density ratio is $\frac{p(x,y)}{p(x)p(y)} = e^{\operatorname{pmi}(x,y)}$, where $\operatorname{pmi}(x, y) = \log\frac{p(x,y)}{p(x)p(y)} = \log\frac{p(y \mid x)}{p(y)}$ is the pointwise mutual information. The critic is a function $T(x, y)$ of a pair.

Where do the samples come from? Take a batch of $N$ pairs $(x_1, y_1), \dots, (x_N, y_N)$ drawn from $p(x, y)$ and form the score matrix $S_{ij} = T(x_i, y_j)$.

• The diagonal $S_{ii}$ uses true pairs: these are positive pairs, samples of $P = p(x,y)$.
• An off-diagonal entry $S_{ij}$, $i \ne j$, pairs $x_i$ with the $y$ of another example. Since the examples are independent, $(x_i, y_j)$ is a sample of $Q = p(x)p(y)$: a negative pair.

NWJ and DV average the diagonal for $\mathbb{E}_P$ and the off-diagonal for $\mathbb{E}_Q$.

Why these estimates are noisy when the MI is large. At the optimal NWJ critic, the $Q$-term averages $e^{T^* - 1} = r$ over negative pairs. One sample of it has variance

by the first inequality of the f-divergence part. So the variance grows exponentially with the mutual information. For correlated Gaussian pairs it is exactly $e^{2I} - 1$: about 164 at $I = 2.55$ nats, and about $1.3 \times 10^{10}$ at $I = 11.6$ nats. The average is dominated by rare negative pairs with a huge ratio, which most batches never see.

Two more warnings. The bounds hold for the expectations; a value computed on one batch is a noisy estimate and can lie above the true MI. And the DV estimate puts the sample average inside a $\log$. By Jensen, $\mathbb{E}[\log \hat m] \le \log \mathbb{E}[\hat m]$, so the DV estimate is biased upward and can overshoot the true value. MINE's estimate is therefore neither an upper nor a lower bound (Poole et al., 2019, section 2.2).

Interactive figure: loading…

InfoNCE: a bound from an N-way classification game #

InfoNCE (van den Oord et al., 2018) uses the same score matrix in a different way. For each row $i$, it asks a classification question: which of the $N$ candidates $y_1, \dots, y_N$ is the true partner of $x_i$? A softmax over the row turns the scores into probabilities, and the loss is the cross-entropy of the correct answer $j = i$:

This is exactly cross_entropy(S, labels = [0, 1, ..., N-1]), the softmax cross-entropy of lesson 04, with the row of scores as logits.

Interactive figure: loading…

Theorem (InfoNCE bound). Let $(x, y_1) \sim p(x, y)$, and let $y_2, \dots, y_N \sim p(y)$ be independent of each other and of $(x, y_1)$. For any critic $T$, let

Then

and the bound never passes the ceiling: $\log N - \mathcal{L}_{\text{InfoNCE}} \le \log N$.

Why it is true, in one picture. Hide the true partner at a random position $c$ among $N$ candidates. Guessing $c$ at random costs $\log N$ nats of cross-entropy. A good critic lowers the loss, and the saving $\log N - \mathcal{L}$ is information about $c$. But the only thing that reveals $c$ is the dependence between $x$ and its partner, so the saving cannot exceed $I(X;Y)$. And a cross-entropy is never negative, so the saving cannot exceed $\log N$.

Proof of the InfoNCE bound (five steps)

1. A classification game. Put the positive at a uniformly random position $c \in \{1, \dots, N\}$ instead of position 1. The softmax treats all positions alike, so the loss does not change. Write $D = (x, y_1, \dots, y_N)$. The softmax $q(c = i \mid D) = e^{T(x, y_i)}/\sum_j e^{T(x, y_j)}$ is a classifier that guesses $c$ from $D$, and $\mathcal{L}$ is its cross-entropy.
2. Cross-entropy $\ge$ entropy (lesson 04), for each $D$: $\mathcal{L} \ge H(c \mid D)$. Since $c$ is uniform, $H(c) = \log N$. So $\log N - \mathcal{L} \le H(c) - H(c \mid D) = I(c; D)$.
3. Chain rule for mutual information: $I(c; D) = I(c; y_{1:N}) + I(c; x \mid y_{1:N})$. Without $x$, the $y_j$ are independent draws of $p(y)$ whatever $c$ is, so $I(c; y_{1:N}) = 0$.
4. $I(c; x \mid y_{1:N}) = H(x \mid y_{1:N}) - H(x \mid c, y_{1:N})$. Given $c$ and all the $y$'s, $x$ depends only on its partner $y_c$, and $(x, y_c) \sim p(x, y)$. So $H(x \mid c, y_{1:N}) = H(X \mid Y)$. Conditioning reduces entropy, so $H(x \mid y_{1:N}) \le H(X)$. Together: $I(c; x \mid y_{1:N}) \le H(X) - H(X \mid Y) = I(X;Y)$. (For continuous $x$ the same steps hold with differential entropies.)
5. Steps 2–4 give $\log N - \mathcal{L} \le I(X;Y)$. A cross-entropy is $\ge 0$, so also $\log N - \mathcal{L} \le \log N$. $\blacksquare$

The optimal critic. Step 2 is tight when the softmax equals the true posterior of the position. That posterior is $p(c = i \mid D) \propto p(y_i \mid x)/p(y_i)$. So the optimal critic is

for any function $a(x)$ of $x$ alone: it adds the same number to a whole row, and the softmax removes it. With $T^*$ the bound tends to $I(X;Y)$ as $N \to \infty$. For a finite $N$ it stays below $I$, and never above $\log N$. To certify $I$ nats you need $N \ge e^{I}$.

Bias and variance, side by side.

Bound	Bias	Variance	Ceiling
InfoNCE	below $I$ for finite $N$, even with $T^*$	low: the softmax normalizes inside each row	$\log N$
NWJ	none at $T^*$: the sample mean is unbiased for the bound	at least $e^{I} - 1$ per negative pair ($e^{2I} - 1$ for Gaussians)	none
DV / MINE	the $\log$ of a sample mean is biased upward	high	none

This is a general limit, not a weakness of one method. McAllester and Stratos (2020) prove that any distribution-free lower bound on mutual information that holds with high confidence from $N$ samples can show at most $O(\log N)$.

Answers to section 1. (1) At most $\log 256 = 5.55$ nats, whatever the critic. The true value is 30 nats, but the bound cannot see more than $\log N$; you would need $N \ge e^{30} \approx 10^{13}$. (2) With equal numbers of each, the best classifier's logit is the log density ratio between true pairs $p(x,y)$ and shuffled pairs $p(x)p(y)$: the pointwise mutual information $\operatorname{pmi}(x, y)$. Its average over true pairs is $I(X;Y)$.

Contrastive learning #

In contrastive representation learning the critic is built from two encoders. Two views of the same example (two augmentations of an image in SimCLR, an image and its caption in CLIP, a context and a future frame in CPC) form a positive pair. The critic is a cosine similarity divided by a temperature $\tau$:

• CPC (van den Oord et al., 2018) introduced the name InfoNCE and the bound $I \ge \log N - \mathcal{L}$, with a log-bilinear critic $z^\top W c$.
• SimCLR (Chen et al., 2020) uses two augmentations of each image and the other $2N - 2$ views of the batch as negatives (the NT-Xent loss).
• CLIP (Radford et al., 2021) uses a symmetric loss: the average of the row loss (each image picks its caption among $N$) and the column loss (each caption picks its image), $\tfrac12\big(\mathcal{L}_{\text{img}\to\text{text}} + \mathcal{L}_{\text{text}\to\text{img}}\big)$.

Why large batches? The ceiling $\log N$ is one reason: with $N = 4096$ the bound stops at 8.3 nats. More negatives also make the classification task harder and the gradient less noisy.

A caution. A high InfoNCE bound does not by itself explain why the features are good. Tschannen et al. (2020) show encoders with the same or higher MI bound and worse downstream accuracy: the architecture of the critic and the augmentations matter as much as the bound. Augmentations, temperature, collapse and the practice of contrastive learning are in deep learning lesson 14.

GANs: the discriminator is a density-ratio classifier #

A GAN trains a generator $G$ with distribution $Q$ against a discriminator $D(x) \in (0, 1)$ that tells real samples (from the data $P$) from generated ones. The value function is

The discriminator maximizes $V$. This is the cross-entropy of picture 1 with a minus sign, so the best discriminator is the Bayes classifier with equal priors:

Theorem (GAN value at the optimal discriminator). $V(D^*, G) = 2\,\mathrm{JS}(P\,\|\,Q) - \log 4$.

Proof. Write $M = \frac12(P + Q)$, so $p + q = 2m$.

1. $V(D^*, G) = \int p \log\frac{p}{p+q} + \int q\log\frac{q}{p+q}$.
2. $\log\frac{p}{p+q} = \log\frac{p}{2m} = \log\frac{p}{m} - \log 2$, and the same for $q$.
3. So $V = \mathrm{KL}(P\,\|\,M) - \log 2 + \mathrm{KL}(Q\,\|\,M) - \log 2 = 2\,\mathrm{JS}(P\,\|\,Q) - \log 4$. $\blacksquare$

So the generator, facing an optimal discriminator, minimizes the Jensen–Shannon divergence. The minimum $-\log 4$ is reached exactly when $Q = P$ (then $D^* = \frac12$ everywhere). When $P$ and $Q$ hardly overlap, JS is close to its maximum $\log 2$ and hardly changes when $Q$ moves: the generator gets almost no gradient. This is one cause of unstable GAN training, and a motivation for the non-saturating loss and for Wasserstein GANs.

Interactive figure: loading…

f-GAN, and why TV leads to Wasserstein GANs

f-GAN (Nowozin, Cseke and Tomioka, 2016) trains a generator with any f-divergence. The critic is $T_\omega(x) = g_f(V_\omega(x))$: a network $V_\omega$ with a real output, followed by an output activation $g_f$ that maps $\mathbb{R}$ into the domain of $f^*$ (for example $g_f(v) = 1 - e^{-v} < 1$ for squared Hellinger). Training solves the saddle-point problem

$$\min_\theta\, \max_\omega\; \mathbb{E}_{x\sim P}[T_\omega(x)] - \mathbb{E}_{x\sim Q_\theta}\big[f^*(T_\omega(x))\big].$$

The inner max estimates $D_f(P\,\|\,Q_\theta)$ by the variational bound; the outer min moves $Q_\theta$ to reduce it.

The original GAN is one case, with $f_{\text{GAN}}(u) = u\log u - (u+1)\log(u+1)$ and $f^*_{\text{GAN}}(t) = -\log(1 - e^t)$ for $t < 0$. Write the critic as $T = \log D$ with $D \in (0, 1)$. Then $-f^*_{\text{GAN}}(T) = \log(1 - D)$, and the bound $\mathbb{E}_P[T] - \mathbb{E}_Q[f^*_{\text{GAN}}(T)]$ is exactly $V(D, G)$. The optimal critic $T^* = f'_{\text{GAN}}(r) = \log\frac{r}{1+r}$ gives $D^* = \frac{p}{p+q}$, and $D_{f_{\text{GAN}}}(P\,\|\,Q) = 2\,\mathrm{JS} - \log 4$. Here $f_{\text{GAN}}(1) = -\log 4 \ne 0$; this is why the minimum is $-\log 4$ and not 0.

TV is also an integral probability metric. Its conjugate is $f^*(t) = t$ on $|t| \le \frac12$, so the theorem gives $\mathrm{TV}(P, Q) = \sup_{|T| \le 1/2}\,\big(\mathbb{E}_P[T] - \mathbb{E}_Q[T]\big)$. Replace the constraint "$|T| \le \frac12$" by "$T$ is 1-Lipschitz" and you get the Wasserstein-1 distance of WGAN. Unlike JS, it keeps changing when two non-overlapping distributions move apart. GAN training in practice is in the generative models track.

4. Worked example #

(a) Six f-divergences of one pair. $P = (\tfrac12, \tfrac12)$ and $Q = (\tfrac14, \tfrac34)$ on two points. The ratios are $r = p/q = (2, \tfrac23)$, so every f-divergence is $D_f = \tfrac14 f(2) + \tfrac34 f(\tfrac23)$.

• KL: $\tfrac14 \cdot 2\log 2 + \tfrac34 \cdot \tfrac23\log\tfrac23 = \tfrac12\log 2 + \tfrac12\log\tfrac23 = \tfrac12\log\tfrac43 = 0.1438$ nats.
• Reverse KL: $\tfrac14(-\log 2) + \tfrac34(-\log\tfrac23) = -0.1733 + 0.3041 = 0.1308$ nats. Not equal to KL.
• TV: $\tfrac14 \cdot \tfrac12|2 - 1| + \tfrac34 \cdot \tfrac12|\tfrac23 - 1| = \tfrac18 + \tfrac18 = \tfrac14$.
• $\chi^2$: $\tfrac14(2-1)^2 + \tfrac34(\tfrac23 - 1)^2 = \tfrac14 + \tfrac1{12} = \tfrac13 = 0.3333$.
• Squared Hellinger: $(\sqrt{1/2} - \sqrt{1/4})^2 + (\sqrt{1/2} - \sqrt{3/4})^2 = 0.0429 + 0.0253 = 0.0681$ (or $0.0341$ with the factor $\tfrac12$).
• JS: $M = (\tfrac38, \tfrac58)$, $\mathrm{KL}(P\,\|\,M) = 0.0323$, $\mathrm{KL}(Q\,\|\,M) = 0.0354$, so $\mathrm{JS} = \tfrac12(0.0323 + 0.0354) = 0.0338$ nats, far below its maximum $\log 2 = 0.693$.

Checks: $\mathrm{KL} = 0.1438 \le \log(1 + \chi^2) = \log\tfrac43 = 0.2877$, and Pinsker: $\mathrm{TV} = 0.25 \le \sqrt{0.1438/2} = 0.268$.

(b) The KL bounds with three critics. A critic on two points is a pair $T = (T(1), T(2))$. NWJ is $\mathbb{E}_P[T] - \mathbb{E}_Q[e^{T-1}]$ and DV is $\mathbb{E}_P[T] - \log\mathbb{E}_Q[e^{T}]$.

Critic $T$	NWJ	DV
$1 + \log r = (1.693, 0.595)$	$1.1438 - 1 = 0.1438$	$1.1438 - \log e = 0.1438$
$\log r = (0.693, -0.405)$	$0.1438 - e^{-1} = -0.2240$	$0.1438 - \log 1 = 0.1438$
$(0, 0)$	$0 - e^{-1} = -0.3679$	$0 - \log 1 = 0$

For the first row: $\mathbb{E}_P[T] = \tfrac12(1.693) + \tfrac12(0.595) = 1.1438$, and $\mathbb{E}_Q[e^{T-1}] = \mathbb{E}_Q[r] = \tfrac14 \cdot 2 + \tfrac34 \cdot \tfrac23 = 1$. Every value is $\le \mathrm{KL} = 0.1438$. NWJ needs the exact constant in $1 + \log r$; DV ignores constants.

(c) The GAN on this pair. $D^* = p/(p+q) = (\tfrac{1/2}{3/4}, \tfrac{1/2}{5/4}) = (\tfrac23, \tfrac25)$. Then

and $2\,\mathrm{JS} - \log 4 = 0.0676 - 1.3863 = -1.3187$. The two agree.

(d) InfoNCE by hand. A batch of $N = 3$ pairs has the score matrix

The row loss is $\ell_i = -S_{ii} + \log\sum_j e^{S_{ij}}$.

• Row 1: $\log(e^3 + e^1 + e^0) = \log 23.80 = 3.1698$, so $\ell_1 = 3.1698 - 3 = 0.1698$. The softmax is $(0.844, 0.114, 0.042)$: the critic picks the right partner with probability 0.844.
• Row 2: $\log(e^0 + e^2 + e^1) = \log 11.107 = 2.4076$, so $\ell_2 = 2.4076 - 2 = 0.4076$.
• Row 3: all scores are equal, so the softmax is uniform and $\ell_3 = \log 3 = 1.0986$. This row gives no information: its saving $\log 3 - \ell_3$ is 0.

$\mathcal{L} = \tfrac13(0.1698 + 0.4076 + 1.0986) = 0.5587$, and the bound is $\log 3 - 0.5587 = 1.0986 - 0.5587 = 0.5399$ nats. Adding any constant to a whole row (a function $a(x_i)$) gives the same number. The bound can never pass $\log 3 = 1.0986$.

5. In code #

NumPy and SciPy check the numbers above, and a small simulation shows the $\log N$ ceiling. The notebook for this lesson has these cells and the exercises.

import numpy as np
from scipy.special import logsumexp, rel_entr
from scipy.spatial.distance import jensenshannon

p, q = np.array([0.5, 0.5]), np.array([0.25, 0.75])
r = p / q
gens = {"KL": lambda u: u * np.log(u), "reverse KL": lambda u: -np.log(u),
        "TV": lambda u: 0.5 * np.abs(u - 1), "chi2": lambda u: (u - 1) ** 2,
        "Hellinger2": lambda u: (np.sqrt(u) - 1) ** 2,
        "JS": lambda u: 0.5 * (u * np.log(u) - (u + 1) * np.log((u + 1) / 2))}
for name, f in gens.items():
    print(name, round(np.sum(q * f(r)), 4))
# KL 0.1438, reverse KL 0.1308, TV 0.25, chi2 0.3333, Hellinger2 0.0681, JS 0.0338
print(rel_entr(p, q).sum(), jensenshannon(p, q) ** 2)   # 0.1438 0.0338 (SciPy agrees)

# NWJ and DV with the critic T = log r: DV is tight, NWJ is not
T = np.log(r)
print(p @ T - q @ np.exp(T - 1), p @ T - np.log(q @ np.exp(T)))   # -0.2240  0.1438

# GAN value at the optimal discriminator = 2 JS - log 4
D = p / (p + q)
print(p @ np.log(D) + q @ np.log(1 - D), 2 * jensenshannon(p, q) ** 2 - np.log(4))  # -1.3187 -1.3187

# InfoNCE from a score matrix: log N - mean of the row cross-entropies
def infonce_bound(S):
    N = S.shape[0]
    return np.log(N) - np.mean(-np.diag(S) + logsumexp(S, axis=1))

S = np.array([[3., 1, 0], [0, 2, 1], [1, 1, 1]])
print(infonce_bound(S))                                   # 0.5399

# Correlated Gaussians in d dimensions: true MI = -d/2 log(1 - rho^2); optimal critic = pmi
def pmi_scores(x, y, rho):
    v = 1 - rho**2
    diff = y[None, :, :] - rho * x[:, None, :]
    return (-0.5 * (diff**2).sum(-1) / v - 0.5 * x.shape[1] * np.log(v)
            + 0.5 * (y**2).sum(-1)[None, :])

rng = np.random.default_rng(0)
d, rho, N = 10, 0.95, 128
x = rng.standard_normal((N, d))
y = rho * x + np.sqrt(1 - rho**2) * rng.standard_normal((N, d))
print(-d / 2 * np.log(1 - rho**2), np.log(N), infonce_bound(pmi_scores(x, y, rho)))
# 11.64 (true MI)  4.852 (log N)  4.851 (InfoNCE: stuck just below log N)

jensenshannon returns $\sqrt{\mathrm{JS}}$, so we square it. logsumexp computes $\log\sum_j e^{S_{ij}}$ without overflow, even for scores near 1000.

6. Where it shows up #

• Contrastive pretraining. CPC (audio, images), SimCLR and MoCo (images), CLIP (image–text), and contrastive sentence embeddings all minimize an InfoNCE loss. The learned critic $\cos(g(x), h(y))/\tau$ is used later for zero-shot classification and retrieval: CLIP picks the caption with the largest score, which is the softmax classifier of this lesson.
• Mutual information estimation. MINE (DV), NWJ and InfoNCE estimate the MI between high-dimensional variables, for example between layers of a network or between a representation and a label. Know their biases: InfoNCE is capped at $\log N$, NWJ and DV are very noisy at high MI.
• GANs and f-GANs. The discriminator is a critic of a variational bound on $2\,\mathrm{JS} - \log 4$; f-GAN does the same for any $f$. See the generative models track.
• Density-ratio estimation. A probabilistic classifier between two data sets gives $\log r$ as its logit. Uses: importance weights under covariate shift, and classifier two-sample tests: the best balanced accuracy is $\frac12(1 + \mathrm{TV})$, so a test accuracy $a$ shows $\mathrm{TV} \ge 2a - 1$ (up to the noise of the test).
• ML for science. Neural ratio estimation in simulation-based inference (Hermans et al., 2020, building on the classifier trick of Cranmer et al., 2015) trains a classifier between simulator pairs $(\theta, x)$ and shuffled pairs. Its logit estimates $\log\frac{p(x \mid \theta)}{p(x)}$, the pmi of this lesson with $P = p(\theta, x)$ and $Q = p(\theta)p(x)$. Multiplied by the prior, the ratio gives the posterior: $p(\theta \mid x) = p(\theta)\,e^{\operatorname{pmi}(\theta, x)}$. Checking an emulator: to compare a generative model of weather fields or seismic waveforms with the simulator, train a classifier between their samples; an accuracy near 50% means a small TV, as far as the classifier can tell. Cross-modal contrastive models such as AstroCLIP align galaxy images with spectra by a CLIP-style InfoNCE loss.

7. Common mistakes #

• Reading $\log N - \mathcal{L}$ as the mutual information. It is a lower bound, capped at $\log N$. A value near $\log N$ says "at least this much", not "exactly this much".
• Swapping $P$ and $Q$. $D_f(P\,\|\,Q)$ averages under $Q$ and uses $r = p/q$. With $f(u) = u\log u$ you get $\mathrm{KL}(P\,\|\,Q)$, not $\mathrm{KL}(Q\,\|\,P)$.
• Using the DV critic in NWJ. NWJ is not invariant to constant shifts: $T = \log r$ is optimal for DV but gives $\mathrm{KL} - e^{-1}$ in NWJ. The NWJ optimum is $1 + \log r$.
• Ignoring the domain of $f^*$. Reverse KL needs $T < 0$, squared Hellinger $T < 1$, TV $|T| \le \frac12$, the GAN generator $T < 0$. Outside, $f^*(t) = +\infty$ and the bound is $-\infty$. f-GAN uses an output activation to stay inside.
• Treating a one-batch estimate as a bound. Only the expected value is a bound. A minibatch DV estimate is biased upward; an NWJ estimate is unbiased but very noisy. Both can lie above the true value.
• Comparing generator tables without checking conventions. $f$ and $f + c(u - 1)$ give the same divergence, and factors of $\frac12$ (JS, Hellinger) differ between papers. Check $f(1) = 0$, or compute $D_f(P\,\|\,P) = 0$.
• Forgetting the class sizes. With $n_P \ne n_Q$ the classifier's logit is $\log r(x) + \log(n_P/n_Q)$, not $\log r(x)$.

8. Check yourself #

1. Show that the Pearson $\chi^2(P\,\|\,Q)$ equals $\mathrm{Var}_Q(r)$.
   Answer

   $\chi^2 = \mathbb{E}_Q[(r - 1)^2]$, and $\mathbb{E}_Q[r] = 1$. So it is the variance of $r$ under $Q$.
2. Derive $f^*$ and the optimal critic for $f(u) = (u - 1)^2$, and write the variational bound.
   Answer

   The derivative of $tu - (u-1)^2$ is $t - 2(u - 1)$, zero at $u = 1 + t/2$ (this needs $u \ge 0$, so $t \ge -2$). Then $f^*(t) = t(1 + t/2) - t^2/4 = t + t^2/4$. $T^* = f'(r) = 2(r - 1)$. Bound: $\chi^2(P\,\|\,Q) \ge \mathbb{E}_P[T] - \mathbb{E}_Q[T + T^2/4]$.
3. $P = \mathcal{N}(\mu, 1)$ and $Q = \mathcal{N}(0, 1)$. Give the KL, the optimal NWJ critic, and the variance of $e^{T^* - 1}$ under $Q$. About how many samples of $Q$ give a standard error of 0.1 on the $Q$-term when $\mathrm{KL} = 5$ nats?
   Answer

   $\mathrm{KL} = \mu^2/2$. $\log r(x) = \mu x - \mu^2/2$, so $T^*(x) = 1 + \mu x - \mu^2/2$. $\mathrm{Var}_Q(r) = \chi^2 = e^{\mu^2} - 1 = e^{2\,\mathrm{KL}} - 1$. For $\mathrm{KL} = 5$: $e^{10} - 1 \approx 22{,}025$. The standard error of a mean of $n$ samples is $\sqrt{22{,}025/n}$, which is 0.1 for $n \approx 2.2 \times 10^6$.
4. You classify samples of $P$ against samples of $Q$ (equal priors). What is the best possible accuracy? What is it for $P = (\tfrac12, \tfrac12)$ and $Q = (\tfrac14, \tfrac34)$?
   Answer

   Predict $P$ where $p > q$. The accuracy is $\frac12\sum_x \max(p, q) = \frac12(1 + \mathrm{TV})$, because $\max(p, q) = \frac{p+q}{2} + \frac{|p-q|}{2}$. Here $\frac12(\frac12 + \frac34) = 0.625 = \frac12(1 + 0.25)$.
5. A contrastive model with batch size 1024 reports $\log N - \mathcal{L} = 6.85$ nats. What can you conclude about $I(X;Y)$?
   Answer

   $\log 1024 = 6.93$. You can conclude $I \gtrsim 6.85$ nats (up to estimation noise), nothing more. The bound is near its ceiling, so the true MI may be much larger. A larger $N$ or another estimator is needed to say more.
6. Why is DV $\ge$ NWJ for the same critic $T$? When are they equal?
   Answer

   $\mathrm{DV}(T) = \max_c \mathrm{NWJ}(T + c) \ge \mathrm{NWJ}(T + 0)$. They are equal when the best shift is $c = 1 - \log\mathbb{E}_Q[e^T] = 0$, that is, when $\mathbb{E}_Q[e^T] = e$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What is the logit of the best classifier between samples of $P$ (label 1) and $Q$ (label 0), with equal class sizes?

   Show answer

2. Define the f-divergence $D_f(P\,\|\,Q)$.

   Show answer

3. Give the generators of KL and of total variation.

   Show answer

4. State the variational lower bound for an f-divergence.

   Show answer

5. What is the optimal critic in the variational bound, and why?

   Show answer

6. Write the Donsker–Varadhan bound on KL.

   Show answer

7. Why do NWJ and MINE estimates of MI get very noisy when the MI is large?

   Show answer

8. State the InfoNCE bound and its ceiling.

   Show answer

9. What is the optimal InfoNCE critic?

   Show answer

10. What is the GAN value $V(D^*, G)$ at the optimal discriminator $D^* = p/(p+q)$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How does contrastive learning relate to mutual information?" — The InfoNCE loss is a softmax cross-entropy that picks the positive among $N$ candidates; $\log N - \mathcal{L}$ is a lower bound on $I(X;Y)$ between the two views, capped at $\log N$.
• "Why do contrastive methods need many negatives?" — The bound cannot exceed $\log N$; to certify $I$ nats you need $N \ge e^{I}$. More negatives also make the task harder and the gradient less noisy.
• "What is the optimal discriminator of a GAN, and what does the generator minimize?" — $D^* = p/(p+q)$; the generator then minimizes $2\,\mathrm{JS} - \log 4$.
• "How do you estimate KL or MI when you only have samples?" — Train a critic and maximize a variational bound (NWJ, DV/MINE, InfoNCE), or train a classifier and use its logit as the log density ratio.
• "What is an f-divergence? Name four." — $\mathbb{E}_Q[f(p/q)]$ with convex $f$, $f(1) = 0$: KL ($u\log u$), reverse KL ($-\log u$), TV ($\frac12|u - 1|$), $\chi^2$ ($(u - 1)^2$).
• "Why can GAN training stall?" — When the supports hardly overlap, JS is close to $\log 2$ and flat, so the generator gets almost no gradient. Non-saturating losses and Wasserstein critics help.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. Compute six f-divergences from their generators for $P = (\frac12, \frac12)$, $Q = (\frac14, \frac34)$ and check them against SciPy.
2. Write infonce_lower_bound(scores) with logsumexp, and show that it never passes $\log N$.
3. Write nwj_lower_bound and dv_lower_bound from a score matrix, and compare the bias and variance of the three bounds on correlated Gaussians.
4. Estimate a log density ratio with scikit-learn's LogisticRegression on quadratic features and compare it with the true one.
5. Learn a quadratic critic $T(x, y) = a\,x^\top y - b\,\|y\|^2$ by gradient steps on the InfoNCE loss, and compare $(a, b)$ with the optimal critic.

11. Watch and read #

The Variational Lower Bound: Estimating a Divergence from Samples | Generative AI 05MyCampus · 12 min

The tangent-line argument of section 3, step by step, with a check on several critics.

Contrastive Learning with SimCLR | Deep Learning AnimatedDeepia · 15 min

How SimCLR builds positive pairs with augmentations and uses the InfoNCE (NT-Xent) loss.

OpenAI CLIP: ConnectingText and Images (Paper Explained)Yannic Kilcher · 48 min

The symmetric contrastive loss over an image–text similarity matrix, and zero-shot classification.

f-GAN: Training Generative Neural Samplers using Variational Divergence Minimization, NIPS 2016Preserve Knowledge · 34 min

The f-GAN talk: divergences, the variational bound with a critic, and output activations.

• Poole, Ozair, van den Oord, Alemi, Tucker, On Variational Bounds of Mutual Information (2019): §2.2 (DV, MINE, NWJ), §2.3 (InfoNCE and its $\log K$ ceiling), and Figures 2–3 (bias and variance on correlated Gaussians).
• van den Oord, Li, Vinyals, Representation Learning with Contrastive Predictive Coding (2018): the InfoNCE loss and the bound $I \ge \log N - \mathcal{L}$.
• Nowozin, Cseke, Tomioka, f-GAN (2016): §2 (f-divergences, conjugates and the variational bound), Table 1 (its JS generator gives $2\,\mathrm{JS}$).
• Nguyen, Wainwright, Jordan, Estimating divergence functionals and the likelihood ratio by convex risk minimization (2010): the origin of the NWJ bound.
• Belghazi et al., MINE: Mutual Information Neural Estimation (2018).
• McAllester, Stratos, Formal Limitations on the Measurement of Mutual Information (2020): the $O(\log N)$ limit.
• Chen, Kornblith, Norouzi, Hinton, A Simple Framework for Contrastive Learning of Visual Representations (SimCLR, 2020) and Radford et al., Learning Transferable Visual Models From Natural Language Supervision (CLIP, 2021).
• Tschannen, Djolonga, Rubenstein, Gelly, Lucic, On Mutual Information Maximization for Representation Learning (2020).
• Goodfellow et al., Generative Adversarial Networks (2014): the optimal discriminator and the JS result (Proposition 1, Theorem 1).
• Murphy, Probabilistic Machine Learning: Advanced Topics, free PDF: §2.7.1 (f-divergences), §5.3.6 (variational bounds on mutual information).