Back to all posts
Distributional Distillation · Part III

From Score Differences to Particle Flows

DMD, W-Flow, and Drifting Models all tell generated samples where to move. We build that idea from the ground up, then ask the key question: was the motion derived from a global objective, or designed directly because it works?

Fu-Yun Wang · 2026 · Mathematical notes

1. A global score does not yet tell each sample where to move

The idea before the notation A distribution objective is a single report-card score for the whole generator. Training, however, must move individual generated samples. The derivation below turns the global report card into one arrow attached to each sample, then backpropagates those arrows through the generator.

Sample a latent \(\mathbf z\), pass it through the generator, and call the resulting output distribution \(q_\theta\):

\[ \mathbf x_\theta=f_\theta(\mathbf z), \qquad q_\theta=(f_\theta)_\#p_{\mathrm{ref}}, \qquad J_\theta=\frac{\partial f_\theta(\mathbf z)}{\partial\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?

\[ \phi_q(\mathbf x) \equiv \frac{\delta\mathcal F_p}{\delta q}(q)(\mathbf x). \]

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:

\[ \begin{aligned} \nabla_\theta\mathcal F_p(q_\theta) &=\int\nabla_\theta q_\theta(\mathbf x) \phi_{q_\theta}(\mathbf x)d\mathbf x\\ &=\mathbb E_{\mathbf z}\!\left[ J_\theta^\top \nabla_{\mathbf x}\phi_{q_\theta}(\mathbf x_\theta) \right]. \end{aligned} \] (1)

Name the negative spatial gradient

\[ \boldsymbol\xi_q(\mathbf x) \equiv-\nabla_{\mathbf x}\phi_q(\mathbf x). \]

Equation (1) becomes

\[ \boxed{ \mathcal F_p(q) \longrightarrow \phi_q=\frac{\delta\mathcal F_p}{\delta q} \xrightarrow{W_2} \boldsymbol\xi_q=-\nabla\phi_q \longrightarrow \nabla_\theta\mathcal F_p(q_\theta) =-\mathbb E[J_\theta^\top\boldsymbol\xi_q].} \] (2)

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:

\[ \frac{d\mathbf x_\tau}{d\tau} =\boldsymbol\xi_\tau(\mathbf x_\tau). \]

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:

\[ \boxed{ \partial_\tau q_\tau +\nabla\!\cdot(q_\tau\boldsymbol\xi_\tau)=0.} \] (3)

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

1

Functional chain rule

\[ \frac{d}{d\tau}\mathcal F_p(q_\tau) =\int\phi_\tau(\mathbf x) \partial_\tau q_\tau(\mathbf x)d\mathbf x. \]
2

Insert probability conservation

\[ \frac{d}{d\tau}\mathcal F_p(q_\tau) =-\int\phi_\tau \nabla\!\cdot(q_\tau\boldsymbol\xi_\tau)d\mathbf x. \]
3

Integrate by parts

Assuming the boundary term vanishes at infinity,

\[ \frac{d}{d\tau}\mathcal F_p(q_\tau) =\int q_\tau \nabla\phi_\tau^\top\boldsymbol\xi_\tau d\mathbf x =\mathbb E_{q_\tau}[ \nabla\phi_\tau^\top\boldsymbol\xi_\tau]. \] (4)

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:

\[ \begin{aligned} \arg\min_{\boldsymbol\xi} \mathbb E_{q_\tau}\!\left[ \nabla\phi_\tau^\top\boldsymbol\xi +\tfrac12\|\boldsymbol\xi\|^2\right] &= \arg\min_{\boldsymbol\xi} \mathbb E_{q_\tau}\!\left[ \tfrac12\|\boldsymbol\xi+\nabla\phi_\tau\|^2 -\tfrac12\|\nabla\phi_\tau\|^2 \right]\\ &=-\nabla\phi_\tau. \end{aligned} \]

Thus the Wasserstein gradient-flow velocity is

\[ \boxed{ \boldsymbol\xi_\tau =-\nabla\frac{\delta\mathcal F_p}{\delta q}(q_\tau).} \] (5)

Substitute it back into Equation (4):

\[ \frac{d}{d\tau}\mathcal F_p(q_\tau) =-\mathbb E_{q_\tau}\!\left[ \left\| \nabla\frac{\delta\mathcal F_p}{\delta q}(q_\tau) \right\|^2\right]\leq0. \] (6)
“It stops when \(p=q\)” is not enough Many hand-designed fields become zero when model and data match. A Wasserstein gradient flow promises more: its arrows come from an explicit scalar objective, and following them is guaranteed to decrease that objective according to Equation (6).

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

\[ \widetilde{\mathbf x} =\operatorname{sg}\!\left[ f_{\theta_0}(\mathbf z) +\eta\boldsymbol\xi_{q_{\theta_0}} (f_{\theta_0}(\mathbf z)) \right], \] \[ \mathcal L_{\mathrm{sg}}(\theta;\theta_0) =\frac{1}{2\eta} \mathbb E_{\mathbf z} \|f_\theta(\mathbf z)-\widetilde{\mathbf x}\|^2. \] (7)

At the point where the target was formed,

\[ \begin{aligned} \nabla_\theta\mathcal L_{\mathrm{sg}} \big|_{\theta=\theta_0} &=\frac1\eta\mathbb E\!\left[ J_{\theta_0}^\top (f_{\theta_0}(\mathbf z)-\widetilde{\mathbf x}) \right]\\ &=-\mathbb E\!\left[ J_{\theta_0}^\top \boldsymbol\xi_{q_{\theta_0}} (f_{\theta_0}(\mathbf z)) \right]\\ &=\nabla_\theta\mathcal F_p(q_\theta) \big|_{\theta=\theta_0}. \end{aligned} \] (8)

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.

Training time is not diffusion time The clock \(\tau\) describes how the model distribution changes over optimization steps. It is neither the diffusion noise level \(t\) nor inference time. A trained one-step generator does not integrate this continuity equation at inference; the local distribution updates have been compressed into its weights.

4. DMD's two scores are an attraction–spreading flow

Fix a diffusion level \(t\) while the optimization clock \(\tau\) runs, and choose

\[ \mathcal F_t(q)=\mathrm{KL}(q\|p_t). \]

Its first variation is \(\log(q/p_t)+1\). Applying Equation (5) gives

\[ \boxed{ \boldsymbol\xi_{\tau,t}(\mathbf x) =-\nabla\left(\log\frac{q_{\tau,t}}{p_t}+1\right) =\underbrace{\nabla\log p_t}_{s_{p_t}} -\underbrace{\nabla\log q_{\tau,t}}_{s_{q_{\tau,t}}}.} \] (9)

Put the field back into the population equation

Substitute the field into Equation (3):

\[ \begin{aligned} \partial_\tau q_{\tau,t} &=-\nabla\!\cdot[ q_{\tau,t}(s_{p_t}-s_{q_{\tau,t}})]\\ &=-\nabla\!\cdot(q_{\tau,t}s_{p_t}) +\nabla\!\cdot(q_{\tau,t}\nabla\log q_{\tau,t})\\ &=-\nabla\!\cdot(q_{\tau,t}s_{p_t}) +\nabla\!\cdot(\nabla q_{\tau,t})\\ &=\boxed{ \Delta q_{\tau,t} -\nabla\!\cdot(q_{\tau,t}s_{p_t}).} \end{aligned} \] (10)

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

\[ \frac{d}{d\tau}\log p_t(\mathbf x_\tau) =\nabla\log p_t(\mathbf x_\tau)^\top \dot{\mathbf x}_\tau =\|s_{p_t}(\mathbf x_\tau)\|^2\geq0, \]

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.

The book's LaTeX figure mapping reverse KL from a distributional energy to a local potential and particle velocity, with particles moving downhill across the potential landscape.
The original LaTeX Wasserstein-flow figure from the book. Reverse KL becomes a local potential, whose negative spatial gradient moves particles downhill; after they move, the student density and its field must be recomputed.

At equilibrium \(q=p_t\), the two terms cancel exactly:

\[ \Delta p_t-\nabla\!\cdot(p_ts_{p_t}) =\Delta p_t-\nabla\!\cdot(\nabla p_t)=0. \]

DMD averages this KL-WGF over noise levels and projects each field through the noisy generator Jacobian:

\[ \mathcal F_{\mathrm{DMD}}(\theta) =\mathbb E_t[\gamma(t)\mathrm{KL}(q_{\theta,t}\|p_t)], \] \[ \nabla_\theta\mathcal F_{\mathrm{DMD}} =-\mathbb E_{t,\mathbf z,\boldsymbol\epsilon}\!\left[ \gamma(t)\boldsymbol\xi_t(\mathbf x_t)^\top \frac{\partial\mathbf x_t}{\partial\theta}\right]. \] (11)

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:

\[ \mathrm{OT}_{\varepsilon_{\mathrm{OT}}}(q,p) = \inf_{\pi\in\Pi(q,p)} \left\{ \int\tfrac12\|\mathbf x-\mathbf y\|^2d\pi +\varepsilon_{\mathrm{OT}}\mathrm{KL}(\pi\|q\otimes p) \right\}, \] \[ S_{\varepsilon_{\mathrm{OT}}}(q,p) =\mathrm{OT}_{\varepsilon_{\mathrm{OT}}}(q,p) -\tfrac12\mathrm{OT}_{\varepsilon_{\mathrm{OT}}}(q,q) -\tfrac12\mathrm{OT}_{\varepsilon_{\mathrm{OT}}}(p,p). \] (12)

Let the conditional barycentric maps of the optimal couplings be

\[ T^\varepsilon_{q,p}(\mathbf x) =\int\mathbf y\,\pi^\varepsilon_{q,p}(d\mathbf y\mid\mathbf x), \qquad T^\varepsilon_{q,q}(\mathbf x) =\int\mathbf x'\,\pi^\varepsilon_{q,q}(d\mathbf x'\mid\mathbf x). \]

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

