← Back to all blogs

How to Do Distribution Matching

Differences, ratios, and the awkward fact that you only have samples.

A surprising number of problems in deep learning reduce to the same statement: make this distribution look like that one. Train a generator so its samples look like data. Train an encoder so features from two domains are indistinguishable. Train a one-step student so it produces the same distribution as a fifty-step diffusion teacher. Train a small language model to match a large one. Align a policy with a reward-tilted reference. Compress a dataset into a handful of synthetic examples with the same statistics.

All of these are the problem

$$\min_{\theta}\; D\big(P \,\Vert\, Q_\theta\big),$$

where $P$ is a target we can only sample from and $Q_\theta$ is a model we control. The whole difficulty is that $D$ is defined in terms of densities we do not have. This post is about the machinery people use to get around that, starting from the taxonomy that organizes it — comparing distributions by difference or by ratio — and ending with where each choice shows up in systems today.

Contents

  1. The problem and the notation
  2. Two ways to compare: difference and ratio
  3. Why the choice matters: a two-line counterexample
  4. Integral probability metrics in detail
  5. $\phi$-divergences and the density-ratio trick
  6. Estimation: where the real difficulty lives
  7. Algorithms: GMMN, WGAN, MMD-GAN, f-GAN, Sinkhorn
  8. Critic-free matching: scores and transport
  9. Distribution matching beyond generative modeling
  10. A practical guide to choosing
  11. References

1. The problem and the notation

Let $\mathcal{X}$ be the sample space. We have two distributions on it: a target $P$, given through samples $x_1,\dots,x_n \sim P$, and a model $Q_\theta$ that we can sample from, usually by pushing a simple prior through a network, $y = g_\theta(z)$ with $z \sim \mathcal{N}(0,I)$. We want to move $\theta$ so that $Q_\theta$ approaches $P$.

Three properties determine whether a given discrepancy $D$ is usable for this. It must be a valid divergence ($D(P,Q) \ge 0$ with equality iff $P = Q$, at least for the distributions we care about); it must be estimable from finite samples, because that is all we have; and — most restrictive of all — the map $\theta \mapsto D(P, Q_\theta)$ must be continuous and differentiable with useful gradients. Section 3 is about how badly the third condition can fail.

Notation

SymbolMeaning
$P, Q$Two probability distributions on $\mathcal{X}$; $p, q$ their densities when they exist.
$Q_\theta$, $g_\theta$Model distribution and the generator that produces it, $y = g_\theta(z)$, $z \sim p(z)$.
$\mathcal{F}$The witness or critic class: a set of functions $f : \mathcal{X}\to\mathbb{R}$.
$f_\psi$A critic parameterized by a network with weights $\psi$.
$\Vert f\Vert_L$Lipschitz seminorm, $\sup_{x \ne y}\, |f(x)-f(y)| / \Vert x-y\Vert$.
$k$, $\mathcal{H}_k$A positive-definite kernel and its reproducing-kernel Hilbert space (RKHS).
$\mu_P$Kernel mean embedding, $\mu_P = \mathbb{E}_{X\sim P}[k(\cdot, X)] \in \mathcal{H}_k$.
$\Pi(P,Q)$Couplings: joint distributions on $\mathcal{X}\times\mathcal{X}$ with marginals $P$ and $Q$.
$\phi$, $\phi^*$A convex generator with $\phi(1)=0$, and its convex conjugate $\phi^*(t) = \sup_u\{ut - \phi(u)\}$.
$r(x)$The density ratio $p(x)/q(x)$ — the central object of §5.
$s_P(x)$The score $\nabla_x \log p(x)$ — the central object of §8.

2. Two ways to compare: difference and ratio

To compare two numbers you either subtract them or divide them. Remarkably, that naive dichotomy is a real taxonomy for probability discrepancies, and it is the one Müller formalized [1].

2.1 Integral probability metrics: take the difference

You cannot subtract two distributions directly, but you can subtract the expectations they assign to a test function. Pick a class $\mathcal{F}$ of witness functions and look for the one that separates them most:

$$D_{\mathcal{F}}(P,Q) \;=\; \sup_{f \in \mathcal{F}}\;\Big|\, \mathbb{E}_{X\sim P}[f(X)] - \mathbb{E}_{Y\sim Q}[f(Y)] \,\Big|.$$

If $\mathcal{F}$ is symmetric ($f \in \mathcal{F} \Rightarrow -f \in \mathcal{F}$), which all the classes below are, the absolute value can be dropped. The maximizing $f$ is the witness function: it is large where $P$ has more mass than $Q$ and small where the reverse holds, so it is literally a picture of how the two differ.

Everything about an IPM is decided by $\mathcal{F}$, and the choices are exactly the well-known metrics:

Witness class $\mathcal{F}$Resulting IPMComment
$\Vert f\Vert_L \le 1$Wasserstein-1 $W_1$Kantorovich–Rubinstein dual of earth-mover's distance.
$\Vert f\Vert_{\mathcal{H}_k} \le 1$Maximum mean discrepancyThe sup has a closed form. This is the crucial practical difference.
$\Vert f\Vert_\infty \le 1$Total variation (times 2)Too strong: sees any difference in support.
$\Vert f\Vert_\infty + \Vert f\Vert_L \le 1$Dudley metricMetrizes weak convergence.
Indicators of half-spacesKolmogorov–SmirnovThe classical one-dimensional two-sample test.
$\{f_\psi : \psi \in \Psi\}$, a network"Neural net distance"What a GAN critic actually computes. See §4.4 for the catch.

