Mutual Information as an Objective: InfoNCE and Contrastive Bounds

The maximizer is the log density ratio, $f^\star(x,y) = \log\frac{p(y|x)}{p(y)} + c$, and at that critic the bound touches $I(X;Y)$ exactly. One obstacle remains: $\log Z(y)$ needs an expectation over the entire marginal $p(x)$, and no closed form exists.

Replacing the partition function with negatives

Sample $N-1$ negatives from $p(x)$. Because $\log$ is concave, Jensen gives $\log Z \ge \mathbb{E}[\log \frac{1}{N-1}\sum_j e^{f(x,y_j)}]$, and substituting that lower estimate for $\log Z$ keeps the inequality valid while adding slack. The leftovers collapse into a softmax:

$$ I(X;Y) \ge \log(N-1) - \mathcal{L}{\text{InfoNCE}}, \qquad \mathcal{L}{\text{InfoNCE}} = -\mathbb{E}\left[\log \frac{e^{f(x,y)}}{\sum_{j=1}^{N} e^{f(x,y_j)}}\right] $$

$\mathcal{L}_{\text{InfoNCE}}$ is a cross-entropy over $N$ candidates with the positive as the label. From the previous lesson it is therefore entropy plus a KL fit to the true candidate posterior. Two consequences follow at once: the loss is non-negative, so no value of it can certify more than $\log(N-1)$ nats, and the bound tightens with $N$ even when the loss does not move.

flowchart TD
    A["🔗<br/>MI = KL(joint, marginals)"] --> B["🧩<br/>Critic f, density p(x)e^f over Z"]
    B --> C["📉<br/>E[f] minus E[log Z]"]
    C --> D["🎲<br/>log Z from N-1 negatives"]
    D --> E["🎯<br/>bound = log(N-1) - L"]

Worked example. Take $N = 4$, $\tau = 1$, and critic scores $[2.0, 0.5, -1.0, 0.0]$ with the positive first. The exponentials are $[7.389, 1.649, 0.368, 1.000]$, summing to $10.406$, so the positive holds mass $0.710$ and $\mathcal{L} = 0.343$ nats. The certified MI is $\log 3 - 0.343 = 0.756$ nats. Now move the positive to the score $-1.0$ slot: its mass drops to $0.035$, $\mathcal{L} = 3.342$, and the bound becomes $-2.24$. The inequality still holds, since MI is non-negative, but the bound is vacuous. A badly ranked critic does not produce a wrong bound, it produces an uninformative one.

Worked example, batch size. The ceiling $\log(N-1)$ is the whole story of why numbers from different runs are not comparable:

$N$ candidates	$\log(N-1)$ nats	max certifiable MI	bound if $L = 3.0$
64	4.14	5.98 bits	1.14 nats
256	5.54	8.00 bits	2.54 nats
1024	6.93	10.00 bits	3.93 nats
8192	9.01	13.00 bits	6.01 nats

A loss that plateaus at $\log(\text{batch})$ has saturated the estimator, not the data.

Check: a run reports $L = 2.1$ nats with 512 candidates. What does it certify, and with 64 candidates?

With 512: $\log 511 - 2.1 = 6.24 - 2.1 = 4.14$ nats (5.97 bits). With 64: $\log 63 - 2.1 = 4.14 - 2.1 = 2.04$ nats (2.95 bits). Same loss, half the certified information.

Reading an implementation as a bound

def info_nce(z1, z2, tau):
    z1, z2 = F.normalize(z1, dim=-1), F.normalize(z2, dim=-1)
    logits = (z1 @ z2.T) / tau                            # critic family: scaled cosine
    target = torch.arange(z1.shape[0], device=z1.device)   # diagonal holds the positive
    return F.cross_entropy(logits.float(), target)         # fp32 log-sum-exp

bound_nats = math.log(batch_size - 1) - loss.item()        # the number that means something

The critic family is restricted to scaled cosine similarity, so the bound is only as tight as that family permits. Temperature $\tau$ rescales the critic; it is not free sharpening. Since the optimum is a log ratio, shrinking $\tau$ emphasizes hard negatives and inflates logit magnitude, trading bound looseness against gradient variance.