\[ \nabla\frac{\delta\mathrm{OT}_{\varepsilon_{\mathrm{OT}}}(\cdot,p)} {\delta q}(q)(\mathbf x) =\int(\mathbf x-\mathbf y) \pi^\varepsilon_{q,p}(d\mathbf y\mid\mathbf x) =\mathbf x-T^\varepsilon_{q,p}(\mathbf x). \]

The self-transport term appears twice because \(q\) occupies both arguments; symmetry cancels the factor \(1/2\) in Equation (12). Therefore

\[ \nabla\frac{\delta S_{\varepsilon_{\mathrm{OT}}}(\cdot,p)} {\delta q}(q)(\mathbf x) =T^\varepsilon_{q,q}(\mathbf x) -T^\varepsilon_{q,p}(\mathbf x). \]

Negating it yields the exact Sinkhorn-WGF velocity:

\[ \boxed{ \boldsymbol\xi^\varepsilon_{q,p}(\mathbf x) =T^\varepsilon_{q,p}(\mathbf x) -T^\varepsilon_{q,q}(\mathbf x).} \] (13)

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)\).

One-line unification \[ \text{DMD}=\text{KL-WGF on noisy marginals }q_t, \qquad \text{W-Flow}=\text{Sinkhorn-WGF on the clean law }q. \] The energy and field estimator differ; the functional → field → Jacobian pipeline is the same.

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

