Back to all posts
Distributional Distillation · Part I

From a Marginal KL to Two Scores

How do we train a one-step generator when there is no correct image assigned to each noise input? We begin with that practical question and derive the familiar two-score update one line at a time.

Fu-Yun Wang · 2026 · Mathematical notes
This article answers one question
  1. Why match noisy marginals?
  2. What does the direction of KL change?
  3. How do we derive the two-score identity?
  4. Why is every smooth \(f\)-divergence the same direction with a different weight?
  5. Why is the fake score an entropy gradient?

1. Stop grading samples one by one

Let a one-step generator be

\[ \mathbf x=G_\theta(\mathbf z), \qquad \mathbf z\sim\mathcal N(\mathbf0,\mathbf I), \qquad \mathbf x\sim q_\theta. \]

Suppose the teacher turns noise \(z_1\) into a cat and \(z_2\) into a dog. Trajectory distillation asks the student to reproduce those two assignments. Distributional distillation allows the student to swap them. What matters is that the collection of outputs still contains the right cats and dogs in the right proportions.

The formal name for the noise-to-image assignment is a coupling. Once we forget that assignment and keep only the output distribution, we keep the marginal. This extra freedom is useful when a one-step student is too small or too constrained to imitate the teacher's exact map.

Why not compare the clean output distributions directly? In image space, valid images occupy an extremely narrow part of all possible pixel arrays. The teacher and student can therefore miss each other almost completely, even if both produce recognizable images. In that case the KL can be infinite or provide almost no useful direction. The standard repair is simple: add the same type of Gaussian noise to both sides before comparing them.

\[ \mathbf x_t =\alpha_tG_\theta(\mathbf z)+\sigma_t\boldsymbol\epsilon, \qquad \boldsymbol\epsilon\sim\mathcal N(\mathbf0,\mathbf I), \] \[ \boxed{ \mathcal L_{\mathrm{marginal}}(\theta) =\mathbb E_t\!\left[\gamma(t) \mathrm{KL}\!\left(q_{\theta,t}\|p_t\right)\right].} \] (1)

At every \(t>0\), the added noise spreads both distributions out, so they overlap and can “see” each other. Large \(t\) mainly compares coarse layout; small \(t\) compares fine detail. Here noise is not merely damage—it turns an almost impossible clean-space comparison into a sequence of easier comparisons from coarse to fine.

The freedom–coverage trade-off A marginal objective frees the student from reproducing the teacher's exact latent-to-image map. But once no target is assigned to each latent, coverage is no longer automatic. This is the structural reason marginal methods can be sharp and still drop modes.

2. Forward and reverse KL fail differently

Start with one Gaussian component \(p_m=\mathcal N(m,\varsigma^2)\) and a single-Gaussian student \(q_\mu=\mathcal N(\mu,\varsigma^2)\). They share the same variance, so their cross-entropy is

\[ \boxed{ H(p_m,q_\mu) = \frac12\log(2\pi\varsigma^2) +\frac{\varsigma^2+(m-\mu)^2}{2\varsigma^2}.} \] (G)

Because \(\mathrm{KL}(p\|q)=H(p,q)-H(p)\), minimizing forward KL over \(q\) is the same as minimizing this cross-entropy.

Optional: derive Equation (G) from the Gaussian density
\[ -\log q_\mu(x) =\frac12\log(2\pi\varsigma^2) +\frac{(x-\mu)^2}{2\varsigma^2}. \]

Take the expectation under \(p_m\), then add and subtract its mean \(m\):

\[ \begin{aligned} H(p_m,q_\mu) &=\frac12\log(2\pi\varsigma^2) +\frac{\mathbb E_{p_m}[(x-\mu)^2]}{2\varsigma^2},\\ \mathbb E_{p_m}[(x-\mu)^2] &=\underbrace{\mathbb E_{p_m}[(x-m)^2]}_{\varsigma^2} +(m-\mu)^2. \end{aligned} \]

The two-mode example

Let the target contain two narrow modes, while the student can represent only one:

\[ p(x)=\tfrac12p_-(x)+\tfrac12p_+(x), \qquad p_\pm(x)=\mathcal N(\pm a,\varsigma^2), \qquad q_\mu(x)=\mathcal N(\mu,\varsigma^2), \qquad \varsigma\ll a. \]

Forward KL straddles the modes

The target \(p\) is not itself Gaussian. Instead, use the linearity of expectation to split its cross-entropy into the two Gaussian components:

\[ H(p,q_\mu) =\tfrac12H(p_-,q_\mu)+\tfrac12H(p_+,q_\mu). \]

Apply Equation (G) once with \(m=-a\) and once with \(m=+a\):

\[ \begin{aligned} H(p,q_\mu) &= \frac12\log(2\pi\varsigma^2) +\frac{ 2\varsigma^2+(-a-\mu)^2+(a-\mu)^2 }{4\varsigma^2}\\ &= \frac12\log(2\pi\varsigma^2) +\frac12+\frac{a^2+\mu^2}{2\varsigma^2}. \end{aligned} \]

Every term except \(\mu^2/(2\varsigma^2)\) is constant in \(\mu\), so the minimum is \(\mu=0\): a location between the modes where the target places almost no mass. A capacity-limited student covers both modes slightly and represents neither well.

Reverse KL locks onto one mode

Unlike forward KL, this reverse KL is difficult to evaluate directly: the logarithm contains a sum of two Gaussian densities, so it cannot be split across the two mixture components. Fortunately, gradient descent does not require a closed-form value; it only needs the derivative with respect to the student location \(\mu\). Write

\[ \begin{aligned} K(\mu) &=\mathrm{KL}(q_\mu\|p)\\ &=\mathbb E_{X\sim q_\mu} [\log q_\mu(X)-\log p(X)]. \end{aligned} \]

A sample from \(q_\mu=\mathcal N(\mu,\varsigma^2)\) can be written as \(X=\mu+\varsigma Z\), where \(Z\sim\mathcal N(0,1)\). Substituting this into the expectation gives

\[ K(\mu)=\mathbb E_Z\!\left[ \log q_\mu(\mu+\varsigma Z) -\log p(\mu+\varsigma Z) \right]. \]

Now expand the student term:

\[ \begin{aligned} \log q_\mu(\mu+\varsigma Z) &=-\frac12\log(2\pi\varsigma^2) -\frac{((\mu+\varsigma Z)-\mu)^2}{2\varsigma^2}\\ &=-\frac12\log(2\pi\varsigma^2)-\frac{Z^2}{2}. \end{aligned} \]

This expression no longer contains \(\mu\), so its total derivative is zero for every fixed \(Z\)—the expectation is not needed for this cancellation. Written with the chain rule, the two contributions are

\[ \begin{aligned} \frac{\partial}{\partial\mu}\log q_\mu(x) &=\frac{x-\mu}{\varsigma^2}, & \frac{\partial}{\partial x}\log q_\mu(x) &=-\frac{x-\mu}{\varsigma^2},\\ \frac{d}{d\mu}\log q_\mu(\mu+\varsigma Z) &=\left.\frac{\partial}{\partial\mu}\log q_\mu(x) \right|_{x=\mu+\varsigma Z} +\left.\frac{\partial}{\partial x}\log q_\mu(x) \right|_{x=\mu+\varsigma Z} \frac{d(\mu+\varsigma Z)}{d\mu}\\ &=\frac{Z}{\varsigma}-\frac{Z}{\varsigma}\cdot1=0. \end{aligned} \]

If \(x\) were held fixed, \(\partial_\mu\log q_\mu(x)\) would not be zero pointwise; only its expectation would vanish. Here \(x=\mu+\varsigma Z\) moves with \(\mu\), so the explicit parameter derivative and the sample-motion derivative cancel pointwise. This special simplification holds because the example changes only the mean while keeping the variance fixed.

The distribution of \(Z\) does not depend on \(\mu\), so we can now differentiate inside the expectation. Applying the chain rule to the remaining target term gives

