- Functional → first variation → field → generator gradient
- Why Wasserstein geometry selects \(-\nabla\delta\mathcal F/\delta q\)
- Why detached regression implements the same local update
- DMD as KL gradient flow on noisy marginals
- W-Flow as Sinkhorn gradient flow
- Drifting Models: the precise overlap and the precise limit
1. A global score does not yet tell each sample where to move
Sample a latent \(\mathbf z\), pass it through the generator, and call the resulting output distribution \(q_\theta\):
Let the scalar \(\mathcal F_p(q)\) grade how far the whole generated distribution is from a target \(p\). Its first variation asks a local question: if we add or remove a tiny amount of probability near \(\mathbf x\), how does that global grade change?
Call that local sensitivity \(\phi_q(\mathbf x)\). The chain rule pairs it with the way the density changes; the identity from Part I then rewrites the same change in terms of moving generated samples:
Name the negative spatial gradient
Equation (1) becomes
The first and last arrows are ordinary chain-rule steps. The middle arrow requires an extra choice. Knowing that one distribution is better than another does not uniquely specify the cheapest way to move probability mass between them. A geometry—or distance rule—on distributions supplies that missing choice.
2. Wasserstein geometry chooses the cheapest useful motion
Imagine training as a continuous clock \(\tau\). At every location \(\mathbf x\), place an arrow \(\boldsymbol\xi_\tau(\mathbf x)\). A generated sample reads the arrow at its current location as its velocity:
When all samples follow these arrows, probability cannot appear or disappear; it can only flow across the boundary of a region. This conservation law is the continuity equation:
Equation (3) describes the effect of any chosen arrow field; it does not yet tell us which arrows are best. To measure how a field changes the objective, let \(\phi_\tau=\delta\mathcal F_p/\delta q(q_\tau)\). Then
Functional chain rule
Insert probability conservation
Integrate by parts
Assuming the boundary term vanishes at infinity,
Equation (4) holds for any possible motion. Wasserstein geometry chooses a compromise: decrease the objective quickly, but penalize large sample movements. Completing the square makes the winning field explicit:
Thus the Wasserstein gradient-flow velocity is
Substitute it back into Equation (4):
3. Why these methods often look like ordinary regression in code
If every generated sample were independent, we could simply move it by \(\eta\boldsymbol\xi_q(\mathbf x)\). A neural generator cannot move samples independently because they share the same weights. Backpropagation finds the parameter update that best approximates all requested motions together. We can implement this directly, or turn each arrow into a frozen regression target.
Freeze the current iterate \(\theta_0\) and define
At the point where the target was formed,
This identity explains the shared code pattern: take the current sample, add one small field step, freeze the result, and regress toward it. But seeing an MSE loss in code does not tell us where the arrow came from. The same regression wrapper can implement a principled gradient flow or a directly designed heuristic field.
4. DMD's two scores are an attraction–spreading flow
Fix a diffusion level \(t\) while the optimization clock \(\tau\) runs, and choose
Its first variation is \(\log(q/p_t)+1\). Applying Equation (5) gives
Put the field back into the population equation
Substitute the field into Equation (3):
Equation (10) separates the two-score update into two population effects. First consider only the target score. If particles follow \(\dot{\mathbf x}=s_{p_t}(\mathbf x)\), then
so each particle climbs toward a region the target considers more likely. Now keep only the fake-score contribution. It becomes the heat equation \(\partial_\tau q=\Delta q\), whose solution spreads the model law like Gaussian smoothing. This does not mean DMD adds another random-noise step; the apparent diffusion is the population effect of deterministic arrows that depend on the current distribution.
At equilibrium \(q=p_t\), the two terms cancel exactly:
DMD averages this KL-WGF over noise levels and projects each field through the noisy generator Jacobian:
Using \(s_{p_t}=-\boldsymbol\epsilon_\varphi/\sigma_t\) and \(s_{q_t}=-\boldsymbol\epsilon_\psi/\sigma_t\) recovers the usual DMD teacher-minus-fake denoiser gradient. The WGF view does not replace the two-score derivation; it explains the distributional motion induced by it.
5. W-Flow uses the same recipe with a transport objective
DMD asks score models how to move each sample. W-Flow instead solves a softened batch transport problem: which real samples should generated samples move toward, while avoiding a biased collapse caused by entropy regularization? Its global objective is the Sinkhorn divergence:
Let the conditional barycentric maps of the optimal couplings be
By the envelope theorem, the optimal coupling need not be differentiated explicitly. For the cross-transport term, the spatial gradient of the first variation is
The self-transport term appears twice because \(q\) occupies both arguments; symmetry cancels the factor \(1/2\) in Equation (12). Therefore
Negating it yields the exact Sinkhorn-WGF velocity:
The first map says where real mass lies relative to the current generated sample. The second asks how the generated batch transports to itself; subtracting it removes the bias introduced by softened transport and supplies self-interaction. W-Flow estimates both maps from batches, then regresses toward \(\mathbf x+\eta\boldsymbol\xi^\varepsilon_{q,p}(\mathbf x)\).
6. Drifting Models build the arrows directly in feature space
Drifting Models reverse the order of reasoning. Instead of choosing a global objective and deriving its arrows, they begin by constructing useful arrows from neighboring real and generated samples. A frozen visual encoder \(\phi\) first maps images into a semantic feature space, where distance is more meaningful than raw pixel distance. For a generated feature \(\mathbf u=\phi(G_\theta(\mathbf z))\), they construct
and backpropagate through the frozen encoder to the generator. This is the same detached-target wrapper as Equation (7). The remaining question is whether the chosen \(\mathbf V\) is the gradient field of a known global objective.
The cleanest connection: a local average is a score estimator
Take one generated feature \(\mathbf u\). Average nearby real features, giving closer neighbors more weight; do the same with a separate generated batch. With a Gaussian weighting rule, those local averages have an exact connection to scores. Let \(p\) be the target feature distribution and \(q\) the generated feature distribution, and define
Let \(Z_r^h(\mathbf u)=\mathbb E_{\mathbf v\sim r}[k_h(\mathbf u,\mathbf v)]\). Up to a constant independent of \(\mathbf u\), \(Z_r^h\) is the Gaussian-smoothed density \(r_h\). Differentiate:
Equation (16) says something intuitive: the vector from \(\mathbf u\) to its local weighted average points uphill in the smoothed density. It is a nonparametric score estimate. Subtract the generated-neighborhood vector from the real-neighborhood vector:
This is exactly the familiar target-score minus self-score shape, now estimated from neighbors instead of denoisers. The real neighborhood attracts the sample; subtracting the generated neighborhood prevents the whole batch from following its own density and makes the field vanish when the two local distributions agree.
Why this does not make every Drift update “the same as DMD”
Equation (17) is the precise reason the methods look related. It becomes the same field as DMD only when all of the following match:
- the query is distributed according to the same smoothed \(q_t\) whose score appears in the KL objective;
- the Gaussian kernel bandwidth represents the same corruption scale;
- target and model scores live in the same space;
- the update is contracted through the corresponding Jacobian with the same stop-gradient convention.
So the practical placement is precise. Drifting Models share DMD's target-minus-self idea and W-Flow's detached regression implementation. Their strong image model also uses semantic features and multiple scales, because raw pixels give a poor notion of similarity. But unless the practical field is derived from a scalar objective, it is best described as a carefully designed field that stops at the right equilibrium—not automatically as “DMD without a diffusion model” or as a Sinkhorn/KL gradient flow.
7. The comparison that survives the algebra
| DMD | W-Flow | Drifting Models | |
|---|---|---|---|
| Starting point | Noisy marginal KL | Sinkhorn divergence | Designed equilibrium field |
| Field | \(s_{p_t}-s_{q_t}\) | \(T^\varepsilon_{q,p}-T^\varepsilon_{q,q}\) | Target neighborhood minus model neighborhood |
| Estimator | Teacher and fake denoisers | Mini-batch Sinkhorn plans | Kernel/feature-space batch statistics |
| Comparison space | Gaussian-noised image/latent marginals | Usually clean sample space | Pretrained semantic feature space |
| Explicit global functional | Yes | Yes | Not required for the practical field |
| Generator update | Field–Jacobian contraction | Detached field-step regression | Detached feature-target regression |
The useful unification is therefore not “all three are the same method.” It is a shared workflow:
DMD and W-Flow derive the field from explicit energies under Wasserstein geometry. Drifting Models begin with a useful field in a learned representation. The Gaussian mean-shift special case exposes why the formulas can look identical; the first-variation and pullback conditions explain why the methods are not interchangeable in general.
References
[1] Yin et al. “One-step Diffusion with Distribution Matching Distillation.” CVPR 2024.
[2] Han et al. “One-Step Generative Modeling via Wasserstein Gradient Flows.” 2026.
[3] Deng et al. “Generative Modeling via Drifting.” 2026.
[4] Gretton et al. “On the Wasserstein Gradient Flow Interpretation of Drifting Models.” 2026.