The reason IPMs are attractive for deep learning is visible in the definition: it only ever involves expectations, which are exactly what samples estimate. No density appears anywhere.

2.2 $\phi$-divergences: take the ratio

The other option is to compare $P$ and $Q$ pointwise through their ratio, then average the result under $Q$ with some convex penalty $\phi$ (convex, with $\phi(1)=0$):

$$D_\phi(P,Q) \;=\; \int_{\mathcal{X}} q(x)\, \phi\!\left(\frac{p(x)}{q(x)}\right) dx \;=\; \mathbb{E}_{Q}\!\left[\phi\big(r(X)\big)\right].$$

Convexity plus Jensen gives $D_\phi \ge \phi(1) = 0$ immediately. Different $\phi$ recover every classical divergence:

DivergenceGenerator $\phi(u)$Behaviour
Kullback–Leibler $\mathrm{KL}(P\Vert Q)$$u\log u$Mass-covering; $+\infty$ if $Q$ misses support of $P$.
Reverse KL $\mathrm{KL}(Q \Vert P)$$-\log u$Mode-seeking; $+\infty$ if $Q$ puts mass where $P$ has none.
Jensen–Shannon$u\log\frac{2u}{u+1} + \log\frac{2}{u+1}$Symmetric, bounded by $\log 4$. What the vanilla GAN minimizes.
Squared Hellinger$(\sqrt{u}-1)^2$Bounded, symmetric, a true metric after a square root.
Pearson $\chi^2$$(u-1)^2$Heavily penalizes $q \ll p$; high variance to estimate.
Total variation$\tfrac12|u-1|$The one member of both families.

Note the structural problem, which is the mirror image of the IPM's virtue: $D_\phi$ is written in terms of $r = p/q$, a quantity that cannot be estimated from samples without an additional model. Section 5 is about how to get it anyway.

2.3 Total variation sits in the intersection

Total variation is the largest possible disagreement in probability that $P$ and $Q$ assign to any single event:

$$\delta(P,Q) = \sup_{A \in \mathcal{F}_{\!\sigma}} \big|P(A) - Q(A)\big| = \tfrac12\int |p(x)-q(x)|\,dx .$$

It is an IPM (take $\mathcal{F}$ to be indicators of measurable sets, or bounded functions) and a $\phi$-divergence ($\phi(u) = \frac12|u-1|$). Sriperumbudur et al. [2] showed this is essentially the only overlap: the two families meet at TV and nowhere else. That is worth internalizing, because it means the two constructions really are different tools and not two views of one thing.

TV also shows why "valid divergence" is not the same as "useful objective." TV is a perfectly good metric that is useless for gradient descent, for the reason in the next section.

3. Why the choice matters: a two-line counterexample

Here is the example from the WGAN paper [3], and it is the single clearest argument in this whole literature. Let $Z \sim \mathcal{U}[0,1]$ and consider two distributions on $\mathbb{R}^2$: $P_0$ is the law of $(0, Z)$ — a vertical unit segment at the origin — and $P_\theta$ is the law of $(\theta, Z)$, the same segment shifted horizontally by $\theta$. Two parallel lines. As $\theta \to 0$ they visibly converge, so any sane notion of distance should shrink continuously to zero.

Compute what each discrepancy actually reports. For $\theta \ne 0$ the two segments have disjoint support, and:

$$D_{\mathrm{JS}}(P_\theta, P_0) = \begin{cases}\log 2 & \theta \ne 0\\ 0 & \theta = 0\end{cases} \qquad \mathrm{KL}(P_\theta \Vert P_0) = \mathrm{KL}(P_0 \Vert P_\theta) = \begin{cases}+\infty & \theta \ne 0\\ 0 & \theta = 0\end{cases}$$ $$\delta(P_\theta, P_0) = \begin{cases}1 & \theta \ne 0\\ 0 & \theta = 0\end{cases} \qquad\text{but}\qquad \boxed{\,W_1(P_\theta, P_0) = |\theta|\,}$$

The $\phi$-divergences are constant in $\theta$ (or infinite), so $\partial D/\partial\theta = 0$ almost everywhere: gradient descent receives no information about which way to move. Wasserstein gives $|\theta|$, whose derivative is $\pm 1$ — exactly the signal you want. The reason is structural. A $\phi$-divergence reads $P$ and $Q$ pointwise through $p/q$, so it is blind to how far apart the supports are; an IPM's witness function is constrained to be smooth, so it cannot spike arbitrarily and must report distance in the geometry of $\mathcal{X}$.

Why this is not a contrived example. Natural image data is widely believed to lie on a low-dimensional manifold in pixel space, and a generator $g_\theta : \mathbb{R}^{d}\to\mathbb{R}^{D}$ with $d \ll D$ produces samples supported on a $d$-dimensional manifold too. Two such manifolds in a high-dimensional ambient space generically intersect in a set of measure zero — that is, they are almost always disjoint. The pathological case is the typical case, which is why the original GAN's discriminator saturates: once $D$ can separate real from fake perfectly, the generator's gradient vanishes.