\[ \begin{aligned} K'(\mu) &=\mathbb E_Z\!\left[ 0-\frac{d}{d\mu}\log p(\mu+\varsigma Z) \right]\\ &=-\mathbb E_Z\!\left[ \left.\frac{d}{dx}\log p(x)\right|_{x=\mu+\varsigma Z} \underbrace{\frac{d(\mu+\varsigma Z)}{d\mu}}_{1} \right]\\ &=\boxed{-\mathbb E_{X\sim q_\mu}[\nabla_x\log p(X)]}. \end{aligned} \] (R)

The last line merely renames \(\mu+\varsigma Z\) as \(X\). Since this random variable follows \(q_\mu\), the two expectations are identical. Equation (R) says that gradient descent moves the student along the average target score:

\[ \mu\leftarrow\mu-\eta K'(\mu) =\mu+\eta\, \mathbb E_{X\sim q_\mu}[\nabla_x\log p(X)]. \]

Now read the direction directly from the two-mode curve. Since \(p(x)>0\), the score \(\nabla_x\log p(x)=p'(x)/p(x)\) has the same sign as the slope \(p'(x)\):

Because the student is narrow compared with the distance between the modes, most of its samples lie in the same basin. Their average score therefore pulls the whole student toward one mode rather than leaving it between the two.

target \(p\) one-mode student \(q_\mu\). Forward KL compromises between the modes; reverse KL renders one mode well and drops the other.
Operational interpretation The reverse-KL update is computed at samples the student already produces. If the student almost never generates one mode, training almost never visits that mode either, so there is little evidence telling the student to recover it. Gaussian noise means recovery is still possible, but finite batches make it unlikely. This is sample starvation, and it is why initialization and coverage anchors matter.

3. Deriving the two-score identity

Fix one noise level \(t\). Write the reparameterized noisy sample as

\[ \mathbf x_t=T_\theta(\mathbf z,\boldsymbol\epsilon) =\alpha_tG_\theta(\mathbf z)+\sigma_t\boldsymbol\epsilon, \qquad \mathbf x_t\sim q_{\theta,t}. \]

The target \(p_t\) stays fixed while the student distribution \(q_{\theta,t}\) moves. The difficulty is that KL grades an entire density, whereas backpropagation moves individual generated samples. We first need an identity that connects these two views.

1

Convert density motion into particle motion

First distinguish two similar-looking symbols. In an integral, \(\mathbf x\) is only a dummy coordinate: it runs over all possible locations and does not itself have a distribution. The density \(q_{\theta,t}(\mathbf x)\) supplies the weight at each location. By contrast, \(\mathbf x_t\) is a random sample whose distribution is \(q_{\theta,t}\).

We obtain that random sample from base noise whose distribution does not depend on \(\theta\):

\[ \begin{aligned} \mathbf z&\sim\mathcal N(\mathbf0,\mathbf I), &\boldsymbol\epsilon&\sim\mathcal N(\mathbf0,\mathbf I),\\ \mathbf x_t &=T_\theta(\mathbf z,\boldsymbol\epsilon) =\alpha_tG_\theta(\mathbf z)+\sigma_t\boldsymbol\epsilon, &\mathbf x_t&\sim q_{\theta,t}. \end{aligned} \]

Now let \(g\) be any differentiable test function held fixed during the perturbation. Each equality below performs one separate conversion:

\[ \begin{aligned} \int \nabla_\theta q_{\theta,t}(\mathbf x)g(\mathbf x)\,d\mathbf x &=\nabla_\theta\int q_{\theta,t}(\mathbf x)g(\mathbf x)\,d\mathbf x\\ &=\nabla_\theta\, \mathbb E_{\mathbf x_t\sim q_{\theta,t}}[g(\mathbf x_t)]\\ &=\nabla_\theta\, \mathbb E_{\mathbf z,\boldsymbol\epsilon} [g(T_\theta(\mathbf z,\boldsymbol\epsilon))]\\ &=\mathbb E_{\mathbf z,\boldsymbol\epsilon}\!\left[ \nabla_{\mathbf x}g(\mathbf x_t)^\top \frac{\partial\mathbf x_t}{\partial\theta} \right]. \end{aligned} \] (2)

The first line moves the derivative through the integral. The second uses the definition of expectation: integrating \(g(\mathbf x)\) against the density \(q_{\theta,t}(\mathbf x)\) is the same as averaging \(g(\mathbf x_t)\) over samples \(\mathbf x_t\sim q_{\theta,t}\). The third is the actual reparameterization step: it replaces sampling from the moving distribution \(q_{\theta,t}\) with sampling fixed noises \((\mathbf z,\boldsymbol\epsilon)\) and passing them through \(T_\theta\). The final line is the ordinary chain rule.

\[ \frac{\partial\mathbf x_t}{\partial\theta} =\alpha_t\frac{\partial G_\theta(\mathbf z)}{\partial\theta}. \]

Here we used the fact that the sampled \(\mathbf z\), \(\boldsymbol\epsilon\), and the fixed-time coefficients \(\alpha_t,\sigma_t\) do not change when differentiating with respect to \(\theta\). Thus the left side of Equation (2) describes how the whole density moves, while the right side describes how its generated particles move.

2

Compute the first variation of reverse KL

Temporarily write \(q=q_{\theta,t}\) and \(p=p_t\). Perturb the density to \(q+\eta\,\delta q\):

\[ \begin{aligned} \frac{d}{d\eta}\mathrm{KL}(q+\eta\delta q\|p)\Big|_{\eta=0} &=\frac{d}{d\eta} \int(q+\eta\delta q)\log\frac{q+\eta\delta q}{p}\,d\mathbf x \Big|_{\eta=0}\\ &=\int\left[ \delta q\log\frac{q+\eta\delta q}{p} +(q+\eta\delta q)\frac{\delta q}{q+\eta\delta q} \right]_{\eta=0}d\mathbf x\\ &=\int\delta q(\mathbf x) \left(\log\frac{q(\mathbf x)}{p(\mathbf x)}+1\right)d\mathbf x. \end{aligned} \]

This is ordinary product rule, applied to a function whose input is now a density. Differentiating the leading \(q\) gives \(\delta q\log(q/p)\); differentiating the \(q\) inside the logarithm gives one more \(\delta q\). Hence

\[ \boxed{ \frac{\delta}{\delta q}\mathrm{KL}(q\|p) =\log\frac qp+1.} \] (3)
3

Let the generator produce the perturbation

This step uses the functional version of the chain rule. For an ordinary function \(f(u)\), where the intermediate variable \(u\) depends on \(\theta\), the usual chain rule is

\[ \frac{d}{d\theta}f(u(\theta)) =\frac{df}{du}\frac{du}{d\theta}. \]

A density is different only because the intermediate object is no longer one number. It contains one value \(q(\mathbf x)\) at every location \(\mathbf x\). If space were a finite grid, the multivariable chain rule would sum the contribution from every grid point:

\[ \nabla_\theta F(q_1(\theta),\ldots,q_n(\theta)) =\sum_{i=1}^n \frac{\partial F}{\partial q_i} \nabla_\theta q_i. \]

For a continuous density, the sum becomes an integral and the partial derivative becomes a functional derivative:

\[ \boxed{ \nabla_\theta F[q_{\theta,t}] =\int \frac{\delta F}{\delta q}(\mathbf x) \nabla_\theta q_{\theta,t}(\mathbf x) \,d\mathbf x.} \]

The quantity \(\delta F/\delta q(\mathbf x)\) asks how sensitive \(F\) is to a small density change at location \(\mathbf x\). The integral adds the effects from all locations. Thus this is the same chain rule as before, with an integral replacing the finite sum.

For reverse KL, Equation (3) supplies the functional derivative. Evaluate it at the current density, then apply Equation (2) with \(g=\log(q_{\theta,t}/p_t)+1\):

\[ \begin{aligned} \nabla_\theta\mathrm{KL}(q_{\theta,t}\|p_t) &=\int \underbrace{\frac{\delta\,\mathrm{KL}}{\delta q}(\mathbf x)}_{ \log(q_{\theta,t}(\mathbf x)/p_t(\mathbf x))+1} \nabla_\theta q_{\theta,t}(\mathbf x)\,d\mathbf x\\ &=\int\nabla_\theta q_{\theta,t}(\mathbf x) \left(\log\frac{q_{\theta,t}(\mathbf x)}{p_t(\mathbf x)}+1\right)d\mathbf x\\ &=\mathbb E\!\left[ \nabla_{\mathbf x}\left( \log\frac{q_{\theta,t}(\mathbf x_t)}{p_t(\mathbf x_t)}+1 \right)^\top \frac{\partial\mathbf x_t}{\partial\theta} \right]. \end{aligned} \]

The spatial gradient kills the constant \(1\) and splits the log ratio into two scores:

\[ \boxed{ \nabla_\theta\mathrm{KL}(q_{\theta,t}\|p_t) =\mathbb E\!\left[ \bigl( \underbrace{\nabla\log q_{\theta,t}(\mathbf x_t)}_{\text{fake score}} -\underbrace{\nabla\log p_t(\mathbf x_t)}_{\text{real score}} \bigr)^\top \frac{\partial\mathbf x_t}{\partial\theta} \right].} \] (4)
Why we hold the score field fixed during this step The first variation already tells us how the objective changes when the density changes. We then evaluate that answer at the current density and ask how the generator moves its samples. Differentiating the first variation again would ask a second-order question; it is not a missing first-order gradient.
4

Translate scores into denoiser predictions

For Gaussian corruption, noise prediction and score prediction are related by

\[ \nabla\log p_t(\mathbf x_t)=-\frac{\boldsymbol\epsilon_\varphi(\mathbf x_t,t)}{\sigma_t}, \qquad \nabla\log q_{\theta,t}(\mathbf x_t)=-\frac{\boldsymbol\epsilon_\psi(\mathbf x_t,t)}{\sigma_t}. \]

Putting the time expectation back gives the familiar computational form:

\[ \boxed{ \nabla_\theta\mathcal L =\mathbb E_{t,\mathbf z,\boldsymbol\epsilon}\!\left[ \frac{\gamma(t)}{\sigma_t} \bigl(\boldsymbol\epsilon_\varphi-\boldsymbol\epsilon_\psi\bigr)^\top \frac{\partial\mathbf x_t}{\partial\theta} \right].} \] (5)

This equation is not the starting assumption. It is the end of the density-to-particle calculation.

4. General \(f\)-divergence: same direction, different weight

The previous result is not special to KL. Many common distribution distances belong to the \(f\)-divergence family. The useful surprise is that they all point in the same teacher-minus-student score direction; they mainly disagree about how strongly to weight each location.

\[ D_f(p_t\|q_{\theta,t}) =\int q_{\theta,t}(\mathbf x) f\!\left(\frac{p_t(\mathbf x)}{q_{\theta,t}(\mathbf x)}\right)d\mathbf x, \qquad r_t(\mathbf x)=\frac{p_t(\mathbf x)}{q_{\theta,t}(\mathbf x)}. \]

At one spatial point, hold \(p\) fixed and differentiate \(qf(p/q)\) with respect to \(q\):

\[ \frac{d}{dq}\left[qf\!\left(\frac pq\right)\right] =f(r)+qf'(r)\left(-\frac{p}{q^2}\right) =f(r)-rf'(r). \]

Therefore the first variation is \(f(r)-rf'(r)\). Its spatial gradient is

\[ \begin{aligned} \nabla\bigl[f(r)-rf'(r)\bigr] &=-rf''(r)\nabla r\\ &=-f''(r)r^2 \bigl(\nabla\log p-\nabla\log q\bigr), \end{aligned} \]

where \(\nabla r=r(\nabla\log p-\nabla\log q)\). Applying the same density-to-particle identity gives

\[ \boxed{ \nabla_\theta D_f(p_t\|q_{\theta,t}) =-\mathbb E\!\left[ \underbrace{h(r_t)}_{f''(r_t)r_t^2} \bigl(\nabla\log p_t-\nabla\log q_{\theta,t}\bigr)^\top \frac{\partial\mathbf x_t}{\partial\theta} \right].} \] (6)

So every smooth \(f\)-divergence uses the same teacher-minus-student score direction. The scalar \(h(r)\) decides where training effort goes. If \(r=p/q\) is large, the target has much more probability than the student at that location.

Divergence\(f(r)\)\(h(r)\)Effect
Reverse KL\(-\log r\)\(1\)No density-ratio estimate; mode-seeking
Jensen–Shannon\(r\log r-(r+1)\log\frac{r+1}{2}\)\(\frac{r}{r+1}\)Bounded, saturating coverage weight
Squared Hellinger\(1-\sqrt r\)\(\frac14\sqrt r\)Intermediate emphasis
Forward KL\(r\log r\)\(r\)Amplifies teacher-heavy, student-light regions

Forward KL is coverage-seeking because \(r=p/q\) becomes large precisely where the teacher has mass and the student does not. But that same factor can have high variance. Any nonconstant \(h(r)\) also requires a density-ratio estimator—often a discriminator—in addition to the fake-score estimator. Reverse KL is computationally attractive partly because \(h\equiv1\) makes the ratio disappear.

5. The fake score is the student's entropy gradient

Split reverse KL into negative entropy and cross-entropy:

\[ \mathrm{KL}(q\|p) =\underbrace{\mathbb E_q[\log q]}_{-H(q)} +\underbrace{\mathbb E_q[-\log p]}_{\text{cross-entropy}}. \]

Applying the same first-variation argument to the two pieces separately gives

\[ \nabla_\theta[-H(q_{\theta,t})] =\mathbb E\!\left[ \nabla\log q_{\theta,t}(\mathbf x_t)^\top \frac{\partial\mathbf x_t}{\partial\theta} \right], \] \[ \nabla_\theta\mathbb E_{q_{\theta,t}}[-\log p_t] =-\mathbb E\!\left[ \nabla\log p_t(\mathbf x_t)^\top \frac{\partial\mathbf x_t}{\partial\theta} \right]. \]

In plain language, the real score says, “move toward places the target considers likely.” The fake score says, “do not let all student samples pile into the same place.” It is not merely a variance-reduction baseline: it is the gradient of the student's negative entropy. If it is missing or badly estimated, attraction remains but the pressure that spreads samples out becomes too weak.

A one-dimensional Gaussian check Let \(p=\mathcal N(0,s^2)\) and \(q_u=\mathcal N(0,u^2)\). Then \[ \mathrm{KL}(q_u\|p) =\tfrac12\left(\frac{u^2}{s^2}-1-\log\frac{u^2}{s^2}\right), \qquad \frac{d}{du}\mathrm{KL}(q_u\|p) =\frac{u}{s^2}-\frac1u. \] The first term is target attraction; the second is entropy spreading. If \(u>s\), the derivative is positive and gradient descent contracts the student. If \(u<s\), it is negative and gradient descent expands it. The two scores cancel exactly at \(u=s\).

What Part I establishes

The complete chain is now visible:

\[ \text{density functional} \longrightarrow \text{first variation} \longrightarrow \text{spatial score field} \longrightarrow \text{generator Jacobian}. \]

What remains is estimation. The teacher score is available from a pretrained diffusion model; the student marginal score is not. SDS, VSD, Diff-Instruct, DMD, and adversarial variants differ primarily in what they substitute or learn for that missing quantity, what weight they apply, and at which state the comparison is formed.

References

[1] Yin et al. “One-step Diffusion with Distribution Matching Distillation.” CVPR 2024.

[2] Xu, Nie, and Vahdat. “One-step Diffusion Models with \(f\)-Divergence Distribution Matching.” 2025.

[3] Wang et al. “ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation.” NeurIPS 2023.