Warning: logits.float() is load bearing. With $\tau = 0.07$ and cosine similarities in $[-1, 1]$, logits reach about $\pm 14$; $e^{14.3} \approx 1.6\times10^{6}$, and summing 1024 such terms exceeds fp16's 65504 limit, producing inf and then NaN through $\log(\text{inf}) - \text{inf}$. Max-subtraction keeps the largest term at 1 and the sum at most $N$, so the log-sum-exp must run in fp32 regardless of the autocast context.

Two more silent looseners. Under data-parallel training with a per-device softmax, $N$ is the local batch, so the bound uses $\log(\text{local}-1)$, not the global batch. And negatives drawn from the marginal can be other views of the same underlying item, which biases the estimate downward.

Rules of thumb for contrastive objectives

• Convert before comparing: report $\log(N-1) - L$, never $L$ alone.
• State the effective negatives: local batch, gathered batch, or queue size.
• Treat $\tau$ as a bias-variance dial on the critic, and tune it with the bound in view.
• If the loss approaches zero while downstream metrics still improve, the critic has saturated the finite-batch ceiling.
• Verify numerics by feeding deliberately large logits: if the loss is not finite, the log-sum-exp is in the wrong dtype.

Carrying the bound forward

InfoNCE is a variational lower bound on a KL between a joint and a product of marginals, estimated by a softmax over in-batch negatives, capped at $\log(N-1)$ nats and loosened by every shortcut in the critic family, the batch sharing, and the precision.

Mutual information is the KL divergence between the _____ and the product of the marginals. The optimal critic in the variational bound is the log density ratio _____. Replacing the log partition function with a sample average over negatives makes the bound _____, because log is concave. InfoNCE lower bounds mutual information by _____, so no loss value can certify more than _____ nats. Under a per-device softmax, the effective number of negatives is the _____.

The softmax inside the loss is itself a KL between the candidate posterior and the model's scores, and which way that divergence pulls on the learned distribution is a choice, not a detail. The next lesson takes up forward versus reverse KL and shows why one covers modes while the other seeks a single one.

Front: Write mutual information as a KL divergence. Back: I(X;Y) = KL(p(x,y) || p(x)p(y)), the divergence between the joint and the product of marginals.

Front: When is I(X;Y) exactly zero? Back: Exactly when X and Y are independent, so the joint factorizes into the marginals.

Front: What is the optimal critic for the variational MI bound? Back: f*(x,y) = log [p(y|x)/p(y)] plus a constant, the log density ratio.

Front: State the InfoNCE lower bound in terms of the loss. Back: I(X;Y) >= log(N-1) - L_InfoNCE, with N total candidates including the positive.

Front: Why does the InfoNCE bound saturate at log N? Back: The loss is a cross-entropy over N candidates and is never negative, so no loss value can certify more than log N nats.

Front: What role does Jensen's inequality play in the InfoNCE derivation? Back: Log is concave, so replacing log Z by the log of a sample average of exp(f) yields a further, looser lower bound.

Front: Which distribution does the InfoNCE softmax approximate at the optimum? Back: The posterior over candidates, p(y_j|x) proportional to p(y_j) exp(f(x,y_j)).

Front: Does a larger batch size lower the InfoNCE loss by itself? Back: No. It tightens the bound: the certified MI rises even when the reported loss is unchanged.

Front: Compute log(N-1) for N = 256, and its ceiling in bits. Back: About 5.54 nats, or 8 bits, the maximum MI a 256-way contrastive task can certify.

Front: What is the numerical risk of a small temperature in fp16? Back: Logits scale with 1/tau, so exp and the negative sum overflow to inf; use max-subtracted log-sum-exp and compute the loss in fp32.

Front: Why does a per-device softmax underestimate the certified MI on multiple GPUs? Back: Negatives are not shared across devices, so the effective N is the local batch, not the global one.

Front: What does the symmetric form of InfoNCE add? Back: The transpose term, a second log-softmax over the other axis, averaging two estimates of the same bound.

Front: Is mutual information invariant under invertible reparametrization of X? Back: Yes, so it measures only shared information, not the scale or coordinates of a representation.

Front: What does the temperature do in a score/tau cross-entropy? Back: It rescales the critic, trading bound looseness against gradient variance and sensitivity to hard negatives.