The same argument disqualifies MMD if the kernel is chosen badly. A Gaussian kernel $k(x,y) = \exp(-\Vert x-y\Vert^2 / 2\sigma^2)$ with $\sigma$ much smaller than the typical inter-point distance evaluates to $\approx 0$ for every cross pair, so $\widehat{\mathrm{MMD}}^2 \to$ constant and the witness function is a sum of narrow bumps with no gradient in between. Bandwidth is not a tuning detail; it is what decides whether the metric is smooth on the scale of your data.

4. Integral probability metrics in detail

4.1 Wasserstein distance and optimal transport

The primal definition is a transport problem: how much work does it take to rearrange the pile of dirt $P$ into the pile $Q$?

$$W_p(P,Q)^p = \inf_{\gamma \in \Pi(P,Q)}\; \mathbb{E}_{(x,y)\sim\gamma}\big[\Vert x - y\Vert^p\big].$$

The infimum runs over all couplings — joint distributions whose marginals are $P$ and $Q$ — each of which is a transport plan saying how much mass moves from $x$ to $y$. This is a linear program, and for $p = 1$ its dual is Kantorovich–Rubinstein duality [4]:

$$W_1(P,Q) = \sup_{\Vert f\Vert_L \le 1}\; \mathbb{E}_P[f(X)] - \mathbb{E}_Q[f(Y)], \qquad \Vert f\Vert_L := \sup_{x\ne y}\frac{|f(x)-f(y)|}{\Vert x - y\Vert}.$$

The dual is what makes $W_1$ trainable: the transport plan (an $n\times n$ object) disappears and is replaced by a single scalar function that a network can represent. Two things follow that shape everything downstream. First, the critic must be Lipschitz, and enforcing that constraint is the entire engineering problem of §7.2. Second, the sup is over an enormous class, so in practice the network never attains it, and what we optimize is a lower bound of unknown tightness.

4.2 Maximum mean discrepancy

Now take $\mathcal{F}$ to be the unit ball of an RKHS $\mathcal{H}_k$. Something very convenient happens. By the reproducing property $f(x) = \langle f, k(\cdot,x)\rangle$, so

$$\mathbb{E}_P[f(X)] - \mathbb{E}_Q[f(Y)] = \big\langle f,\; \mu_P - \mu_Q \big\rangle_{\mathcal{H}_k}, \qquad \mu_P := \mathbb{E}_{X\sim P}\big[k(\cdot, X)\big].$$

Maximizing a linear functional over a unit ball is trivial — align $f$ with the vector — so the supremum is attained in closed form by $f^\star \propto \mu_P - \mu_Q$, giving

$$\mathrm{MMD}(P,Q) = \sup_{\Vert f\Vert_{\mathcal{H}_k}\le 1} \mathbb{E}_P[f] - \mathbb{E}_Q[f] = \big\Vert \mu_P - \mu_Q \big\Vert_{\mathcal{H}_k}.$$

Expanding the squared norm turns it into three expectations of the kernel, with no optimization at all [5]:

$$\mathrm{MMD}^2(P,Q) = \mathbb{E}_{X,X'\sim P}\,k(X,X') + \mathbb{E}_{Y,Y'\sim Q}\,k(Y,Y') - 2\,\mathbb{E}_{X\sim P, Y\sim Q}\,k(X,Y).$$

This is the single most useful formula in the sample-based-comparison toolbox: an IPM whose supremum you never have to solve for. Given a batch of $n$ real and $m$ generated points, the plug-in (V-statistic) estimator is

$$\widehat{\mathrm{MMD}}^2 = \frac{1}{n^2}\sum_{i,j} k(x_i, x_j) + \frac{1}{m^2}\sum_{i,j}k(y_i,y_j) - \frac{2}{nm}\sum_{i,j}k(x_i,y_j),$$

and dropping the diagonal terms $i = j$ from the first two sums gives the standard unbiased U-statistic version. Everything is differentiable in $\theta$ through $y_j = g_\theta(z_j)$.

MMD is a genuine metric — $\mathrm{MMD}(P,Q) = 0 \iff P = Q$ — precisely when $k$ is characteristic, meaning $P \mapsto \mu_P$ is injective. Gaussian and Laplace kernels are characteristic on $\mathbb{R}^d$; a linear kernel is not (it only matches means), and a polynomial kernel of degree $d$ only matches moments up to order $d$. There is a nice way to see what MMD is doing: expand the Gaussian kernel and it compares an infinite weighted collection of moments at once. Distribution matching by moment matching, with the kernel choosing the weighting.

4.3 Energy distance, Cramér, and the sliced trick

Two useful relatives. The energy distance

$$\mathcal{E}(P,Q) = 2\,\mathbb{E}\Vert X - Y\Vert - \mathbb{E}\Vert X - X'\Vert - \mathbb{E}\Vert Y - Y'\Vert$$

looks like a different object but is exactly an MMD with the kernel $k(x,y) = \Vert x\Vert + \Vert y\Vert - \Vert x - y\Vert$ [6] — the statistics and machine learning literatures independently invented the same quantity. Its one-dimensional version is the Cramér distance, which motivated Cramér GAN [7]; the selling point there is that unlike the Wasserstein estimator, the Cramér estimator has unbiased sample gradients.

The sliced Wasserstein distance exploits the fact that in one dimension, optimal transport is just sorting ($W_p$ between 1-D empirical distributions is a sort followed by a sum). Project onto random directions and average:

$$\mathrm{SW}_p^p(P,Q) = \int_{\mathbb{S}^{d-1}} W_p^p\big(\pi^\omega_\# P,\; \pi^\omega_\# Q\big)\, d\omega \;\approx\; \frac{1}{L}\sum_{\ell=1}^{L} W_p^p\big(\pi^{\omega_\ell}_\# P,\; \pi^{\omega_\ell}_\# Q\big).$$

Cost is $O(Ln\log n)$ with no critic and no adversarial game. It is a valid metric, but the number of random projections needed to detect a difference grows badly with dimension, so it is used mostly in latent spaces or as an auxiliary loss.

4.4 The neural net distance and its catch

In practice nobody optimizes over all 1-Lipschitz functions; they optimize over $\{f_\psi\}$, a finite-capacity network with some constraint. Arora et al. [8] analyzed this "neural net distance" and found a genuinely worrying result: with a critic of bounded capacity, $D_{\text{nn}}(P, Q_\theta)$ can be near zero even when $P$ and $Q_\theta$ are wildly different. Concretely, a distribution supported on only $O(\text{capacity}/\epsilon^2)$ points can fool the critic into reporting $\epsilon$-closeness. The metric you optimize is not the metric you named, and a low critic loss is not evidence of a good generator — it may be evidence of mode collapse the critic cannot see.

This is a general lesson about adversarial distribution matching. The sup in an IPM is a promise you cannot keep with a finite network, and the gap between the true metric and the achieved one is unknown and changes during training. It is one motivation for MMD-style methods, where the sup is exact by construction and the only approximation is the sample average.

5. $\phi$-divergences and the density-ratio trick

If $\phi$-divergences need $r = p/q$ and we only have samples, they look unusable. The escape is that a classifier trained to separate the two samples estimates the ratio for you. This one observation is the foundation of GANs, noise-contrastive estimation, contrastive representation learning, and most modern mutual-information estimators.

5.1 Classifiers estimate density ratios

Label real samples $c=1$ and generated samples $c=0$, with equal class priors, and train a classifier $D_\psi(x) \approx \Pr(c=1\mid x)$ with the usual cross-entropy loss. Bayes' rule gives the optimum in closed form:

$$D^\star(x) = \frac{p(x)}{p(x) + q(x)} \qquad\Longrightarrow\qquad \log \frac{D^\star(x)}{1 - D^\star(x)} = \log\frac{p(x)}{q(x)} = \log r(x).$$

The logit of an optimal discriminator is the log density ratio. Once you have $r$, every $\phi$-divergence is available as a sample average, $D_\phi = \mathbb{E}_Q[\phi(r)]$. Substituting $D^\star$ into the GAN value function recovers $V(\theta,\psi^\star) = 2\,D_{\mathrm{JS}}(P\Vert Q_\theta) - \log 4$: the GAN was performing $\phi$-divergence minimization all along, with the classifier as a plug-in ratio estimator.

5.2 The variational representation and f-GAN

There is a cleaner and more general route. Since $\phi$ is convex, it equals the conjugate of its conjugate, $\phi(u) = \sup_t \{ut - \phi^*(t)\}$. Substituting into the definition and exchanging sup and integral gives a variational lower bound valid for any function $T$:

$$D_\phi(P \Vert Q) \;\ge\; \sup_{T}\;\Big\{ \mathbb{E}_{X\sim P}\big[T(X)\big] - \mathbb{E}_{Y\sim Q}\big[\phi^*(T(Y))\big]\Big\}, \qquad T^\star(x) = \phi'\!\big(r(x)\big).$$

Both terms are expectations, so both are estimable from samples. This is the f-GAN [9]: pick your favourite $\phi$, look up $\phi^*$, parameterize $T$ with a network, and you get a GAN variant that minimizes that specific divergence. Taking $\mathrm{KL}$ recovers the Donsker–Varadhan bound used by MINE [10] for mutual-information estimation.

Notice the shape shared with an IPM: both are a supremum over functions of a difference of two expectations. The difference is that an IPM constrains the function class and takes the raw difference, while a $\phi$-divergence leaves the class unconstrained and passes the second term through $\phi^*$. That extra nonlinearity is what makes the objective sensitive to ratios — and what makes it blow up when supports are disjoint.

5.3 Forward or reverse: the direction is a design decision

$\phi$-divergences are asymmetric, and the asymmetry has consequences you can see in samples.

Whenever you see a paper choose reverse KL (score distillation, MiniLLM, RLHF regularization) it is buying sharpness and accepting reduced diversity, and vice versa. It is usually the most consequential single decision in the method.

6. Estimation: where the real difficulty lives

All of §4 and §5 assumed population quantities. With $n$ samples, the picture changes, and three separate problems appear.

Sample complexity depends violently on the metric

For the empirical measure $\hat{P}_n$, the Wasserstein distance converges at $\mathbb{E}\,W_1(\hat P_n, P) = O(n^{-1/d})$ in $d$ dimensions [11]. At $d = 100$ that is hopeless: the empirical Wasserstein distance between two samples from the same distribution is large. MMD, in contrast, converges at $O(n^{-1/2})$ independent of dimension, because it is an average of a bounded kernel rather than a combinatorial matching. This is the strongest argument for kernel methods — and it comes with a caveat: dimension-free estimation rate does not mean dimension-free test power. With a poorly chosen kernel, MMD's ability to distinguish two distributions still degrades badly, it just does so with a small variance.

Minibatch estimates are biased, and so are their gradients

$\widehat{\mathrm{MMD}}^2$ has an unbiased estimator, but $\widehat{\mathrm{MMD}} = \sqrt{\cdot}$ does not, and neither do minibatch Wasserstein or Sinkhorn values (the minibatch OT plan is not a restriction of the full plan). In practice one optimizes $\mathrm{MMD}^2$ rather than $\mathrm{MMD}$ for this reason, and kernel methods need large batches: the cross term averages $nm$ kernel evaluations, and with small $n$ the estimator is too noisy to provide a useful gradient. This is the practical reason GMMN-style methods were considered expensive.

The critic is never optimal

Every method with a sup inside it — WGAN, f-GAN, MMD-GAN with a learned kernel — alternates between improving the critic and improving the generator, so the generator receives gradients from a stale, suboptimal critic. The theory (optimal $D^\star$, exact JS, exact $W_1$) holds only at the inner optimum, which is never reached. Combined with §4.4, this means the reported loss curve is not a distance and generally does not correlate with sample quality — a notorious frustration with GAN training, and something WGAN partly fixed by making the critic loss at least roughly proportional to a real distance.

The recurring trade-off. You may have a closed-form supremum (MMD), in which case you must choose a kernel by hand and you pay in large batches and weak test power in high dimension. Or you may learn the witness function (WGAN, f-GAN), in which case you get an adaptive, data-driven notion of difference and you pay with a minimax game, unknown approximation gap, and a loss you cannot trust. Every algorithm in §7 is a position on this axis, and MMD-GAN is the attempt to have both.

7. Algorithms

7.1 Generative moment matching networks

The simplest thing that works. GMMN [12] (and concurrently [13]) generates $y = g_\theta(u)$ from a uniform prior $u \sim \mathcal{U}[-1,1]^d$ and minimizes the plug-in MMD estimate directly:

$$\min_\theta\; \widehat{\mathrm{MMD}}^2\big(\{x_i\}_{i=1}^N,\; \{g_\theta(u_j)\}_{j=1}^M\big), \qquad k(x,y) = \exp\!\left(-\frac{\Vert x - y\Vert_2^2}{2\sigma^2}\right).$$

There is no discriminator and no minimax — it is a plain minimization, which is the whole appeal. Because a characteristic kernel's feature map contains all moments, this is moment matching to infinite order.

Algorithm 1 — GMMN training
Given: kernel k (usually a mixture of Gaussians with bandwidths sigma_1..sigma_S)
repeat:
  1. X <- minibatch of N real samples
  2. U <- N samples from Uniform[-1,1]^d
  3. Y <- g_theta(U)                                  # M = N generated samples
  4. Kxx, Kyy, Kxy <- pairwise kernel matrices
  5. mmd2 <- mean(Kxx_offdiag) + mean(Kyy_offdiag) - 2*mean(Kxy)
  6. theta <- Adam step on mmd2                       # note: minimization only
# batches must be large; the estimator is the bottleneck, not the network

The weaknesses are exactly the ones §3 and §6 predict. Performance is highly sensitive to $\sigma$ — the standard fix is a mixture of kernels over several bandwidths — large batches are mandatory, and a fixed kernel in pixel space measures a notion of similarity that has nothing to do with perceptual similarity. GMMN partly compensates by running MMD in the latent space of a pretrained autoencoder, which is already an admission that the kernel should be learned.

7.2 WGAN, and the fight to enforce Lipschitzness

WGAN [3] takes the dual from §4.1 and parameterizes the witness function with a network:

$$\min_\theta \; \sup_{\Vert f\Vert_L \le 1}\; \mathbb{E}_{x\sim P}\big[f(x)\big] - \mathbb{E}_{z\sim p(z)}\big[f(g_\theta(z))\big].$$

No kernel, no bandwidth: the critic learns which directions matter. Everything then hinges on imposing the Lipschitz constraint, and the history of that is the history of GAN stability work.

Algorithm 2 — WGAN with gradient penalty
repeat:
  # ---- critic: n_critic inner steps (typically 5)
  for i = 1..n_critic:
      x     <- real minibatch ;  z <- N(0, I)
      y     <- g_theta(z)
      eps   ~ Uniform[0,1] ;  xh <- eps*x + (1-eps)*y      # interpolate
      gp    <- ( ||grad_xh f_psi(xh)||_2 - 1 )^2
      L_psi <- -( mean f_psi(x) - mean f_psi(y) ) + lambda * mean(gp)
      psi   <- Adam step on L_psi

  # ---- generator: one step
  z       <- N(0, I)
  L_theta <- -mean f_psi(g_theta(z))                        # push critic score up
  theta   <- Adam step on L_theta

# mean f_psi(x) - mean f_psi(y) is an estimate of W_1; it should decrease
# and it (unlike the vanilla GAN loss) loosely tracks sample quality

The unresolved objection is capacity. Arora et al. [8] showed a network critic cannot represent the true Wasserstein supremum, so WGAN is optimizing a weaker surrogate than advertised. If a single learned function is too weak, an obvious response is to compare in a richer space — match higher-order moments through a kernel — which is the next section.

7.3 MMD-GAN: learn the kernel

MMD-GAN [17] is the synthesis. GMMN's problem is a fixed kernel in the wrong space; WGAN's problem is a single-function critic. So keep MMD's closed-form supremum but make the kernel adversarial:

$$\min_\theta\; \max_{\kappa \in \mathcal{K}}\; \mathrm{MMD}_{\kappa}\big(P,\,Q_\theta\big).$$

The construction rests on a small, elegant fact. If $k$ is characteristic and $f$ is injective, then the composition $\tilde k(x,x') = k\big(f(x), f(x')\big)$ is still characteristic. So we can learn an arbitrary feature map and compose it with a fixed Gaussian kernel, and the resulting object is still a valid metric:

$$\tilde k_w(x, x') = \exp\!\big(-\Vert f_w(x) - f_w(x')\Vert_2^2\big), \qquad \min_\theta \max_{w}\; \mathrm{MMD}_{k \circ f_w}(P, Q_\theta).$$

Why compose with a Gaussian rather than let the network output a similarity directly? Because an arbitrary learned $f_w(x,x')$ need not be positive definite, and then none of the RKHS theory — the closed form, the metric property, the mean embedding — applies. The composition gets a learned representation while keeping the guarantees for free.

Injectivity is enforced approximately by making $f_w$ the encoder of an autoencoder, $w = \{w_e, w_d\}$, with a reconstruction penalty:

$$\min_\theta \max_{w}\; \mathrm{MMD}_{k\circ f_{w_e}}(P, Q_\theta) \;-\; \lambda\, \mathbb{E}_{y \in X \cup g_\theta(Z)} \big\Vert y - f_{w_d}\big(f_{w_e}(y)\big)\big\Vert_2^2 .$$
Algorithm 3 — MMD-GAN
repeat:
  # ---- adversarial kernel: maximize MMD in the learned feature space
  for i = 1..n_critic:
      x <- real batch ; y <- g_theta(z)
      hx, hy <- f_we(x), f_we(y)                    # encoder features
      mmd2   <- MMD2_gaussian(hx, hy)
      rec    <- ||x - f_wd(hx)||^2 + ||y - f_wd(hy)||^2     # keep f_we injective
      L_w    <- -( mmd2 - lambda * rec )
      w      <- step on L_w      (then clip / spectral-normalize f_we)

  # ---- generator: minimize the same MMD
  y       <- g_theta(z)
  L_theta <- MMD2_gaussian(f_we(x), f_we(y))
  theta   <- step on L_theta

Empirically MMD-GAN beats GMMN decisively and is competitive with WGAN at smaller network size and shorter training time [14]. Two refinements matter. First, the same Lipschitz problem reappears — an unconstrained $f_w$ can inflate MMD without bound — and the fix that works is to regularize the features rather than the witness function: gradient-constrained MMD [18] and Sobolev GAN [19] both penalize the gradient of the critic/features, effectively "weakening" the MMD into something with better-behaved gradients. Second, this line gave the field a good evaluation metric: Kernel Inception Distance [14] is the unbiased $\mathrm{MMD}^2$ estimator with a polynomial kernel on Inception features, and unlike FID it is unbiased and has a usable variance estimate.

7.4 f-GAN and the divergence zoo

f-GAN [9] operationalizes §5.2: choose $\phi$, get an algorithm. The critic is written $T_\psi(x) = a_\phi(V_\psi(x))$ with an output activation $a_\phi$ that keeps $T$ in the domain of $\phi^*$.

Algorithm 4 — f-GAN
Given: convex phi with phi(1)=0; conjugate phi*; output activation a_phi
repeat:
  x <- real batch ; y <- g_theta(z)
  T_x, T_y <- a_phi(V_psi(x)), a_phi(V_psi(y))
  L_psi   <- -( mean(T_x) - mean(phi_star(T_y)) )   # maximize the variational bound
  psi     <- step on L_psi
  L_theta <-  ( mean(T_x) - mean(phi_star(T_y)) )   # minimize the divergence
  theta   <- step on L_theta

# phi = u log u        -> KL          a_phi = identity,     phi*(t) = exp(t-1)
# phi = -log u         -> reverse KL  a_phi = -exp(-v),     phi*(t) = -1-log(-t)
# phi = JS generator   -> vanilla GAN a_phi = -softplus(-v),phi*(t) = -log(2-e^t)
# phi = (u-1)^2        -> Pearson     a_phi = identity,     phi*(t) = t^2/4 + t

The honest finding from this line of work is that the choice of $\phi$ matters far less than the regularization of the critic. Least-squares GAN (Pearson $\chi^2$), hinge-loss GAN, and JS-GAN all behave similarly once spectral normalization is applied. The divergence determines the idealized objective; the constraint on the critic determines the actual optimization landscape, and in practice the latter dominates.

7.5 Entropic optimal transport: Sinkhorn divergences

A fifth position on the axis: keep optimal transport's geometry but make the primal problem solvable on a minibatch by adding an entropy term

$$\mathrm{OT}_\varepsilon(P,Q) = \min_{\gamma\in\Pi(P,Q)} \mathbb{E}_\gamma\big[c(x,y)\big] + \varepsilon\, \mathrm{KL}\big(\gamma \Vert P\otimes Q\big),$$

which turns the linear program into a strictly convex one solvable by a few dozen iterations of the Sinkhorn algorithm — alternating row and column rescalings of $K = e^{-C/\varepsilon}$, all differentiable and all on the GPU [20]. Since $\mathrm{OT}_\varepsilon(P,P) \ne 0$, one uses the debiased Sinkhorn divergence [21]

$$S_\varepsilon(P,Q) = \mathrm{OT}_\varepsilon(P,Q) - \tfrac12 \mathrm{OT}_\varepsilon(P,P) - \tfrac12\mathrm{OT}_\varepsilon(Q,Q),$$

which has the pleasing property of interpolating between the two families: as $\varepsilon \to 0$ it recovers Wasserstein, and as $\varepsilon\to\infty$ it converges to MMD with kernel $-c$. The regularization parameter is a dial between "sharp geometry, hard to estimate" and "smooth statistics, easy to estimate," which is a rather satisfying way to see the trade-off from §6. OT-GAN [22] combines it with a learned cost, since the choice of ground metric $c$ is the OT analogue of choosing a kernel.

8. Critic-free matching: scores and transport

Everything so far compares $P$ and $Q$ through function values. A different idea: compare them through the gradients of their log-densities. This buys a large advantage — $\nabla_x \log p(x)$ is invariant to the normalizing constant of $p$ — so it works when the target is known only up to a constant, which is exactly the situation in Bayesian inference, energy-based models, and reward-tilted policies.

8.1 Fisher divergence and score matching

$$D_F(P \Vert Q) = \mathbb{E}_{x\sim P}\big\Vert \nabla_x\log p(x) - \nabla_x \log q(x)\big\Vert^2 .$$

With integration by parts this becomes computable from samples of $P$ alone without knowing $\nabla\log p$, which is Hyvärinen's score matching [23]; its denoising variant is the objective underlying diffusion models. It is a genuine distribution matching criterion, just one that compares vector fields rather than expectations.

8.2 Stein discrepancies and SVGD

The kernelized Stein discrepancy is an IPM with a clever twist. Define the Stein operator $\mathcal{A}_p f(x) = f(x)\nabla_x\log p(x)^\top + \nabla_x f(x)$; under mild conditions $\mathbb{E}_{P}[\mathcal{A}_p f] = 0$ for all $f$ in a suitable class. So the quantity

$$\mathrm{KSD}(Q, P) = \sup_{\Vert f\Vert_{\mathcal{H}_k}\le 1} \mathbb{E}_{x\sim Q}\big[\mathcal{A}_p f(x)\big]$$

is zero iff $Q = P$, has a closed form like MMD, and — the important part — requires only samples from $Q$ and the score of $P$ [24]. Turning its gradient into a particle update gives Stein variational gradient descent [25], which moves a set of particles to match a target with an attraction term (toward high density) and a kernel repulsion term (enforcing diversity).

8.3 Diffusion and flow matching as distribution matching without a critic

It is worth stating plainly that modern generative models do distribution matching without ever computing a discrepancy. A diffusion model regresses the score of noisy data; a flow matching model regresses a conditional velocity field. Both objectives are plain least squares whose minimizer provably transports the prior onto the data distribution. There is no adversary, no kernel, no density ratio — and this is a large part of why they displaced GANs. I wrote about the details in The Evolution of Generative Models.

The relevance here is that the explicit-discrepancy machinery did not disappear; it moved to wherever a pretrained model needs to be matched by another model. That is §9.

9. Distribution matching beyond generative modeling

9.1 Domain adaptation and invariance

Given labelled source data and unlabelled target data, train a feature extractor $h_\phi$ so that $h_\phi(X_{\text{src}})$ and $h_\phi(X_{\text{tgt}})$ are indistinguishable; then a classifier trained on the source features transfers. This is distribution matching in feature space, and all three families appear:

The same template covers fairness (make representations independent of a protected attribute), batch-effect correction in single-cell biology, and sim-to-real transfer.

9.2 Diffusion distillation: matching a teacher's distribution

This is where distribution matching is most active today. A diffusion teacher needs 50 network evaluations per sample; we want a one-step student $g_\theta$. Regressing on the teacher's outputs pairwise gives blurry averages, so instead we match distributions: minimize the reverse KL between the student's output distribution and the teacher's.

The problem is that neither density is available — but their scores are, and the gradient of reverse KL only needs scores. For the noised student distribution $q_{\theta,t}$ and noised teacher $p_t$,

$$\nabla_\theta\, \mathrm{KL}\big(q_{\theta,t} \Vert p_t\big) = -\,\mathbb{E}\Big[\big(\underbrace{s_{p}(x_t,t)}_{\text{teacher score}} - \underbrace{s_{q_\theta}(x_t,t)}_{\text{student score}}\big)^{\!\top} \frac{\partial x_t}{\partial \theta}\Big].$$

The teacher score is the frozen pretrained model. The student score is unavailable, so a second "fake" diffusion model is trained online, by ordinary denoising score matching, on the student's own samples. Distribution Matching Distillation [29], Diff-Instruct [30], and Variational Score Distillation [31] are all this idea; the older Score Distillation Sampling of DreamFusion [32] is the special case where the student score is replaced by the noise itself, which is why SDS needs enormous guidance weights and produces oversaturated results.

Algorithm 5 — Distribution matching distillation (sketch)
Frozen: teacher score s_real(x_t, t) from the pretrained diffusion model
Learn:  one-step generator g_theta ; auxiliary "fake" score s_fake(x_t, t; xi)

repeat:
  # ---- fit the auxiliary score model to the student's current distribution
  z  <- N(0, I) ;  x0 <- g_theta(z).detach()
  t  <- Uniform ; eps <- N(0, I)
  xt <- alpha_t * x0 + sigma_t * eps
  xi <- step on || eps_xi(xt, t) - eps ||^2          # plain denoising score matching

  # ---- update the generator along the score difference
  z  <- N(0, I) ;  x0 <- g_theta(z)
  t  <- Uniform ; xt <- alpha_t * x0 + sigma_t * N(0, I)
  grad <- ( s_fake(xt, t) - s_real(xt, t) )          # = -grad of reverse KL wrt x
  theta <- step using backprop of  (x0 * stopgrad(grad)).sum()

# reverse KL is mode-seeking, so practical systems (DMD2) add a GAN loss
# and/or a small regression term to recover the diversity it drops

Notice the composition: a reverse-KL objective estimated through scores, stabilized by an adversarial term from §7, distilling a model trained by score matching from §8. Every family in this post appears in one loop.

9.3 Language model distillation: the direction of KL, again

Standard knowledge distillation minimizes forward KL between teacher and student token distributions, which is mass-covering: the student wastes capacity trying to reproduce the teacher's entire long tail, which it lacks the capacity to represent, and ends up producing samples the teacher would never generate. MiniLLM [33] switched to reverse KL for exactly this reason, so the student concentrates on the teacher's high-probability modes; because the objective is now an expectation under the student, it requires policy-gradient machinery. Generalized KD [34] makes both choices explicit — a divergence parameter (interpolated JSD) and a data source (teacher-generated vs student-generated sequences) — and shows the second matters at least as much, since training on the student's own outputs eliminates the train/inference distribution mismatch.

9.4 RLHF and DPO as constrained distribution matching

The RLHF objective is explicitly a distribution matching problem in disguise:

$$\max_\pi\; \mathbb{E}_{y\sim\pi(\cdot|x)}\big[r(x,y)\big] - \beta\, \mathrm{KL}\big(\pi \Vert \pi_{\text{ref}}\big) \qquad\Longrightarrow\qquad \pi^\star(y|x) \;\propto\; \pi_{\text{ref}}(y|x)\,\exp\!\Big(\tfrac{1}{\beta} r(x,y)\Big).$$

The optimum is a known, unnormalized target distribution — a reward-tilted reference model — and RL is just one way of matching it. DPO [35] inverts the relation to write $r$ in terms of $\pi^\star$ and $\pi_{\text{ref}}$, eliminating both the reward model and the sampling loop. The intractable normalizer cancels because the loss only involves differences of log-ratios between a preferred and dispreferred response, which is the same "ratios are available even when densities are not" move as §5.1.

9.5 Dataset distillation and evaluation metrics

Two smaller but instructive uses. Dataset distillation by distribution matching [36] synthesizes a tiny training set whose feature-mean embedding matches the real dataset's under randomly initialized encoders — MMD with a random-feature kernel, and dramatically cheaper than the gradient-matching alternatives.

And every metric used to evaluate generative models is a discrepancy from this post:

10. A practical guide to choosing

Situation Reasonable choice Why
Matching in a low-dimensional latent or feature space MMD with a mixture of Gaussian bandwidths, or Sinkhorn No adversary, no extra network, stable. Dimension is low enough that the kernel is meaningful.
Matching in pixel space from scratch Learned critic: WGAN-GP or hinge loss with spectral norm — or, better, don't; use diffusion / flow matching Fixed kernels do not encode perceptual similarity; the critic has to learn the right notion of difference.
Both supports are (nearly) disjoint An IPM — $W_1$, MMD, Sinkhorn $\phi$-divergences saturate and give no gradient (§3).
The target is known only up to a normalizer Reverse KL via scores, or KSD / SVGD Score-based criteria are invariant to the partition function (§8).
You want sharpness and can sacrifice diversity Reverse KL, or add an adversarial term Mode-seeking. Standard in distillation.
You must not drop modes Forward KL / maximum likelihood, or a symmetric divergence Mass-covering. Standard in pretraining.
Small batches are forced on you Avoid MMD and minibatch OT Their estimators need $O(n^2)$ pairs; the gradient is too noisy otherwise (§6).
You need a number you can trust for evaluation KID (unbiased MMD$^2$), plus precision/recall Never trust a learned critic's loss as a distance (§4.4).

Things I'd want a reader to leave with

The taxonomy is real and predictive. Differences (IPMs) stay finite and differentiable when supports are disjoint because the witness function is constrained to be smooth; ratios ($\phi$-divergences) blow up or saturate because they read the two distributions pointwise. Total variation is the only member of both families, and it behaves like a ratio-based divergence in practice.

MMD's real advantage is not statistical power, it is that the supremum is closed-form. No inner optimization means no minimax, no stale critic, no unknown approximation gap — you pay for it by having to choose a kernel and use large batches. MMD-GAN's contribution is recognizing that the kernel can be learned while keeping the closed form, by composing a learned injective feature map with a characteristic kernel.

Constraining the critic matters more than choosing the divergence. Once spectral normalization or a gradient penalty is in place, most GAN objectives behave similarly. The theory says the divergence determines the optimum; the practice says the constraint determines whether you get anywhere near it.

And the direction of the divergence is the decision to think hardest about. Forward versus reverse KL is the difference between blurry-but-complete and sharp-but-narrow, and it explains, with one idea, why likelihood models are diverse and hazy, why score distillation is sharp and mode-collapsed, why MiniLLM flipped the KL, and why practical distillation systems add a GAN loss to buy back what reverse KL threw away.

11. References