6. Strategy B.1: Compare Samples Directly¶
Chapter 5 preserved forward KL by changing the model until likelihood became usable. Strategy B preserves the opposite choice: keep a free-form generator, give up density evaluation, and build the training signal from samples.
This chapter studies the first way to spend that signal. We compute one scalar discrepancy between a generated batch and a data batch, then backpropagate it through the generated samples. Chapter 7 will keep the same sample access but turn such comparisons into explicit drifting fields.
We begin with three comparisons that are written directly from the two batches: kernel averages, sorted projections, and a soft transport plan. GAN comes last as the more advanced variation. It starts from a desired distribution divergence too, but learns the sample-based machinery that exposes that divergence to the generator.
1. The Strategy B.1 bargain¶
Draw independent source and data batches,
$$ \begin{aligned} z_i &\sim P_Z,\\ x_i &=G_\theta(z_i),\\ y_j &\sim P. \end{aligned} \tag{6.1} $$
and form the empirical laws
$$ \begin{aligned} \widehat Q_\theta &= \frac{1}{m}\sum_{i=1}^{m}\delta_{x_i},\\ \widehat P &= \frac{1}{n}\sum_{j=1}^{n}\delta_{y_j}. \end{aligned} \tag{6.2} $$
The training loss must be computable from these two point clouds and differentiable with respect to the generated points:
$$ \min_\theta\; \widehat D\!\left(\widehat Q_\theta,\widehat P\right). \tag{6.3} $$
We use the same unconstrained one-step MLP generator for all four methods. Only the comparison changes:
| Method | $D$ in its computable form | What reads the batches | The bill |
|---|---|---|---|
| MMD | kernel averages | fixed RBF kernels | bandwidth and high-dimensional sensitivity |
| Sliced Wasserstein | sorted one-dimensional projections | random directions | finite projections lose geometry |
| Sinkhorn divergence | debiased entropic transport cost | soft transport plan | quadratic batch memory and an $\varepsilon$ tradeoff |
| GAN | JS divergence through adversarial classification | learned discriminator | coupled min-max optimization |
The generator never exposes $q_\theta(x)$. It only needs to produce samples and support gradients through those samples. The first three rows prescribe the comparison before training begins. The GAN instead learns its comparison from the same real and generated batches.
2. The shared pretraining example¶
We use Chapter 4's two-dimensional Gaussian source and eight-mode data target. Alongside energy distance and mode coverage, we report stray mass: the fraction of generated samples farther than three target standard deviations from every mode center. It makes probability left between the modes visible as one number.
Energy distance is used only as the common scoreboard. None of the four methods is trained on it. The setup code is folded because the same source, target, and plotting conventions were introduced in Chapters 4 and 5.
from __future__ import annotations
import math
import numpy as np
import torch
import matplotlib.pyplot as plt
torch.manual_seed(0)
np.random.seed(0)
N_MODES = 8
RING_RADIUS = 4.0
MODE_STD = 0.30
EVAL_N = 3000
angles = torch.arange(N_MODES) * (2 * math.pi / N_MODES)
MODE_CENTERS = torch.stack(
[RING_RADIUS * torch.cos(angles), RING_RADIUS * torch.sin(angles)],
dim=1,
)
MODE_COLORS = [plt.get_cmap("tab10")(k) for k in range(N_MODES)]
def sample_ring(n, weights=None):
'''Draw samples from the eight-component Gaussian ring.'''
if weights is None:
component = torch.randint(0, N_MODES, (n,))
else:
component = torch.multinomial(weights, n, replacement=True)
return MODE_CENTERS[component] + MODE_STD * torch.randn(n, 2)
def sample_source(n):
'''Draw samples from the two-dimensional standard Gaussian source.'''
return torch.randn(n, 2)
def assign_modes(x):
'''Return the nearest ring-mode index for each sample.'''
return torch.cdist(x, MODE_CENTERS).argmin(1)
def plot_samples(ax, x, title="", color="#1d4ed8", by_mode=False):
'''Plot a two-dimensional sample cloud on the shared canvas.'''
x = x.detach().cpu()
colors = (
[MODE_COLORS[k] for k in assign_modes(x).tolist()]
if by_mode
else color
)
ax.scatter(
x[:, 0],
x[:, 1],
s=5,
alpha=0.35,
c=colors,
edgecolors="none",
)
ax.set_xlim(-6.5, 6.5)
ax.set_ylim(-6.5, 6.5)
ax.set_aspect("equal")
ax.set_title(title, fontsize=10)
ax.set_xticks([])
ax.set_yticks([])
DATA = sample_ring(EVAL_N)
import torch.nn as nn
from tqdm import tqdm
def mode_histogram(x):
'''Return the sample fraction assigned to each ring mode.'''
return torch.bincount(assign_modes(x), minlength=N_MODES).float() / len(x)
def modes_hit(x, threshold=0.02):
'''Count modes containing more than threshold of the sample mass.'''
return int((mode_histogram(x) > threshold).sum())
def energy_distance(x, y):
'''Compute the empirical energy distance between two sample batches.'''
mean_distance = lambda a, b: torch.cdist(a, b).mean()
return (
2 * mean_distance(x, y)
- mean_distance(x, x)
- mean_distance(y, y)
).item()
def stray_fraction(x, threshold=3 * MODE_STD):
'''Return the mass farther than threshold from every target mode.'''
nearest = torch.cdist(x, MODE_CENTERS).min(dim=1).values
return nearest.gt(threshold).float().mean().item()
def mlp(inp, out, hidden=128, depth=3, activation=nn.SiLU):
'''Build a small multilayer perceptron.'''
layers = [nn.Linear(inp, hidden), activation()]
for _ in range(depth - 1):
layers.extend([nn.Linear(hidden, hidden), activation()])
return nn.Sequential(*layers, nn.Linear(hidden, out))
class Generator(nn.Module):
'''A free-form one-step map from Gaussian noise to the sample space.'''
def __init__(self, latent_dim=2, hidden=128):
super().__init__()
self.latent_dim = latent_dim
self.network = mlp(latent_dim, 2, hidden)
def forward(self, z):
return self.network(z)
def sample(self, n):
return self(torch.randn(n, self.latent_dim))
def mmd2(x, y, sigmas=(0.5, 1.0, 2.0, 4.0, 8.0)):
'''Biased multi-bandwidth RBF maximum mean discrepancy squared.'''
def kernel(d2):
return sum(torch.exp(-d2 / (2 * sigma**2)) for sigma in sigmas)
return (
kernel(torch.cdist(x, x).square()).mean()
+ kernel(torch.cdist(y, y).square()).mean()
- 2 * kernel(torch.cdist(x, y).square()).mean()
)
def sliced_w2_squared(x, y, n_projections=128):
'''Average squared 1D Wasserstein distance over random projections.'''
directions = torch.randn(
x.shape[1],
n_projections,
dtype=x.dtype,
device=x.device,
)
directions = directions / directions.norm(dim=0, keepdim=True)
x_projected = (x @ directions).sort(dim=0).values
y_projected = (y @ directions).sort(dim=0).values
return (x_projected - y_projected).square().mean()
def sinkhorn_divergence(x, y, eps=0.1, n_iter=10):
'''Debiased entropic OT with Euclidean cost and log-domain scaling.'''
def regularized_ot(a, b):
cost = torch.cdist(a, b)
n, m = cost.shape
f = a.new_zeros(n)
g = b.new_zeros(m)
for _ in range(n_iter):
f = -eps * torch.logsumexp(
(g[None, :] - cost) / eps - math.log(m),
dim=1,
)
g = -eps * torch.logsumexp(
(f[:, None] - cost) / eps - math.log(n),
dim=0,
)
return f.mean() + g.mean()
return (
regularized_ot(x, y)
- 0.5 * regularized_ot(x, x)
- 0.5 * regularized_ot(y, y)
)
SAMPLE_RUNS = {}
def record_run(key, samples):
'''Measure and retain one sample-matching model.'''
result = {
"samples": samples.detach(),
"energy": energy_distance(samples, DATA),
"modes": modes_hit(samples),
"stray": stray_fraction(samples),
}
SAMPLE_RUNS[key] = result
print(
f"{key:20s} ED={result['energy']:.4f} "
f"modes={result['modes']}/8 stray={100 * result['stray']:.1f}%"
)
def show_run(key, label):
'''Plot one trained generator beside the shared target.'''
result = SAMPLE_RUNS[key]
fig, ax = plt.subplots(1, 2, figsize=(7.4, 3.6))
title = (
f"{label}\nED={result['energy']:.3f} "
f"modes={result['modes']}/8 stray={100 * result['stray']:.1f}%"
)
plot_samples(ax[0], result["samples"], title, by_mode=True)
plot_samples(ax[1], DATA, "target", color="#b45309", by_mode=True)
plt.tight_layout()
plt.show()
fig, ax = plt.subplots(1, 2, figsize=(7.4, 3.6))
plot_samples(ax[0], sample_source(EVAL_N), r"source $P_Z=\mathcal{N}(0,I)$", color="#64748b")
plot_samples(ax[1], DATA, "target $P$ (samples only)", color="#b45309", by_mode=True)
plt.tight_layout()
plt.show()
3. MMD: fix the comparison with a kernel¶
Maximum mean discrepancy chooses a positive-definite kernel $k$ to compare all within-model, within-data, and cross-batch pairs:
$$ \begin{aligned} \mathrm{MMD}^2(Q_\theta,P) &= \mathbb E[k(x,x')] + \mathbb E[k(y,y')] - 2\mathbb E[k(x,y)],\\ k(x,y) &= \sum_{\sigma\in\mathcal S} \exp\!\left( -\frac{\lVert x-y\rVert^2}{2\sigma^2} \right),\\ \mathcal S &= \{0.5,1,2,4,8\}. \end{aligned} \tag{6.4} $$
This is the exact multi-bandwidth radial-basis-function kernel used below. Small bandwidths inspect local clusters; large bandwidths inspect the global arrangement. The empirical loss replaces the three expectations by averages over the two batches. The resulting generative moment matching network has one fixed differentiable loss that can be optimized directly through the generated samples.
torch.manual_seed(42)
mmd_generator = Generator()
optimizer = torch.optim.Adam(mmd_generator.parameters(), lr=1e-3)
for _ in tqdm(range(4000), desc="MMD"):
loss = mmd2(mmd_generator.sample(256), sample_ring(256))
optimizer.zero_grad()
loss.backward()
optimizer.step()
record_run("MMD", mmd_generator.sample(EVAL_N))
show_run("MMD", "MMD")
MMD: 100%|██████████| 4000/4000 [00:24<00:00, 163.81it/s]
MMD ED=0.0096 modes=8/8 stray=18.0%
Reading the trained MMD generator¶
The MMD generator reaches all eight modes with energy distance $0.0096$. Its stray mass is $18.0\%$: broad arcs and a few radial streaks remain between the compact clusters. The energy-distance score therefore looks stronger than the geometry.
The fixed comparison gives the generator a single optimization objective. Its limitation moves into the kernel: a bandwidth useful at one geometric scale may miss another, and pairwise distances become less discriminative in high-dimensional spaces. Five bandwidths soften that choice without removing it, and the kernel average still charges a thin bridge only in proportion to its mass.
4. Sliced Wasserstein: compare sorted projections¶
Optimal transport is especially simple in one dimension: sort two equally sized batches and pair their order statistics. Sliced Wasserstein projects the samples onto random unit directions $\omega$, solves those one-dimensional problems, and averages:
$$ \mathrm{SW}_2^2(Q_\theta,P) = \mathbb E_{\omega} \left[ W_2^2\!\left( \omega_\#Q_\theta,\, \omega_\#P \right) \right]. \tag{6.5} $$
Every projected loss is differentiable through the generated coordinates. The computation replaces a multidimensional transport problem by matrix multiplication and sorting.
torch.manual_seed(43)
sliced_generator = Generator()
optimizer = torch.optim.Adam(sliced_generator.parameters(), lr=1e-3)
for _ in tqdm(range(4000), desc="sliced Wasserstein"):
loss = sliced_w2_squared(
sliced_generator.sample(512), sample_ring(512), n_projections=128
)
optimizer.zero_grad()
loss.backward()
optimizer.step()
record_run("sliced Wasserstein", sliced_generator.sample(EVAL_N))
show_run("sliced Wasserstein", "sliced Wasserstein")
sliced Wasserstein: 100%|██████████| 4000/4000 [00:15<00:00, 258.98it/s]
sliced Wasserstein ED=0.0152 modes=8/8 stray=30.5%
Reading the sliced-Wasserstein generator¶
All eight destinations remain detectable, but the generator is much closer to a continuous annulus: energy distance is $0.0152$ and stray mass reaches $30.5\%$. Many samples also sit inside the ring. A projection can hide this angular structure because modes that are separate in the plane overlap on many one-dimensional shadows.
More projections reduce that blind spot, but every extra projection adds sorting work and still does not reconstruct a full-dimensional transport plan. The method trades a cheap geometric comparison in one dimension for projection variance and partial views.
5. Sinkhorn divergence: soften full-dimensional transport¶
Entropic optimal transport retains the full pairwise geometry but regularizes the transport plan:
$$ \mathrm{OT}_\varepsilon(Q,P) = \min_{\pi\in\Pi(Q,P)} \mathbb E_{(x,y)\sim\pi}[\lVert x-y\rVert] + \varepsilon\,\mathrm{KL}(\pi\,\|\,Q\otimes P). \tag{6.6} $$
Entropy makes the plan computable by repeated matrix rescaling. Subtracting the two self-costs removes the entropic bias:
$$ \mathrm{S}_\varepsilon(Q,P) = \mathrm{OT}_\varepsilon(Q,P) - \tfrac12\mathrm{OT}_\varepsilon(Q,Q) - \tfrac12\mathrm{OT}_\varepsilon(P,P). \tag{6.7} $$
The regularization $\varepsilon$ is a resolution dial. Smaller values approach sharper transport but are harder to solve; larger values produce a smoother and more biased comparison. A batch of size $n$ also requires an $n\times n$ cost matrix.
torch.manual_seed(44)
sinkhorn_generator = Generator()
optimizer = torch.optim.Adam(sinkhorn_generator.parameters(), lr=1e-3)
for _ in tqdm(range(1500), desc="Sinkhorn divergence"):
loss = sinkhorn_divergence(
sinkhorn_generator.sample(256),
sample_ring(256),
eps=0.1,
n_iter=10,
)
optimizer.zero_grad()
loss.backward()
optimizer.step()
record_run("Sinkhorn", sinkhorn_generator.sample(EVAL_N))
show_run("Sinkhorn", "Sinkhorn divergence")
Sinkhorn divergence: 100%|██████████| 1500/1500 [00:44<00:00, 33.61it/s]
Sinkhorn ED=0.0070 modes=8/8 stray=12.0%
Reading the Sinkhorn generator¶
Sinkhorn gives the best energy distance, $0.0070$, while covering all eight modes. Its clusters are compact, but spoke-like connections leave $12.0\%$ stray mass. That is better than MMD and sliced Wasserstein on this structural metric, although the visible connections show that the endpoint match is still imperfect.
Full-dimensional transport gives more precise geometry than projections, but changing the loss does not remove the continuity constraint of a connected source passed through a continuous generator. The other bill is computational: dense pairwise costs limit batch size, minibatches see only part of the target, and finite $\varepsilon$ smooths the transport problem we actually descend.
6. GAN: learn the comparison¶
The first three methods gave us an explicit formula to evaluate on every pair of batches. A generative adversarial network starts from the same goal—minimize a distribution discrepancy—but handles a discrepancy whose density formula is unavailable.
Let $M_\theta=\tfrac12(P+Q_\theta)$. The Jensen--Shannon divergence is
$$ \begin{aligned} \mathrm{JS}(P\,\|\,Q_\theta) &= \tfrac12\mathrm{KL}(P\,\|\,M_\theta) + \tfrac12\mathrm{KL}(Q_\theta\,\|\,M_\theta),\\ \min_\theta\;&\mathrm{JS}(P\,\|\,Q_\theta). \end{aligned} \tag{6.8} $$
This is a valid measure of the gap between the real and generated distributions, but evaluating its density-ratio terms would require the unavailable $p(x)$ and $q_\theta(x)$. GAN makes the comparison sample-based by learning that ratio. A discriminator $D_\phi(x)\in(0,1)$ receives labeled real and generated samples. The generator and discriminator solve the min--max game
$$ \begin{aligned} \min_\theta\max_\phi\;&V(\phi,\theta),\\ V(\phi,\theta) &= \mathbb E_{y\sim P}[\log D_\phi(y)] + \mathbb E_{z\sim P_Z} \left[\log\!\left(1-D_\phi(G_\theta(z))\right)\right]. \end{aligned} \tag{6.9} $$
For a fixed generator, the best discriminator and the resulting value are
$$ \begin{aligned} D^\star(x) &= \frac{p(x)}{p(x)+q_\theta(x)},\\ \max_\phi V(\phi,\theta) &= 2\,\mathrm{JS}(P\,\|\,Q_\theta)-\log 4. \end{aligned} \tag{6.10} $$
The classifier therefore learns the density-ratio witness from the two sample streams. Minimizing its optimized value with respect to the generator minimizes the desired distribution discrepancy without ever evaluating either density.
In practice we use the non-saturating generator loss $-\mathbb E_z[\log D_\phi(G_\theta(z))]$. It has the same desired equilibrium but provides stronger gradients when the discriminator initially rejects every generated sample. The price is coupled optimization: whenever the generator changes, the learned comparison becomes stale. Our small implementation uses decaying instance noise and one batch-standard-deviation feature to keep the discriminator from winning immediately.
def append_batch_std(x):
'''Append one feature that reports diversity across the current batch.'''
batch_std = x.std(dim=0, unbiased=False).mean().expand(len(x), 1)
return torch.cat([x, batch_std], dim=1)
torch.manual_seed(41)
gan = Generator()
discriminator = mlp(3, 1, hidden=128, depth=2, activation=lambda: nn.LeakyReLU(0.2))
generator_optimizer = torch.optim.Adam(gan.parameters(), lr=2e-4, betas=(0.5, 0.9))
discriminator_optimizer = torch.optim.Adam(
discriminator.parameters(), lr=2e-4, betas=(0.5, 0.9)
)
binary_cross_entropy = nn.BCEWithLogitsLoss()
for step in tqdm(range(6000), desc="GAN"):
noise_scale = max(0.0, 0.3 * (1 - step / 4200))
real = sample_ring(512) + noise_scale * torch.randn(512, 2)
fake = gan.sample(512).detach() + noise_scale * torch.randn(512, 2)
discriminator_optimizer.zero_grad()
discriminator_loss = (
binary_cross_entropy(
discriminator(append_batch_std(real)), torch.ones(512, 1)
)
+ binary_cross_entropy(
discriminator(append_batch_std(fake)), torch.zeros(512, 1)
)
)
discriminator_loss.backward()
discriminator_optimizer.step()
generator_optimizer.zero_grad()
fake = gan.sample(512) + noise_scale * torch.randn(512, 2)
generator_loss = binary_cross_entropy(
discriminator(append_batch_std(fake)), torch.ones(512, 1)
)
generator_loss.backward()
generator_optimizer.step()
record_run("GAN", gan.sample(EVAL_N))
show_run("GAN", "GAN")
GAN: 100%|██████████| 6000/6000 [00:22<00:00, 266.28it/s]
GAN ED=0.0098 modes=8/8 stray=7.1%
Reading the trained GAN¶
The GAN reaches all eight modes with energy distance $0.0098$ and $7.1\%$ stray mass. The clusters are compact, but several are stretched and thin threads still join neighboring modes. The low stray mass fits the density-ratio view: a discriminator can react sharply when generated mass appears in a gap where the target places almost none.
The price is a moving comparison, which can make training unstable. An overly strong discriminator supplies little useful gradient, while a generator that concentrates on a few easy regions may ignore the rest—a failure called mode collapse. Alternating updates, discriminator capacity, learning-rate balance, and random seed therefore matter.
7. Same generator class, four sample comparisons¶
The final figure holds the source, target, and generator architecture fixed, so no method receives a more expressive model. Their stable learning rates, update counts, batch sizes, and per-update costs differ; this is not an equal-compute benchmark. Differences in coverage, sharpness, and mass between modes come from the comparison and the optimization it requires rather than from model capacity.
keys = ["MMD", "sliced Wasserstein", "Sinkhorn", "GAN"]
fig, axes = plt.subplots(1, 5, figsize=(13.6, 3.0))
for ax, key in zip(axes, keys):
result = SAMPLE_RUNS[key]
title = (
f"{key}\nED={result['energy']:.3f} "
f"stray={100 * result['stray']:.1f}%"
)
plot_samples(ax, result["samples"], title, by_mode=True)
plot_samples(
axes[4],
DATA,
f"target\nstray={100 * stray_fraction(DATA):.1f}%",
color="#b45309",
by_mode=True,
)
plt.suptitle("Strategy B.1: sample comparisons used as endpoint losses", fontsize=11)
plt.tight_layout()
plt.show()
print(f"{'method':20s} {'ED':>8s} {'modes':>7s} {'stray':>8s}")
for key in keys:
result = SAMPLE_RUNS[key]
print(
f"{key:20s} {result['energy']:8.4f} "
f"{result['modes']:>5d}/8 {100 * result['stray']:7.1f}%"
)
print(f"{'target':20s} {'':>8s} {'8/8':>7s} {100 * stray_fraction(DATA):7.1f}%")
method ED modes stray MMD 0.0096 8/8 18.0% sliced Wasserstein 0.0152 8/8 30.5% Sinkhorn 0.0070 8/8 12.0% GAN 0.0098 8/8 7.1% target 8/8 1.3%
The sampled target leaves $1.3\%$ outside the three-standard-deviation radius because each Gaussian mode has nonzero width. Against that baseline, all four methods find every destination, yet their stray mass ranges from $7.1\%$ to $30.5\%$. Energy distance alone does not expose that spread: GAN and MMD score almost identically while placing very different amounts of probability between modes.
All four methods train the same free-form generator class from the same two sample streams. MMD fixes the comparison with kernels, sliced Wasserstein reads sorted projections, and Sinkhorn computes a soft transport plan. GAN is the advanced variation: it begins with a desired divergence, then learns the density-ratio witness that makes that comparison available from samples. Each route removes the unavailable model density in a different way and pays for different comparison machinery.
This chapter spent each comparison as a scalar loss. Strategy B.2 keeps the same endpoint and sample access, opens that scalar to read the direction assigned to each generated particle, and uses those directions to build the recent popular drifting methods.