Information Theory for Language advanced 8 min read 10 flashcards

f-Divergences Beyond KL

KL is one member of a family generated by a convex function, the choice of member decides whether your model covers the data or collapses onto a mode, and some tasks need a divergence that is not in the family at all.

Distillation minimises forward KL. RLHF penalises reverse KL against the reference policy. GAN training, in its original derivation, minimises Jensen-Shannon. Diffusion models minimise a weighted sum of KLs. These are four different choices, they produce visibly different failure modes, and they are all points in one family.

The family is due to Csiszár and to Ali and Silvey, independently in the mid-1960s. Given a convex function \(f\) with \(f(1) = 0\),

\[ D_f(P \parallel Q) = \mathbb{E}_{q}\left[ f\!\left( \frac{p(x)}{q(x)} \right) \right] = \int q(x)\, f\!\left(\frac{p(x)}{q(x)}\right) dx \]

Every member is non-negative and vanishes only when \(P = Q\) almost everywhere, both by Jensen's inequality. The choice of \(f\) decides everything else.

\(f(t)\) Divergence Character
\(t \log t\) Forward KL, \(D(P \parallel Q)\) Mass-covering; infinite where \(p > 0, q = 0\)
\(-\log t\) Reverse KL, \(D(Q \parallel P)\) Mode-seeking; infinite where \(q > 0, p = 0\)
\(\frac{1}{2}\lvert t - 1\rvert\) Total variation Bounded by 1; a true metric
\((\sqrt{t}-1)^2\) Squared Hellinger Bounded; a metric; symmetric
\((t-1)^2\) \(\chi^2\) Heavily penalises ratio spikes
\(t\log t - (t{+}1)\log\frac{t+1}{2}\) Jensen-Shannon Symmetric, bounded by 1 bit

Mass-covering versus mode-seeking

The most consequential distinction in practice is the asymmetry of KL, and it follows directly from where each version blows up.

Forward KL, \(D(P \parallel Q)\) with \(P\) the data, integrates against \(p\). Any region where the data has mass and the model has none contributes \(+\infty\). A model minimising forward KL therefore stretches to cover every mode, including regions between modes where no real data lives. This is why maximum likelihood training, which is forward KL, produces language models that will happily generate the average of two incompatible styles.

Reverse KL integrates against \(q\). Regions where the model has mass and the data has none are what cost. A model minimising reverse KL contracts onto one mode and ignores the rest, safely. This is why reverse KL is the right choice for RLHF's regularisation term, where the goal is to keep the policy inside the reference model's support rather than to cover it, and why variational inference with reverse KL under-estimates posterior variance.

Neither is a bug. Ask which failure you prefer: hallucinating in the gaps, or losing diversity.

Why the family is trainable

The practical unlock is the variational representation. Every \(f\)-divergence has a lower bound in terms of the convex conjugate \(f^{\star}\):

\[ D_f(P \parallel Q) \ \ge\ \sup_{T} \left( \mathbb{E}_{p}[T(x)] - \mathbb{E}_{q}[f^{\star}(T(x))] \right) \]

where \(T\) ranges over functions. Parameterise \(T\) with a neural network and you can estimate, and minimise, any \(f\)-divergence from samples alone, without densities. That is exactly the f-GAN construction, which showed the original GAN objective is the Jensen-Shannon member of the family and that swapping \(f\) swaps the training dynamics in predictable ways (Nowozin et al., 2016, f-GAN, arXiv:1606.00709, NeurIPS 2016).

When it breaks

The family has a hard structural limit: \(f\)-divergences depend only on the ratio \(p/q\) pointwise, so they are blind to the geometry of the sample space. If \(P\) and \(Q\) are two Diracs on different points, every \(f\)-divergence returns its maximum value, whether the points are adjacent or a mile apart, and every gradient is zero or undefined. Disjoint supports are common in high dimensions, and this is precisely the vanishing-gradient failure that motivated moving away from JS-based GAN objectives.

The alternative family is the integral probability metrics, which compare expectations of test functions rather than density ratios: Wasserstein distance and maximum mean discrepancy are the two that matter. They see geometry, and they stay finite and differentiable under disjoint support. Total variation is the only divergence in both families. When your problem involves distributions that may not overlap, that is the signal to leave the \(f\)-family entirely. See optimal transport and Wasserstein and cross-entropy and KL.

Check yourself

10 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track