Generative Diffusion Models Explained (26): Distillation Based on Identities (Part 2)
Let's return to our diffusion series. In Generative Diffusion Models Explained (25): Distillation Based on Identities (Part 1)], we introduced SiD (Score identity Distillation), a scheme for distilling diffusion models that requires neither real data nor sampling from the teacher model. It has a form similar to GAN, but with better training stability than GAN.
The core of SiD is to construct a better loss function for the student model through identity transformations — this is genuinely groundbreaking, but it also leaves some open questions. For instance, SiD's identity transformation of the loss function is incomplete; what would happen if we carried it out fully? How can we theoretically justify the necessity of the $\lambda$ that SiD introduces? Last month's paper Flow Generator Matching] (abbreviated FGM) successfully explained the choice of $\lambda=0.5$ from a more fundamental gradient perspective, and inspired by FGM, the author has further found an explanation for $\lambda = 1$.
In what follows we will go through these theoretical developments concerning SiD in detail. more
Recap of the Idea
From the previous post, we know that SiD's approach to distillation rests on the idea that "similar distributions produce similar trained denoising models." In formula form, this is
\begin{align} &\text{teacher diffusion model:}\quad\boldsymbol{\varphi}^* = \mathop{\text{argmin}}_{\boldsymbol{\varphi}} \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t) - \boldsymbol{\varepsilon}\Vert^2\right]\label{eq:tloss} \\[8pt] &\text{student diffusion model:}\quad\boldsymbol{\psi}^* = \mathop{\text{argmin}}_{\boldsymbol{\psi}} \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\varepsilon}\Vert^2\right]\label{eq:dloss}\\[8pt] &\text{student generative model:}\quad\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \underbrace{\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right]}_{\mathcal{L}_1}\label{eq:gloss-1} \end{align}
There's quite a bit of notation here, so let's go through it one by one. The first loss function is the training objective of the diffusion model we want to distill, where $\boldsymbol{x}_t = \bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}$ denotes the noised sample, $\bar{\alpha}_t,\bar{\beta}_t$ is the noise schedule, and $\boldsymbol{x}_0$ is the training sample. The second loss function is for the diffusion model trained on data generated by the student model, where $\boldsymbol{x}_t^{(g)}=\bar{\alpha}_t\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}) + \bar{\beta}_t\boldsymbol{\varepsilon}$, and $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ here denotes a sample generated by the student model, also written as $\boldsymbol{x}_0^{(g)}$. The third loss function attempts to train the student generative model (the generator) by narrowing the gap between the diffusion models trained on real data and on student-generated data.
Here the teacher model can be pretrained in advance, and training the two student models only requires the teacher model itself, not the data used to train the teacher — so as a distillation method, SiD is data-free. The two student models are then trained alternately, GAN-style, progressively improving the quality of the generator's outputs. As far as the author is aware, this training idea first appeared in the paper Learning Generative Models using Denoising Density Estimators], which we also discussed in From Denoising Autoencoders to Generative Models].
However, despite looking unproblematic on paper, in practice the alternating training of equations $\eqref{eq:dloss}$ and $\eqref{eq:gloss-1}$ collapses very easily, to the point where it barely produces usable results. This is due to two gaps between theory and practice:
1. In theory, we should first solve equation $\eqref{eq:dloss}$ to optimality before optimizing equation $\eqref{eq:gloss-1}$, but in practice, for the sake of training cost, we optimize equation $\eqref{eq:gloss-1}$ without having trained it to optimality first.
2. In theory, $\boldsymbol{\psi}^*$ varies with $\boldsymbol{\theta}$ — that is, it should be written as $\boldsymbol{\psi}^*(\boldsymbol{\theta})$ — so when optimizing equation $\eqref{eq:gloss-1}$ we should include an extra term for the gradient of $\boldsymbol{\psi}^*(\boldsymbol{\theta})$ with respect to $\boldsymbol{\theta}$; but in practice, when optimizing equation $\eqref{eq:gloss-1}$, we always treat $\boldsymbol{\psi}^*$ as a constant.
The first problem isn't really so bad, since as training proceeds, $\boldsymbol{\psi}$ can gradually approach the theoretically optimal $\boldsymbol{\psi}^*$. But the second problem is much more difficult and fundamental — one could say that GAN's training instability owes a good deal to this very issue. The core contribution of both SiD and FGM is precisely their attempt to address this second problem.
Identity Transformation
SiD's idea is to use an identity transformation to reduce the generator's loss function $\eqref{eq:gloss-1}$'s dependence on $\boldsymbol{\psi}^*$, thereby weakening the second problem. This idea is indeed groundbreaking, and quite a bit of subsequent work has built on SiD, including FGM, which we'll introduce below.
At the heart of the identity transformation is the following identity:
\begin{equation}\mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{f}(\boldsymbol{x}_t,t), \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)\right\rangle\right] = \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{f}(\boldsymbol{x}_t,t), \boldsymbol{\varepsilon}\right\rangle\right]\label{eq:id}\end{equation}
In short, $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$ can be replaced by $\boldsymbol{\varepsilon}$. Here $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$ is the theoretically optimal solution to equation $\eqref{eq:tloss}$, and $\boldsymbol{f}(\boldsymbol{x}_t,t)$ is an arbitrary vector function depending only on $\boldsymbol{x}_t$ and $t$. Note that "depending only on $\boldsymbol{x}_t$ and $t$" is a necessary condition for the identity to hold — as soon as $\boldsymbol{f}$ involves an independent $\boldsymbol{x}_0$ or $\boldsymbol{\varepsilon}$, the identity may no longer hold, so this must be checked carefully before applying it.
The previous post already gave a proof of this identity, but on reflection that proof was somewhat roundabout. Here is a more direct proof:
Proof:
Rewrite the objective $\eqref{eq:tloss}$ equivalently as
\begin{equation}\boldsymbol{\varphi}^* = \mathop{\text{argmin}}_{\boldsymbol{\varphi}} \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\Big[\mathbb{E}_{\boldsymbol{\varepsilon}\sim p(\boldsymbol{\varepsilon}|\boldsymbol{x}_t)}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t) - \boldsymbol{\varepsilon}\Vert^2\right]\Big]\end{equation}
By $\mathbb{E}[\boldsymbol{x}] = \mathop{\text{argmin}}\limits_{\boldsymbol{\mu}}\mathbb{E}_{\boldsymbol{x}}\left[\Vert \boldsymbol{\mu} - \boldsymbol{x}\Vert^2\right]$ (if unfamiliar, this can be verified by differentiation), we find that the theoretically optimal solution of the above is
\begin{equation}\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t) = \mathbb{E}_{\boldsymbol{\varepsilon}\sim p(\boldsymbol{\varepsilon}|\boldsymbol{x}_t)}[\boldsymbol{\varepsilon}]\end{equation}
so
\begin{equation}\begin{aligned} > \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{f}(\boldsymbol{x}_t,t), \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)\right\rangle\right]=&\, \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\left[\left\langle\boldsymbol{f}(\boldsymbol{x}_t,t), \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)\right\rangle\right] \\ > =&\, \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\left[\left\langle\boldsymbol{f}(\boldsymbol{x}_t,t), \mathbb{E}_{\boldsymbol{\varepsilon}\sim p(\boldsymbol{\varepsilon}|\boldsymbol{x}_t)}[\boldsymbol{\varepsilon}]\right\rangle\right] \\ > =&\, \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t),\boldsymbol{\varepsilon}\sim p(\boldsymbol{\varepsilon}|\boldsymbol{x}_t)}\left[\left\langle\boldsymbol{f}(\boldsymbol{x}_t,t), \boldsymbol{\varepsilon}\right\rangle\right] \\ > =&\, \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{f}(\boldsymbol{x}_t,t), \boldsymbol{\varepsilon}\right\rangle\right] > \end{aligned}\end{equation}
This completes the proof. The "necessary step" in the proof is the first equality, which relies on the condition that "$\boldsymbol{f}(\boldsymbol{x}_t,t)$ depends only on $\boldsymbol{x}_t$ and $t$."
The crux of identity $\eqref{eq:id}$ is the optimality of $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$, and since objective $\eqref{eq:tloss}$ has the same form as $\eqref{eq:dloss}$, the same conclusion also applies to $\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t,t)$. Using this, we can transform $\eqref{eq:gloss-1}$ into
\begin{equation}\begin{aligned} &\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[8pt] =&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\bigg[\Big\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \underbrace{\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)}_{\text{can be replaced with}\boldsymbol{\varepsilon}}\Big\rangle\bigg] \\[5pt] =&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\varepsilon}\right\rangle\right]\triangleq \mathcal{L}_2 \end{aligned}\label{eq:gloss-2}\end{equation}
The final form is exactly the generator loss function $\mathcal{L}_2$ proposed by SiD. It is the key to SiD's successful training — we can understand it as having, via an identity transformation, effectively pre-estimated the value of $\boldsymbol{\psi}^*$ while weakening its dependence on $\boldsymbol{\psi}^*$, so that training the generator with this loss works better than with $\mathcal{L}_1$.
SiD's remaining open questions are:
1. The identity transformation of $\mathcal{L}_2$ isn't complete: expanding $\mathcal{L}_2$, we find there's still a term $\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle]$, whose $\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)$ could likewise be replaced by $\boldsymbol{\varepsilon}$. The question is then whether this fuller transformation, given below, would be a better choice than $\mathcal{L}_2$:
\begin{equation}\mathcal{L}_3 = \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}\Vert^2 - 2\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle + \langle \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\boldsymbol{\varepsilon}\rangle\right]\label{eq:gloss-3}\end{equation}
2. In practice, the loss SiD actually uses is neither $\mathcal{L}_2$ nor $\mathcal{L}_1$, but rather $\mathcal{L}_2 - \lambda\mathcal{L}_1$, where $\lambda > 0$, and experiments find that the optimal value of $\lambda$ is around $1$, with some tasks even performing best at $\lambda=1.2$. This is quite puzzling, because $\mathcal{L}_1,\mathcal{L}_2$ are theoretically equal, so is $\lambda > 1$ effectively optimizing $\mathcal{L}_1$ in reverse? Wouldn't that run counter to the original motivation? Clearly this urgently calls for a theoretical explanation.
Facing the Gradient Head-On
Let's review once more: the fundamental difficulty we face is that, in theory, $\boldsymbol{\psi}^*$ is a function of $\boldsymbol{\theta}$, so when computing $\nabla_{\boldsymbol{\theta}} \mathcal{L}_1$ or $\nabla_{\boldsymbol{\theta}} \mathcal{L}_2$ we need to somehow compute $\nabla_{\boldsymbol{\theta}}\boldsymbol{\psi}^*$. But in practice, the best we can get is $\mathcal{L}_i^{\color{skyblue}{(\text{sg})}} \triangleq \mathcal{L}_i|_{\boldsymbol{\psi}^* \to \color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}$, where $\color{skyblue}{\text{sg}}$ denotes stop-gradient — i.e., we cannot obtain the gradient of $\boldsymbol{\psi}^*$ with respect to $\boldsymbol{\theta}$. So regardless of $\mathcal{L}_1,\mathcal{L}_2,\mathcal{L}_3$, their gradients as computed in practice are biased.
This is where FGM comes in, with an idea that gets closer to the essence of the matter: the losses $\mathcal{L}_1,\mathcal{L}_2,\mathcal{L}_3$ only ensure equality at the level of loss values, but what an optimizer actually needs is equality at the level of gradients. So we need to find a new loss function $\mathcal{L}_4$ satisfying
\begin{equation}\nabla_{\boldsymbol{\theta}}\mathcal{L}_4(\boldsymbol{\theta}, \color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]})= \nabla_{\boldsymbol{\theta}}\mathcal{L}_{1/2/3}(\boldsymbol{\theta}, \boldsymbol{\psi}^*)\end{equation}
that is, $\nabla_{\boldsymbol{\theta}}\mathcal{L}_4^{\color{skyblue}{(\text{sg})}} = \nabla_{\boldsymbol{\theta}}\mathcal{L}_{1/2/3}$, so that using $\mathcal{L}_4$ as the loss function achieves unbiased optimization.
FGM's derivation is likewise based on identity $\eqref{eq:id}$, though its original derivation is somewhat involved. For our purposes here, we can start directly from $\mathcal{L}_3$, i.e., equation $\eqref{eq:gloss-3}$, where the only term related to $\boldsymbol{\psi}^*$ that remains is $\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle]$. Let's directly compute its gradient, applying "identity transform first, then differentiate" and "differentiate first, then identity transform" respectively to the $\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2]$ operation, and compare the results.
Identity transform first, then differentiate:
\begin{equation}\begin{aligned} &\,\nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2] \\[5pt] =&\, \nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle] = \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle] \\[5pt] =&\, \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle] + \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]},t),\boldsymbol{\varepsilon}\rangle] \end{aligned}\label{eq:g-grad-1}\end{equation}
Differentiate first, then identity transform:
\begin{equation}\begin{aligned} &\,\nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2] \\[8pt] =&\, \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\nabla_{\boldsymbol{\theta}}\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2] = 2\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle\nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle] \\[8pt] =&\, 2\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle\nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle] + \underbrace{2\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle\nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]},t), \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle]}_{\text{can apply equation}\eqref{eq:id}} \\[5pt] =&\, 2\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle\nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle] + 2\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle\nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]},t), \boldsymbol{\varepsilon}\rangle] \end{aligned}\label{eq:g-grad-2}\end{equation}
Note the third equality carefully: only the term $\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]},t)$ can have identity $\eqref{eq:id}$ applied to it, because $\nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t)$'s $\boldsymbol{x}_t^{(g)}$ needs to be differentiated with respect to $\boldsymbol{\theta}$, and once we differentiate, the result is no longer necessarily a function of $\boldsymbol{x}_t^{(g)}$, so it no longer satisfies the condition for applying equation $\eqref{eq:id}$.
We now have two results for $\nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2]$. Taking twice equation $\eqref{eq:g-grad-1}$ and subtracting equation $\eqref{eq:g-grad-2}$ gives
\begin{equation}\begin{aligned} &\,\nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle] = \nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2] = \eqref{eq:g-grad-1}\times 2 - \eqref{eq:g-grad-2} \\[5pt] =&\,2 \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle] - 2\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle\nabla_{\boldsymbol{\theta}}\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle] \\[5pt] =&\,2 \nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle] - \nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\Vert\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t)\Vert^2] \\[5pt] =&\,\nabla_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[2\langle \boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle - \Vert\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t)\Vert^2] \end{aligned}\end{equation}
Notice, in the expression whose gradient is finally taken, that every occurrence of $\boldsymbol{\psi}^*$ has been given a stop-gradient $\color{skyblue}{\text{sg}}$ — meaning we no longer need to try to compute its gradient with respect to $\boldsymbol{\theta}$ — yet its gradient equals the exact gradient of $\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}[\langle \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle]$. So by substituting it for the corresponding term in $\mathcal{L}_3$, we obtain $\mathcal{L}_4$:
\begin{equation}\mathcal{L}_4^{\color{skyblue}{(\text{sg})}} = \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}\Vert^2 - 2\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle + 2\langle \boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\rangle - \Vert\boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right]\end{equation}
This is FGM's final result: it depends only on $\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}$, yet satisfies $\nabla_{\boldsymbol{\theta}}\mathcal{L}_4^{\color{skyblue}{(\text{sg})}}=\nabla_{\boldsymbol{\theta}}\mathcal{L}_{1/2/3}$. Looking more closely, we see that $\mathcal{L}_4^{\color{skyblue}{(\text{sg})}}=2\mathcal{L}_2^{\color{skyblue}{(\text{sg})}}-\mathcal{L}_1^{\color{skyblue}{(\text{sg})}}=2(\mathcal{L}_2^{\color{skyblue}{(\text{sg})}}-0.5\times \mathcal{L}_1^{\color{skyblue}{(\text{sg})}})$ holds, so FGM effectively affirms, from a gradient perspective, SiD's choice of $\lambda=0.5$.
Incidentally, the original FGM paper frames its description within the ODE-style diffusion framework (flow matching). But as the author noted in the previous post, neither SiD nor FGM actually relies on the iterative generation process of the diffusion model — they only use the denoising model that the diffusion model was trained to produce. So regardless of whether the framework is ODE, SDE, or DDPM, that's all surface-level; the denoising model is what truly matters. Hence this post can continue using the notation from the SiD post to introduce FGM.
Generalized Divergence
FGM has successfully worked out the most fundamental gradient, but this can only explain SiD's $\lambda=0.5$ — meaning that if we want to explain the viability of other values of $\lambda$, we need to modify our starting point. To this end, let's go back to basics and reconsider the generator's objective $\eqref{eq:gloss-1}$.
Readers familiar with diffusion models will know that the theoretically optimal solution of equation $\eqref{eq:tloss}$ can also be written as $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)=-\bar{\beta}_t\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t)$, and similarly the optimal solution of equation $\eqref{eq:dloss}$ is $\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)=-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})$, where $p(\boldsymbol{x}_t)$ and $p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})$ are respectively the noised distributions of the real data and of the generator's data. If this result is unfamiliar, see Generative Diffusion Models Explained (5): The General SDE Framework] and Generative Diffusion Models Explained (18): Score Matching = Conditional Score Matching] for background.
Substituting these two theoretically optimal solutions back into equation $\eqref{eq:gloss-1}$, we find that the generator is in fact trying to minimize a Fisher divergence:
\begin{equation}\begin{aligned} \mathcal{F}(p, p_{\boldsymbol{\theta}}) =&\, \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})} \left[\Vert \nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}) - \nabla_{\boldsymbol{x}_t^{(g)}}\log p(\boldsymbol{x}_t^{(g)})\Vert^2\right] \\ =&\, \int p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}) \left\Vert \nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}) - \nabla_{\boldsymbol{x}_t^{(g)}}\log p(\boldsymbol{x}_t^{(g)})\right\Vert^2 d\boldsymbol{x}_t^{(g)} \end{aligned}\end{equation}
What we want to reconsider is the reasonableness of the Fisher divergence, and where it might be improved. Notice that $p_{\boldsymbol{\theta}}$ appears twice in the Fisher divergence. Now let's ask the reader to consider a question: which of these two occurrences of $p_{\boldsymbol{\theta}}$ is more important?
The answer is the second one. To understand why, let's consider two scenarios: (1) fix the first occurrence of $p_{\boldsymbol{\theta}}$ and only optimize the second occurrence, $p_{\boldsymbol{\theta}}$; (2) fix the second occurrence of $p_{\boldsymbol{\theta}}$ and only optimize the first occurrence, $p_{\boldsymbol{\theta}}$. What differences arise in the results? In the first scenario, in all likelihood nothing much changes — we would still learn $p_{\boldsymbol{\theta}}=p$. In fact, since the Fisher divergence carries $\Vert\Vert^2$, the following, more general, conclusion is almost self-evidently true:
As long as $r(\boldsymbol{x})$ is a distribution that is nonzero everywhere, $p(\boldsymbol{x})=q(\boldsymbol{x})$ is still the theoretically optimal solution of the following generalized Fisher divergence:
\begin{equation}\mathcal{F}(p,q|r) = \int r(\boldsymbol{x}) \Vert \nabla_{\boldsymbol{x}} p(\boldsymbol{x}) - \nabla_{\boldsymbol{x}} q(\boldsymbol{x})\Vert^2 d\boldsymbol{x}\end{equation}
Put simply, the first occurrence of $p_{\boldsymbol{\theta}}$ isn't important at all — it could be swapped for any other distribution, and $\Vert\Vert^2$ alone still guarantees that the two distributions become equal. But the second scenario is entirely different: fixing the second occurrence of $p_{\boldsymbol{\theta}}$ and optimizing only the first occurrence, $p_{\boldsymbol{\theta}}$, gives the theoretically optimal solution
\begin{equation}p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}) = \delta(\boldsymbol{x}_t^{(g)} - \boldsymbol{x}_t^*),\quad \boldsymbol{x}_t^* = \mathop{\text{argmin}}_{\boldsymbol{x}_t^{(g)}} \,\left\Vert \nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}) - \nabla_{\boldsymbol{x}_t^{(g)}}\log p(\boldsymbol{x}_t^{(g)})\right\Vert^2\end{equation}
where $\delta$ is the Dirac delta distribution — meaning the model only needs to generate the single sample that minimizes $\Vert\Vert^2$ in order to minimize the loss. This, put plainly, is mode collapse! So the role of the first occurrence of $p_{\boldsymbol{\theta}}$ in the Fisher divergence isn't just secondary — it may even be actively harmful.
This suggests that, when training a model with a gradient-based optimizer, it's actually better to simply drop the gradient through the first occurrence of $p_{\boldsymbol{\theta}}$ altogether — i.e., a Fisher divergence of the following form is a better choice:
\begin{equation}\begin{aligned} \mathcal{F}^+(p, p_{\boldsymbol{\theta}}) =&\, \int p_{\color{skyblue}{\text{sg}[}\boldsymbol{\theta}\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)}) \left\Vert \nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}) - \nabla_{\boldsymbol{x}_t^{(g)}}\log p(\boldsymbol{x}_t^{(g)})\right\Vert^2 d\boldsymbol{x}_t^{(g)} \\[5pt] =&\, \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})} \left[\Vert \nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]}) - \nabla_{\boldsymbol{x}_t^{(g)}}\log p(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]})\Vert^2\right] \\[5pt] \propto&\, \underbrace{\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})} \left[\Vert \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\color{skyblue}{\text{sg}[}\boldsymbol{x}_t^{(g)}\color{skyblue}{]},t)\Vert^2\right]}_{\mathcal{L}_5} \end{aligned}\end{equation}
In other words, $\mathcal{L}_5$ here is very likely a better starting point than $\mathcal{L}_1$: numerically it equals $\mathcal{L}_1$, but with part of the gradient removed:
\begin{equation}\nabla_{\boldsymbol{\theta}}\mathcal{L}_5 = \nabla_{\boldsymbol{\theta}}\mathcal{L}_1 - \nabla_{\boldsymbol{\theta}}\underbrace{\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})} \left[\Vert \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\color{skyblue}{\text{sg}[}\boldsymbol{\psi}^*\color{skyblue}{]}}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right]}_{\text{exactly}\mathcal{L}_1^{\color{skyblue}{(\text{sg})}}}\end{equation}
where $\nabla_{\boldsymbol{\theta}}\mathcal{L}_1$ has already been computed by FGM, equal to $\nabla_{\boldsymbol{\theta}}(2\mathcal{L}_2^{\color{skyblue}{(\text{sg})}}-\mathcal{L}_1^{\color{skyblue}{(\text{sg})}})$. Therefore, starting from $\mathcal{L}_5$, the loss function we actually use in practice is $2\mathcal{L}_2^{\color{skyblue}{(\text{sg})}}-\mathcal{L}_1^{\color{skyblue}{(\text{sg})}}-\mathcal{L}_1^{\color{skyblue}{(\text{sg})}}=2(\mathcal{L}_2^{\color{skyblue}{(\text{sg})}}-\mathcal{L}_1^{\color{skyblue}{(\text{sg})}})$, which explains the choice of $\lambda=1$. As for choosing $\lambda$ slightly greater than 1, this is a somewhat more extreme move: it amounts to adding $-\mathcal{L}_1^{\color{skyblue}{(\text{sg})}}$ on top of $\mathcal{L}_5$ as an extra penalty term, further reducing the risk of mode collapse. Of course, here it really is meant purely as a penalty term, so its weight shouldn't be too large — according to SiD's experimental results, training already starts collapsing around $\lambda=1.5$.
Incidentally, an earlier work by the FGM authors, One-Step Diffusion Distillation through Score Implicit Matching], also proposed a similar move of changing the first occurrence of $p_{\boldsymbol{\theta}}$ to $p_{\color{skyblue}{\text{sg}[}\boldsymbol{\theta}\color{skyblue}{]}}$, but it did not explicitly discuss the rationale for this operation starting from the original form of the Fisher divergence, leaving the treatment somewhat incomplete.
Summary
This post has covered the subsequent theoretical developments of SiD (Score identity Distillation), mainly explaining, from a gradient perspective, the choice of the $\lambda$ parameter in SiD. The core piece is the clever approach discovered by FGM (Flow Generator Matching) for accurately estimating SiD's gradient, which affirms the choice of $\lambda=0.5$. Building on this, the author extended the concept of Fisher divergence, thereby explaining the choice of value for $\lambda=1$.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.