\[ \widetilde{\mathbf u} =\operatorname{sg}\!\left[ \mathbf u+\eta\mathbf V_{\phi_\#p,\phi_\#q}(\mathbf u) \right], \qquad \mathcal L_{\mathrm{drift}} =\|\phi(G_\theta(\mathbf z))-\widetilde{\mathbf u}\|^2, \] (14)

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

\[ k_h(\mathbf u,\mathbf v) =\exp\!\left(-\frac{\|\mathbf u-\mathbf v\|^2}{2h^2}\right), \] \[ M_r^h(\mathbf u) =\frac{\mathbb E_{\mathbf v\sim r} [k_h(\mathbf u,\mathbf v)\mathbf v]} {\mathbb E_{\mathbf v\sim r} [k_h(\mathbf u,\mathbf v)]}, \qquad r\in\{p,q\}. \] (15)

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:

\[ \begin{aligned} \nabla_{\mathbf u}\log r_h(\mathbf u) &=\frac{\nabla_{\mathbf u}Z_r^h(\mathbf u)} {Z_r^h(\mathbf u)}\\ &=\frac{ \mathbb E_{\mathbf v\sim r}[ \nabla_{\mathbf u}k_h(\mathbf u,\mathbf v)]} {\mathbb E_{\mathbf v\sim r}[k_h(\mathbf u,\mathbf v)]}\\ &=\frac{ \mathbb E_{\mathbf v\sim r}[ k_h(\mathbf u,\mathbf v)(\mathbf v-\mathbf u)/h^2]} {\mathbb E_{\mathbf v\sim r}[k_h(\mathbf u,\mathbf v)]}\\ &=\frac{M_r^h(\mathbf u)-\mathbf u}{h^2}. \end{aligned} \] (16)

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:

\[ \begin{aligned} \mathbf V_{p,q}(\mathbf u) &= \bigl(M_p^h(\mathbf u)-\mathbf u\bigr) -\bigl(M_q^h(\mathbf u)-\mathbf u\bigr)\\ &=M_p^h(\mathbf u)-M_q^h(\mathbf u)\\ &=\boxed{ h^2\bigl( \nabla\log p_h(\mathbf u) -\nabla\log q_h(\mathbf u) \bigr).} \end{aligned} \] (17)

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:

A useful equilibrium field is not automatically the WGF of a known energy Changing the kernel, normalization, feature scaling, or neighborhood rule may preserve \(\mathbf V_{p,q}=0\) at \(p=q\) while destroying descent of any known global functional. Even applying \(\nabla\log p_h-\nabla\log q_h\) directly to clean particles is not automatically the exact WGF of \(\mathrm{KL}(q_h\|p_h)\) viewed as a functional of the clean law \(q\): the smoothing operator must be pulled back in the first variation.

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

DMDW-FlowDrifting Models
Starting pointNoisy marginal KLSinkhorn divergenceDesigned equilibrium field
Field\(s_{p_t}-s_{q_t}\)\(T^\varepsilon_{q,p}-T^\varepsilon_{q,q}\)Target neighborhood minus model neighborhood
EstimatorTeacher and fake denoisersMini-batch Sinkhorn plansKernel/feature-space batch statistics
Comparison spaceGaussian-noised image/latent marginalsUsually clean sample spacePretrained semantic feature space
Explicit global functionalYesYesNot required for the practical field
Generator updateField–Jacobian contractionDetached field-step regressionDetached feature-target regression

The useful unification is therefore not “all three are the same method.” It is a shared workflow:

\[ \boxed{ \text{choose or design a population field} \;\longrightarrow\; \text{move generated particles} \;\longrightarrow\; \text{project that motion through the generator}.} \]

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.