7. Strategy B.2: Turn Sample Divergences into Drifting Fields¶
Chapter 6 trained a free-form generator by computing one scalar discrepancy between generated and data batches, then backpropagating that number. But the backward pass does not jump directly from the number to the parameters. It differentiates the discrepancy with respect to every generated sample, maps each particle gradient back to the parameters through that sample's generator Jacobian, and sums the contributions across the batch.
Strategy B.2 keeps the same model law, target law, sample access, and distribution comparison. It takes the intermediate particle gradients before that final chain-rule step, points them downhill, and freezes short moves along those directions as regression targets. Differentiating the regression loss then passes through the same generator Jacobians and recovers the same parameter update. This chapter derives that complete reformulation first, then reads three representative drifting methods from familiar Strategy B comparisons:
- the exact field hidden inside Chapter 6's MMD loss,
- the energy-distance field and its one-triplet Three-Body Scattering approximation,
- and the barycentric field hidden inside a Sinkhorn divergence.
1. The Strategy B.2 bargain¶
Given source draws $z_i$ and data particles $y_j$, the generated particle is
$$ x_i = G_\theta(z_i). \tag{7.1} $$
The scalar sample comparison is
$$ \widehat D_\theta = \widehat D(x_1,\ldots,x_m;\,y_1,\ldots,y_n). \tag{7.2} $$
Strategy B.1 optimizes it directly:
$$ \min_\theta\;\widehat D_\theta. \tag{7.3} $$
Before the chain rule can update $\theta$, it differentiates that scalar with respect to every generated position. Point each result downhill:
$$ v_i = -\nabla_{x_i}\widehat D_\theta. \tag{7.4} $$
For a temporary positive scale $\eta$, freeze a short move along each arrow:
$$ \widetilde x_i = \operatorname{sg}\!\left[x_i+\eta v_i\right]. \tag{7.5} $$
Here $\operatorname{sg}$ means stop gradient: the target is treated as a constant when we update the generator. The temporary scale $\eta$ lets us check whether the target displacement introduces a new tuning choice. Write the Jacobian of particle $i$ with respect to the generator parameters as
$$ J_i = \frac{\partial x_i}{\partial\theta}. \tag{7.6} $$
The chain rule for the B.1 scalar loss is
$$ \begin{aligned} \nabla_\theta\widehat D_\theta &= \sum_i J_i^\mathsf T \nabla_{x_i}\widehat D_\theta\\ &= -\sum_i J_i^\mathsf T v_i. \end{aligned} \tag{7.7} $$
B.2 instead regresses onto the frozen targets:
$$ \mathcal L_{\mathrm{move}} = \frac{1}{2\eta} \sum_i \left\lVert x_i-\widetilde x_i\right\rVert^2. \tag{7.8} $$
Because stop gradient makes every $\widetilde x_i$ constant during this update, its chain-rule derivative is
$$ \begin{aligned} \nabla_\theta\mathcal L_{\mathrm{move}} &= \frac{1}{\eta} \sum_i J_i^\mathsf T \left(x_i-\widetilde x_i\right)\\ &= \frac{1}{\eta} \sum_i J_i^\mathsf T \left( x_i-\operatorname{sg}[x_i+\eta v_i] \right)\\ &= \frac{1}{\eta} \sum_i J_i^\mathsf T \left(-\eta v_i\right)\\ &= -\sum_i J_i^\mathsf T v_i\\ &= \nabla_\theta\widehat D_\theta. \end{aligned} \tag{7.9} $$
The last line is exactly the B.1 chain rule in (7.7): $\eta$ has cancelled. We therefore set $\eta=1$ from here on, matching the implementation below. The same unit-target convention appears in Sinkhorn Drifting, TBSM, and Teacher-Feature Drifting.
The detached moved-target objective in (7.8) is the training device at the center of Generative Modeling via Drifting. That paper also designs fields directly from an equilibrium condition. We stay narrower first: every implemented field below is derived from an explicit Strategy B sample discrepancy, so B.2 remains visibly a reformulation of B.1 rather than a new objective.
The form $x-\operatorname{sg}[\widehat x]$, with $\widehat x$ constructed from the current model, also recurs beyond pretraining. In RL post-training, a reward-derived improvement direction replaces the sample-discrepancy field, but the resulting target is again frozen before the model is updated.
2. The shared pretraining example¶
We keep the same Gaussian source, eight-mode target, one-step MLP generator, and three diagnostics used in Chapter 6. Energy distance is again the common scoreboard, mode coverage checks whether every destination is found, and stray mass measures probability left between the modes.
The setup is folded because Chapters 4–6 already introduced it. All drifting code remains visible.
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)
)
DRIFT_RUNS = {}
def record_drift(key, samples):
'''Measure and retain one drifting model.'''
result = {
"samples": samples.detach(),
"energy": energy_distance(samples, DATA),
"modes": modes_hit(samples),
"stray": stray_fraction(samples),
}
DRIFT_RUNS[key] = result
print(
f"{key:20s} ED={result['energy']:.4f} "
f"modes={result['modes']}/8 stray={100 * result['stray']:.1f}%"
)
def show_drift(key, label):
'''Plot one drifting result beside the shared target.'''
result = DRIFT_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: expose the field already inside the loss¶
Start with the exact multi-bandwidth MMD used in Chapter 6:
$$ \begin{aligned} \mathrm{MMD}^2(Q_\theta,P) &= \mathbb E_{x,x'\sim Q_\theta}[k(x,x')] + \mathbb E_{y,y'\sim P}[k(y,y')] - 2\mathbb E_{\substack{x\sim Q_\theta\\y\sim P}}[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{7.10} $$
For one RBF scale,
$$ \nabla_x k(x,y) = k(x,y)\frac{y-x}{\sigma^2}. $$
Here $x,x'$ are independent model draws, $y,y'$ are independent data draws, and the two sides are independent of each other. To evaluate the field, draw one $x\sim Q_\theta$ and hold that particle fixed. Differentiate the cross term and the model self-term, then average over its independent interaction partners $y$ and $x'$. Up to one positive batch factor, the particle field is
$$ \begin{aligned} v_{\mathrm{MMD}}(x) \propto \sum_{\sigma\in\mathcal S} \Bigg[ & \underbrace{ \mathbb E_{y\sim P} \left[ k_\sigma(x,y)\frac{y-x}{\sigma^2} \right] }_{\text{attraction to data}}\\ & \underbrace{ - \mathbb E_{x'\sim Q_\theta} \left[ k_\sigma(x,x')\frac{x'-x}{\sigma^2} \right] }_{\text{repulsion from the generated crowd}} \Bigg]. \end{aligned} \tag{7.11} $$
Attraction pulls generated samples toward real data, while repulsion pushes them apart to discourage mode collapse; this fidelity-coverage balance is the core intuition of drifting methods.
The diagnostic below evaluates the field across the plane rather than only at the initially clumped generator outputs. Its normalized arrows isolate attraction, repulsion, and their sum.
Nothing new has been added to MMD. We have only opened its gradient before the chain rule closes over the generator.
def rbf_pull(x, y, sigmas=(0.5, 1.0, 2.0, 4.0, 8.0)):
'''Kernel-weighted mean vectors pointing from x toward y.'''
delta = y[None, :, :] - x[:, None, :]
squared_distance = delta.square().sum(dim=-1)
pull = torch.zeros_like(x)
for sigma in sigmas:
weight = torch.exp(-squared_distance / (2 * sigma**2)) / sigma**2
pull += (weight[:, :, None] * delta).mean(dim=1)
return pull
def mmd_field_parts(x, model_samples, data_samples):
'''Return the data-attraction and model-repulsion parts at query points x.'''
attraction = rbf_pull(x, data_samples)
repulsion = -rbf_pull(x, model_samples)
return attraction, repulsion
def mmd_field(x, y):
'''Rescaled negative particle gradient of the multi-bandwidth MMD.'''
attraction, repulsion = mmd_field_parts(x, x, y)
return attraction + repulsion
def train_moved_field(
field,
steps,
batch_size,
learning_rate,
seed,
description,
cosine_decay=False,
):
'''Train G by regressing onto detached x + field(x).'''
torch.manual_seed(seed)
generator = Generator()
optimizer = torch.optim.Adam(generator.parameters(), lr=learning_rate)
scheduler = (
torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, steps)
if cosine_decay
else None
)
for _ in tqdm(range(steps), desc=description):
x = generator.sample(batch_size)
with torch.no_grad():
detached_x = x.detach()
target = detached_x + field(detached_x, sample_ring(batch_size))
loss = (x - target).square().sum(dim=1).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if scheduler is not None:
scheduler.step()
return generator
# Check both equalities before training.
torch.manual_seed(7)
probe = Generator()
probe_x = probe(sample_source(64))
probe_y = sample_ring(64)
probe_mmd = mmd2(probe_x, probe_y)
particle_gradient = torch.autograd.grad(probe_mmd, probe_x, retain_graph=True)[0]
analytic_field = mmd_field(probe_x.detach(), probe_y)
field_error = (
analytic_field + 0.5 * len(probe_x) * particle_gradient.detach()
).abs().max()
parameters = tuple(probe.parameters())
direct_gradient = torch.autograd.grad(probe_mmd, parameters, retain_graph=True)
exact_target = probe_x.detach() - particle_gradient.detach()
moved_loss = 0.5 * (probe_x - exact_target).square().sum()
moved_gradient = torch.autograd.grad(moved_loss, parameters)
gradient_error = torch.sqrt(
sum((a - b).square().sum() for a, b in zip(direct_gradient, moved_gradient))
/ sum(a.square().sum() for a in direct_gradient)
)
print(f"analytic field vs rescaled -grad MMD: max abs diff {field_error:.1e}")
print(f"scalar loss vs moved-target parameter gradient: relative diff {gradient_error:.1e}")
with torch.no_grad():
visual_model = probe(sample_source(512))
visual_data = sample_ring(2048)
axis = torch.linspace(-5.5, 5.5, 13)
grid_x, grid_y = torch.meshgrid(axis, axis, indexing="xy")
grid = torch.stack([grid_x.flatten(), grid_y.flatten()], dim=1)
attraction, repulsion = mmd_field_parts(grid, visual_model, visual_data)
total = attraction + repulsion
def draw_field(ax, field, title, color, show_data=False, show_model=False):
'''Plot field directions over the common sample-space grid.'''
if show_data:
ax.scatter(
visual_data[:, 0], visual_data[:, 1],
s=4, color="#d97706", alpha=0.10, edgecolors="none",
)
ax.scatter(
MODE_CENTERS[:, 0], MODE_CENTERS[:, 1],
s=28, color="#b45309", marker="x", linewidths=1.5,
)
if show_model:
ax.scatter(
visual_model[:, 0], visual_model[:, 1],
s=6, color="#475569", alpha=0.22, edgecolors="none",
)
direction = field / field.norm(dim=1, keepdim=True).clamp_min(1e-8)
ax.quiver(
grid[:, 0], grid[:, 1],
direction[:, 0], direction[:, 1],
color=color, angles="xy", scale_units="xy", scale=2.4, width=0.004,
)
ax.set(xlim=(-6.2, 6.2), ylim=(-6.2, 6.2), aspect="equal", title=title)
ax.set_xticks([])
ax.set_yticks([])
fig, ax = plt.subplots(1, 3, figsize=(11.2, 3.6))
draw_field(ax[0], attraction, "attraction\ntoward real data", "#d97706", show_data=True)
draw_field(ax[1], repulsion, "repulsion\naway from model crowd", "#dc2626", show_model=True)
draw_field(
ax[2], total, "total MMD field\nattraction + repulsion", "#2563eb",
show_data=True, show_model=True,
)
plt.tight_layout()
plt.show()
analytic field vs rescaled -grad MMD: max abs diff 1.3e-07 scalar loss vs moved-target parameter gradient: relative diff 1.5e-07
mmd_drifting = train_moved_field(
mmd_field,
steps=4000,
batch_size=256,
learning_rate=1e-3,
seed=42,
description="MMD field",
)
record_drift("MMD field", mmd_drifting.sample(EVAL_N))
show_drift("MMD field", "MMD field regression")
MMD field: 100%|██████████| 4000/4000 [00:15<00:00, 253.86it/s]
MMD field ED=0.0096 modes=8/8 stray=18.0%
Reading the MMD field¶
The two numerical checks settle the formulation before the training run: the analytical arrows equal the rescaled negative particle gradient, and frozen-target regression gives the same parameter gradient as direct scalar MMD descent.
The trained field reaches all eight modes with energy distance $0.0096$ and $18.0\%$ stray mass. Those are exactly the metrics printed by Chapter 6's scalar-MMD run under the same seed and minibatch sequence. The gradients differ only by the fixed positive scale shown in the analytical check, so Adam follows the same update path. Its remaining arcs are therefore not a new B.2 failure mode; they are the same kernel bill seen in B.1. Opening the scalar loss into arrows changed how we wrote the update, not what geometric differences MMD emphasizes.
4. Energy distance: remove the bandwidth¶
The RBF field decides which neighbors matter through $\sigma$. Energy distance replaces that bandwidth choice with Euclidean distance:
$$ \begin{aligned} \mathcal E^2(Q_\theta,P) &= 2\mathbb E\lVert x-y\rVert\\ &\quad- \mathbb E\lVert x-x'\rVert\\ &\quad- \mathbb E\lVert y-y'\rVert. \end{aligned} \tag{7.12} $$
It has the same three-term algebra as an MMD with the generalized distance kernel $k(x,y)=-\lVert x-y\rVert$. Euclidean distance is a negative-type metric; equivalently, it induces an associated positive-definite distance kernel.
Because
$$ \nabla_x\lVert x-y\rVert = \frac{x-y}{\lVert x-y\rVert}, $$
the negative particle gradient of $\tfrac12\mathcal E^2$ is
$$ \begin{aligned} v_{\mathcal E}(x) &= \mathbb E_{y\sim P} \left[ \frac{y-x}{\lVert y-x\rVert} \right]\\ &\quad- \mathbb E_{x'\sim Q_\theta} \left[ \frac{x'-x}{\lVert x'-x\rVert} \right]. \end{aligned} \tag{7.13} $$
Every contribution is a unit vector: attraction toward the data minus attraction toward the model's own crowd. Subtracting the second term is what turns self-attraction into repulsion.
Three-Body Scattering: one triplet instead of all pairs¶
The full energy field compares every generated particle with every data and generated source, costing $O(B^2)$ interactions for batch size $B$. Three-Body Scattering for Generative Modeling keeps one projectile $x$, draws one real source $y\sim P$ and one independent generated source $x'\sim Q_\theta$, and uses
$$ \widehat v_{\mathrm{TBSM}}(x;y,x') = \frac{y-x}{\lVert y-x\rVert} - \frac{x'-x}{\lVert x'-x\rVert}. \tag{7.14} $$
Conditioned on the projectile,
$$ \mathbb E_{y,x'} \left[ \widehat v_{\mathrm{TBSM}}(x;y,x') \mid x \right] = v_{\mathcal E}(x). \tag{7.15} $$
The field and the resulting current-parameter regression gradient are therefore unbiased; the scalar regression-loss value is not an unbiased estimate of energy distance. The independent generated source is essential—reusing the projectile would remove the repulsive draw.
def mean_unit_pull(x, y):
'''Mean unit vectors pointing from each x toward a batch y.'''
delta = y[None, :, :] - x[:, None, :]
return (
delta / delta.norm(dim=-1, keepdim=True).clamp_min(1e-6)
).mean(dim=1)
def energy_field(x, y):
'''Rescaled negative particle gradient of one-half energy distance squared.'''
return mean_unit_pull(x, y) - mean_unit_pull(x, x)
def paired_unit(source, projectile):
'''One unit bearing from each projectile toward its paired source.'''
delta = source - projectile
return delta / delta.norm(dim=1, keepdim=True).clamp_min(1e-6)
# Verify the full-batch field before replacing its expectations by one triplet.
torch.manual_seed(11)
energy_x = (1.5 * sample_source(64)).double()
energy_y = sample_ring(64).double()
gradient_x = energy_x.clone().requires_grad_(True)
energy_squared = (
2 * torch.cdist(gradient_x, energy_y).mean()
- torch.cdist(gradient_x, gradient_x).mean()
- torch.cdist(energy_y, energy_y).mean()
)
negative_gradient = -len(gradient_x) * torch.autograd.grad(
0.5 * energy_squared, gradient_x
)[0]
energy_error = (negative_gradient - energy_field(energy_x, energy_y)).abs().max()
print(f"energy field vs rescaled -grad energy distance: max abs diff {energy_error:.1e}")
# Averaging K independent triplets should reduce source-sampling variance.
torch.manual_seed(13)
projectiles = sample_source(256)
reference_field = (
mean_unit_pull(projectiles, sample_ring(3000))
- mean_unit_pull(projectiles, sample_source(3000))
)
source_counts = [1, 4, 16, 64]
field_mse = []
for count in source_counts:
real = sample_ring(len(projectiles) * count).view(len(projectiles), count, 2)
generated = sample_source(len(projectiles) * count).view(len(projectiles), count, 2)
real_delta = real - projectiles[:, None, :]
generated_delta = generated - projectiles[:, None, :]
estimate = (
real_delta / real_delta.norm(dim=-1, keepdim=True).clamp_min(1e-6)
- generated_delta / generated_delta.norm(dim=-1, keepdim=True).clamp_min(1e-6)
).mean(dim=1)
field_mse.append((estimate - reference_field).square().mean().item())
print(f"{count:2d} source pair(s) per projectile: field MSE={field_mse[-1]:.4f}")
fig, ax = plt.subplots(figsize=(4.8, 3.2))
ax.loglog(source_counts, field_mse, "o-", color="#7c3aed")
ax.set(xlabel="source pairs per projectile", ylabel="field MSE")
ax.grid(alpha=0.3)
plt.tight_layout()
plt.show()
energy field vs rescaled -grad energy distance: max abs diff 6.7e-16 1 source pair(s) per projectile: field MSE=0.7354 4 source pair(s) per projectile: field MSE=0.1909 16 source pair(s) per projectile: field MSE=0.0494 64 source pair(s) per projectile: field MSE=0.0132
energy_drifting = train_moved_field(
energy_field,
steps=2000,
batch_size=200,
learning_rate=1e-3,
seed=52,
description="energy field",
)
record_drift("energy field", energy_drifting.sample(EVAL_N))
def train_tbsm(steps=4000, batch_size=512, learning_rate=1e-3, seed=53):
'''Train with one real and one independent generated source per projectile.'''
torch.manual_seed(seed)
generator = Generator()
optimizer = torch.optim.Adam(generator.parameters(), lr=learning_rate)
for _ in tqdm(range(steps), desc="Three-Body Scattering"):
x = generator.sample(batch_size)
with torch.no_grad():
detached_x = x.detach()
real_source = sample_ring(batch_size)
generated_source = generator.sample(batch_size).detach()
field = (
paired_unit(real_source, detached_x)
- paired_unit(generated_source, detached_x)
)
target = detached_x + field
loss = (x - target).square().sum(dim=1).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
return generator
tbsm = train_tbsm()
record_drift("TBSM", tbsm.sample(EVAL_N))
fig, ax = plt.subplots(1, 3, figsize=(10.8, 3.5))
for axis, key, title in (
(ax[0], "energy field", "full energy field"),
(ax[1], "TBSM", "one-triplet TBSM"),
):
result = DRIFT_RUNS[key]
plot_samples(
axis,
result["samples"],
f"{title}\nED={result['energy']:.3f} stray={100 * result['stray']:.1f}%",
by_mode=True,
)
plot_samples(ax[2], DATA, "target", color="#b45309", by_mode=True)
plt.tight_layout()
plt.show()
energy field: 100%|██████████| 2000/2000 [00:02<00:00, 875.25it/s]
energy field ED=0.0216 modes=8/8 stray=40.8%
Three-Body Scattering: 100%|██████████| 4000/4000 [00:06<00:00, 590.09it/s]
TBSM ED=0.0302 modes=8/8 stray=45.2%
Reading the energy fields¶
The gradient check first confirms that (7.13) is the rescaled negative particle gradient of energy distance. The full field is therefore exactly the B.2 rewrite of direct scalar energy-distance descent; matched runs produce the same result. Both field estimators then find all eight modes, but the unit-vector field leaves substantial mass moving around the destinations instead of settling tightly onto them. The full field reaches energy distance $0.0216$ with $40.8\%$ stray mass. A unit bearing has length one no matter how close its source is, so finite-batch fluctuations do not automatically shrink near the fixed point.
The one-triplet version adds source-sampling variance: its field MSE falls from $0.7354$ with one source pair to $0.0132$ with $64$, close to the expected inverse-sample trend. With one triplet, the trained generator reaches energy distance $0.0302$ and $45.2\%$ stray mass. Its advantage is computational rather than statistical: it removes dense pairwise matrices. TBSM's full image-scale method can learn a tracker for the conditional mean field and mix that prediction with the instantaneous triplet; we leave the raw unbiased tradeoff visible here.
5. Sinkhorn Drifting: transport plans become barycentric arrows¶
The Sinkhorn family lets us choose the ground cost. To obtain the Sinkhorn-Drifting field, use $c(x,y)=\tfrac12\lVert x-y\rVert^2$:
$$ \begin{aligned} \mathrm{OT}_\tau(Q,P) &= \min_{\pi\in\Pi(Q,P)} \mathbb E_{(x,y)\sim\pi} \left[ \tfrac12\lVert x-y\rVert^2 \right]\\ &\quad+ \tau\, \mathrm{KL}(\pi\,\|\,Q\otimes P). \end{aligned} \tag{7.16} $$
Debias it with the two self-costs:
$$ \begin{aligned} \mathrm S_\tau(Q,P) &= \mathrm{OT}_\tau(Q,P)\\ &\quad- \tfrac12\mathrm{OT}_\tau(Q,Q)\\ &\quad- \tfrac12\mathrm{OT}_\tau(P,P). \end{aligned} \tag{7.17} $$
For empirical batches $Q=\tfrac1n\sum_i\delta_{x_i}$ and $P=\tfrac1m\sum_j\delta_{y_j}$, let $\pi^{Q,P}_{ij}$ and $\pi^{Q,Q}_{ij}$ denote the cross and self transport plans. Row-normalize the mass leaving $x_i$, then average its destinations. These barycentric projections are
$$ \begin{aligned} T_\tau^{Q,P}(x_i) &= \sum_{j=1}^{m} \frac{\pi^{Q,P}_{ij}} {\sum_{k=1}^{m}\pi^{Q,P}_{ik}} y_j,\\ T_\tau^{Q,Q}(x_i) &= \sum_{j=1}^{n} \frac{\pi^{Q,Q}_{ij}} {\sum_{k=1}^{n}\pi^{Q,Q}_{ik}} x_j. \end{aligned} \tag{7.18} $$
For a fully converged uniform-marginal plan, each denominator is $1/n$. We keep the row sums explicit because a finite number of Sinkhorn rescalings satisfies the marginals only approximately.
From plans to arrows.
Step 1. At the optimal cross plan, the envelope theorem lets us differentiate only the quadratic cost. Its particle gradient is proportional to $x-T_\tau^{Q,P}(x)$.
Step 2. In the self-cost, the particle appears in both arguments. Those two derivatives cancel the factor $\tfrac12$, leaving $x-T_\tau^{Q,Q}(x)$.
Step 3. Subtract the two terms. The bare $x$ cancels, giving
$$ v_{\mathrm{Sinkhorn}}(x) = T_\tau^{Q,P}(x) - T_\tau^{Q,Q}(x). \tag{7.19} $$
The cross plan says where the data wants this particle; the self plan says where its own local crowd already is. Their difference is again attraction minus repulsion, now weighted by soft full-dimensional transport rather than kernels or unit bearings.
def uniform_sinkhorn_plan(cost, tau=0.05, n_iter=15):
'''Approximate entropic plan from alternating uniform-marginal rescaling.'''
n, m = cost.shape
f = cost.new_zeros(n)
g = cost.new_zeros(m)
for _ in range(n_iter):
f = -tau * torch.logsumexp(
(g[None, :] - cost) / tau - math.log(m),
dim=1,
)
g = -tau * torch.logsumexp(
(f[:, None] - cost) / tau - math.log(n),
dim=0,
)
return torch.exp((f[:, None] + g[None, :] - cost) / tau) / (n * m)
def barycentric_projection(plan, points):
'''Conditional mean of points under each row of a transport plan.'''
conditional = plan / plan.sum(dim=1, keepdim=True).clamp_min(1e-12)
return conditional @ points
def sinkhorn_field(x, y, tau=0.05, n_iter=15):
'''Cross barycentric projection minus the model self-projection.'''
quadratic_cost = lambda a, b: 0.5 * torch.cdist(a, b).square()
cross_plan = uniform_sinkhorn_plan(quadratic_cost(x, y), tau, n_iter)
self_plan = uniform_sinkhorn_plan(quadratic_cost(x, x), tau, n_iter)
cross_barycenter = barycentric_projection(cross_plan, y)
self_barycenter = barycentric_projection(self_plan, x)
return cross_barycenter - self_barycenter
# Check the field against the point gradient with nearly converged frozen plans.
torch.manual_seed(17)
sinkhorn_x = 1.5 * sample_source(64)
sinkhorn_y = sample_ring(64)
quadratic_cost = lambda a, b: 0.5 * torch.cdist(a, b).square()
def frozen_transport_part(a, b, tau=0.5, n_iter=200):
'''Transport-cost part with optimal plans detached by the envelope theorem.'''
plan = uniform_sinkhorn_plan(quadratic_cost(a, b), tau, n_iter).detach()
return (plan * quadratic_cost(a, b)).sum()
gradient_x = sinkhorn_x.clone().requires_grad_(True)
frozen_sinkhorn = (
frozen_transport_part(gradient_x, sinkhorn_y)
- 0.5 * frozen_transport_part(gradient_x, gradient_x)
- 0.5 * frozen_transport_part(sinkhorn_y, sinkhorn_y)
)
negative_gradient = -len(gradient_x) * torch.autograd.grad(
frozen_sinkhorn, gradient_x
)[0]
checked_field = sinkhorn_field(sinkhorn_x, sinkhorn_y, tau=0.5, n_iter=200)
relative_error = (negative_gradient - checked_field).norm() / checked_field.norm()
print(f"barycentric field vs rescaled -grad Sinkhorn: relative diff {relative_error:.1e}")
barycentric field vs rescaled -grad Sinkhorn: relative diff 3.2e-07
sinkhorn_drifting = train_moved_field(
lambda x, y: sinkhorn_field(x, y, tau=0.05, n_iter=20),
steps=4000,
batch_size=256,
learning_rate=1e-3,
seed=54,
description="Sinkhorn Drifting",
cosine_decay=True,
)
record_drift("Sinkhorn drifting", sinkhorn_drifting.sample(EVAL_N))
show_drift("Sinkhorn drifting", "Sinkhorn Drifting")
Sinkhorn Drifting: 100%|██████████| 4000/4000 [01:33<00:00, 42.69it/s]
Sinkhorn drifting ED=0.0107 modes=8/8 stray=9.9%
Reading the Sinkhorn field¶
Sinkhorn Drifting reaches all eight modes with energy distance $0.0107$ and $9.9\%$ stray mass. Its destinations are more compact than the two energy-distance fields. Soft transport uses the whole two-dimensional geometry, while the cross-minus-self subtraction prevents the entropic plan from pulling every generated particle toward the same local average.
The full field is the B.2 rewrite of quadratic-cost Sinkhorn descent. Our finite number of Sinkhorn rescalings makes its barycentric projections a truncated approximation, so a matched scalar B.1 run is close rather than identical. Chapter 6's Sinkhorn number is not that control: it used the unsquared Euclidean ground cost, whereas (7.16) uses squared distance. The other bills remain those of transport: two dense pairwise plans, quadratic batch memory, and a temperature $\tau$ that trades sharp assignments against smooth optimization.
6. Same sample comparisons, read as fields¶
All four runs retain one-step inference: generation is still one evaluation of $G_\theta$. The iteration happens only during training as each batch proposes new local targets.
One-step drifting is a fast-growing area, so this chapter is not a complete catalog. We cover only representative methods that make the B.1-to-B.2 construction visible; related variants may design fields directly or measure them in frozen feature spaces.
keys = ["MMD field", "energy field", "TBSM", "Sinkhorn drifting"]
fig, axes = plt.subplots(1, 5, figsize=(13.8, 3.0))
for axis, key in zip(axes, keys):
result = DRIFT_RUNS[key]
title = (
f"{key}\nED={result['energy']:.3f} "
f"stray={100 * result['stray']:.1f}%"
)
plot_samples(axis, 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.2: sample discrepancies read as particle fields", fontsize=11)
plt.tight_layout()
plt.show()
print(f"{'method':20s} {'ED':>8s} {'modes':>7s} {'stray':>8s}")
for key in keys:
result = DRIFT_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 field 0.0096 8/8 18.0% energy field 0.0216 8/8 40.8% TBSM 0.0302 8/8 45.2% Sinkhorn drifting 0.0107 8/8 9.9% target 8/8 1.3%
The comparison separates three choices that the phrase “drifting method” can hide.
- Which scalar comparison supplies the field? MMD produces kernel-weighted arrows, energy distance produces unit bearings, and Sinkhorn produces barycentric transport arrows.
- How accurately is that field read? Full batches average many interactions; TBSM spends one triplet per projectile and accepts variance in exchange for linear interaction cost.
- Where is the field measured? Coordinates are enough for this two-dimensional example; realistic image methods usually move to frozen representation spaces.
Against the target's $1.3\%$ stray-mass floor, the MMD and Sinkhorn fields finish at $18.0\%$ and $9.9\%$. The energy field and its one-triplet approximation finish at $40.8\%$ and $45.2\%$. All four still find every mode, so the structural metric again reveals differences hidden by mode coverage alone.
The main point is simpler than the method names. Strategy B.1 and B.2 solve the same sample-based distribution-matching problem. B.1 lets automatic differentiation carry a scalar discrepancy all the way to the parameters. B.2 opens that backward pass at the generated samples, freezes the resulting directions, and presents the same update as supervised regression toward locally moved targets.