Diffusion Models Revisited (25): Identity-Based Distillation (Part 1)
Today we're going to discuss the paper Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation. As the title suggests, this is a new paper exploring how to distill diffusion models faster and better.
Even without having done distillation before, you can probably guess the standard procedure: randomly sample a large number of inputs, use the diffusion model to generate corresponding outputs, and then use these input-output pairs as training data to supervise the training of a new model. However, as we all know, the original diffusion model acting as the teacher typically needs many steps (e.g. 1000) of iteration to produce high-quality output, so regardless of the training details, one obvious drawback of this approach is that generating the training data is extremely time-consuming and costly. In addition, the distilled student model usually incurs some loss in quality, to a greater or lesser extent.
Is there a way to solve both drawbacks at once? This is exactly the problem the above paper sets out to address. more
Back Again
The paper calls its proposed scheme "Score identity Distillation (SiD)". This name comes from the fact that the entire framework is designed and derived based on several identities, and this somewhat casual name is meant to highlight the key role that identity transformations play in SiD — which is indeed SiD's core contribution.
As for the training idea behind SiD, it is actually almost identical to the paper Learning Generative Models using Denoising Density Estimators (abbreviated "DDE"), which we introduced earlier in From Denoising Autoencoders to Generative Models — even the final form bears a fair resemblance. The only difference is that back then diffusion models had not yet risen to prominence, so DDE was proposed as a new type of generative model, and at the time it was actually quite niche. Now that diffusion models are popular, it can be reformulated as a diffusion model distillation method, since it requires a trained denoising autoencoder — which is exactly the core component of a diffusion model.
Below, I'll introduce SiD in my own way. Suppose we have a teacher diffusion model $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$ trained on the target dataset, which requires multiple sampling steps to generate high-quality images. Our goal is to train a student model $\boldsymbol{x} = \boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ that can sample in a single step — that is, a GAN-like generator that, given a specified noise input $\boldsymbol{z}$, can directly produce an image meeting the requirements. If we had a large number of $(\boldsymbol{z},\boldsymbol{x})$ pairs, we could just do direct supervised training (of course, the loss function and other details would still need to be worked out, and readers can refer to related work for that). But what if we don't have such pairs? It's certainly not impossible to train without them — after all, GANs are trained without $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$ pairs too — so the key question is how to leverage an already-trained diffusion model to provide a better training signal.
Both SiD and its predecessor DDE adopt an idea that seems roundabout at first glance but is actually quite clever:
If the data distribution produced by $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ closely resembles the target distribution, then training a diffusion model $\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t,t)$ on the dataset generated by $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ should also result in something that closely resembles $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$, shouldn't it?
The Basic Form
The clever part of this idea is that it sidesteps the need for samples generated by the teacher model, and also avoids needing the teacher model's real training samples, because "training a diffusion model on the dataset generated by $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$" only requires data generated by the student model $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ (for short, "student data"), and since $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ is a single-step model, generating data with it is relatively cheap timewise.
Of course, this is still just an idea — there's still some way to go before turning it into a workable training scheme. First let's recall how diffusion models work. Following the formulation from Diffusion Models Revisited (3): DDPM = Bayes + Denoising, we add noise to the input $\boldsymbol{x}_0$ as follows:
\begin{equation}\boldsymbol{x}_t = \bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\end{equation}
In other words, $p(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_t\boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I})$. The training objective for $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$ is then denoising:
\begin{equation}\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}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon},t) - \boldsymbol{\varepsilon}\Vert^2\right] \label{eq:d-real-data}\end{equation}
Here $\tilde{p}(\boldsymbol{x}_0)$ is the teacher model's training data. Likewise, if we want to train a diffusion model on the student data generated by $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$, the training objective becomes:
\begin{equation}\begin{aligned} \boldsymbol{\psi}^* =&\, \mathop{\text{argmin}}_{\boldsymbol{\psi}} \mathbb{E}_{\boldsymbol{x}_0^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)}),\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] \\ =&\, \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] \end{aligned}\label{eq:dloss}\end{equation}
where $\boldsymbol{x}_t^{(g)}=\bar{\alpha}_t\boldsymbol{x}_0^{(g)} + \bar{\beta}_t\boldsymbol{\varepsilon}=\bar{\alpha}_t\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}) + \bar{\beta}_t\boldsymbol{\varepsilon}$, i.e. samples obtained by adding noise to the student data, whose distribution we denote as $p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)})$. The second equality uses the fact that "$\boldsymbol{x}_0^{(g)}$ is fully determined by $\boldsymbol{z}$," so the expectation over $\boldsymbol{x}_0^{(g)}$ is equivalent to an expectation over $\boldsymbol{z}$. Now we have two diffusion models, and the discrepancy between them, to some extent, measures the difference between the data distributions generated by the teacher and student models. So an intuitive idea is to learn the student model by minimizing this discrepancy:
\begin{equation} \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{equation}
Note that the optimization in $\eqref{eq:dloss}$ depends on $\boldsymbol{\theta}$, so when $\boldsymbol{\theta}$ changes via $\eqref{eq:gloss-1}$, the value of $\boldsymbol{\psi}^*$ changes accordingly. As a result, $\eqref{eq:dloss}$ and $\eqref{eq:gloss-1}$ actually need to be optimized in alternation, similar to a GAN.
The Crucial Touch
Speaking of GANs, some readers might already be wincing, since GANs are notoriously prone to collapsing during training. Unfortunately, the alternating training scheme built from $\eqref{eq:dloss}$ and $\eqref{eq:gloss-1}$ proposed above suffers from the same issue. In principle there's nothing wrong with it theoretically — the problem lies in the gap between theory and practice, which manifests in two ways:
1. In theory we should first find the optimal solution to $\eqref{eq:dloss}$ before optimizing $\eqref{eq:gloss-1}$, but in practice, for the sake of training cost, we optimize $\eqref{eq:gloss-1}$ before it has been trained to optimality;
2. In theory $\boldsymbol{\psi}^*$ varies with $\boldsymbol{\theta}$, so it should really be written as $\boldsymbol{\psi}^*(\boldsymbol{\theta})$, which means optimizing $\eqref{eq:gloss-1}$ should include an extra term for the gradient of $\boldsymbol{\psi}^*(\boldsymbol{\theta})$ with respect to $\boldsymbol{\theta}$ — but in practice, when optimizing $\eqref{eq:gloss-1}$, we always treat $\boldsymbol{\psi}^*$ as a constant.
These two issues are quite fundamental — they are also the root cause of unstable GAN training. An earlier paper, Revisiting GANs by Best-Response Constraint: Perspective, Methodology, and Application, specifically improved GAN training by addressing point 2. It might seem like neither issue can really be resolved: for the first, we can basically never afford to always train $\boldsymbol{\psi}$ to optimality — the cost is simply prohibitive; for the second, in an alternating training setup we have no good way to obtain any useful information about $\boldsymbol{\psi}^*(\boldsymbol{\theta})$, let alone its gradient with respect to $\boldsymbol{\theta}$.
Fortunately, for the diffusion model distillation problem at hand, SiD proposes a scheme that effectively mitigates both issues. The idea behind SiD is remarkably "simple": since we cannot avoid using an approximate value for $\boldsymbol{\psi}^*$ and treating $\boldsymbol{\psi}^*$ as a constant, the only remaining option is to use identity transformations to eliminate, as much as possible, the dependence of the objective $\eqref{eq:gloss-1}$ on $\boldsymbol{\psi}^*$. As long as the dependence of $\eqref{eq:gloss-1}$ on $\boldsymbol{\psi}^*$ is weak enough, the negative impact of the two issues above will also be weak enough.
This is SiD's core contribution, and the truly brilliant, "crucial touch" of the whole method.
The Identity Transformation
Now let's look concretely at what identity transformation is applied. Consider $\eqref{eq:d-real-data}$ first — its optimization objective can be equivalently rewritten as:
\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[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon},t) - \boldsymbol{\varepsilon}\Vert^2\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{x}_t\sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)}\left[\left\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t) - \frac{\boldsymbol{x}_t - \bar{\alpha}_t \boldsymbol{x}_0}{\bar{\beta}_t}\right\Vert^2\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{x}_t\sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)}\left[\left\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t) + \bar{\beta}_t\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right\Vert^2\right] \end{aligned}\end{equation}
By the score matching results from Diffusion Models Revisited (5): The General SDE Framework, the optimal solution of the above objective is $\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 $\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)})$. The objective of $\eqref{eq:gloss-1}$ can then be equivalently rewritten as:
\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] \\[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) + \bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})\right\rangle\right] \\[5pt] =&\, \color{green}{\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)\right\rangle\right]} \\[5pt] &\,\qquad+ \color{red}{\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \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)})\right\rangle\right]} \end{aligned}\end{equation}
Next we use an identity proven earlier in Diffusion Models Revisited (18): Score Matching = Conditional Score Matching to simplify the red part above:
\begin{equation}\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t) = \mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right] \label{eq:id}\end{equation}
This is an identity derived from the definition of probability density and Bayes' rule, and it holds regardless of the form of $p(\boldsymbol{x}_t),p(\boldsymbol{x}_t|\boldsymbol{x}_0),p(\boldsymbol{x}_0|\boldsymbol{x}_t)$. Substituting this identity into the red part gives us:
\begin{equation}\color{red}{\begin{aligned} &\,\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \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)})\right\rangle\right] \\[5pt] = &\,\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}),\boldsymbol{x}_0^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)}|\boldsymbol{x}_t^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}} \log p(\boldsymbol{x}_t^{(g)}|\boldsymbol{x}_0^{(g)})\right\rangle\right] \\[5pt] = &\,-\mathbb{E}_{\boldsymbol{x}_0^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)}),\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}|\boldsymbol{x}_0^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \frac{\boldsymbol{x}_t - \bar{\alpha}_t \boldsymbol{x}_0}{\bar{\beta}_t}\right\rangle\right] \\[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{\varepsilon}\right\rangle\right] \end{aligned}}\end{equation}
Combining this with the green part yields the new loss function for the student model:
\begin{equation}\mathcal{L}_2 = \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]\label{eq:gloss-2}\end{equation}
This is SiD's core result. The paper's experiments show that this achieves efficient distillation, whereas $\eqref{eq:gloss-1}$ fails to produce meaningful results in training.
Compared with $\eqref{eq:gloss-1}$, the expression $\eqref{eq:gloss-2}$ above clearly contains $\boldsymbol{\psi}^*$ fewer times — that is, it depends more weakly on $\boldsymbol{\psi}^*$. Moreover, since the above expression is derived from an identity transformation based on the optimal solution $\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)})$, it effectively (partially) peeks in advance at the exact value of $\boldsymbol{\psi}^*$, which is one of the reasons for its superiority.
Other Details
Up to this point, the derivation in this article is essentially a repetition of the original paper's derivation. But aside from a few notational inconsistencies, there are some other differences in detail worth clarifying, to avoid confusing readers.
First, the paper's derivation assumes $\bar{\alpha}_t=1$ by default, following the setting used in Elucidating the Design Space of Diffusion-Based Generative Models. However, although $\bar{\alpha}_t=1$ is quite representative and simplifies the notation, it doesn't cover all types of diffusion models very well, so this article's derivation retains $\bar{\alpha}_t$. Second, the paper presents its results using $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t) = \frac{\boldsymbol{x}_t - \bar{\beta}_t \boldsymbol{\epsilon}(\boldsymbol{x}_t,t)}{\bar{\alpha}_t}$ as the standard, which is clearly at odds with the more common convention of using $\boldsymbol{\epsilon}(\boldsymbol{x}_t,t)$ in diffusion models — I have not yet fully grasped what advantage the original paper's notation offers.
Finally, the original paper found that the loss function $\mathcal{L}_1$, i.e. $\eqref{eq:gloss-1}$, is far too unstable and often ends up hurting performance, so SiD ultimately takes the negative of $\eqref{eq:gloss-1}$ as an additional loss term, weighted and added to the improved loss $\eqref{eq:gloss-2}$, giving a final loss of $\mathcal{L}_2 - \lambda\mathcal{L}_1$ (note: in the original paper the weighting symbol was $\alpha$, but in this article $\alpha$ has already been used to denote the noise schedule, so we use $\lambda$ instead). This can achieve better distillation results in some cases. For the specific experimental details and data, readers are encouraged to consult the original paper.
Compared with other distillation methods, one drawback of SiD is its relatively heavy GPU memory requirement, since it needs to maintain three models simultaneously — $\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t)$, $\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t,t)$, and $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ — all of the same size. Although they are not backpropagated through simultaneously, the combined footprint still roughly doubles the total memory usage. To address this, at the end of the paper SiD proposes that in the future one could try adding LoRA to the pretrained model to serve as the two additionally introduced models, further reducing memory requirements.
Further Thoughts
I believe that many readers with a solid theoretical background who have thought deeply about this topic would have a good chance of arriving at the initial "basic form" — i.e., the alternating optimization of $\eqref{eq:dloss}$ and $\eqref{eq:gloss-1}$ — on their own, especially given that DDE had already paved the way; it doesn't seem too hard to anticipate. But what makes SiD brilliant is that it doesn't stop there — it goes on to propose the subsequent identity transformation, making training more stable and efficient, which reflects the authors' deep understanding of both diffusion models and optimization theory.
At the same time, SiD leaves quite a few questions worth further thought and exploration. For instance, has the identity-based simplification of the student model's loss $\eqref{eq:gloss-2}$ been pushed to its limit? Not quite, because there is still a $\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)$ on the left side of the inner product, which can be simplified in the same way. Specifically, we have:
\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] \\[5pt] =&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2 - 2\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\right\rangle + \Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[5pt] =&\,\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})}\left[ \scriptsize{\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2 - 2\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})\right\rangle + \left\langle\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)})\right\rangle} \right] \\[5pt] \end{aligned}\end{equation}
Here, each $-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})$ can be transformed via the same identity $\eqref{eq:id}$ into a single $\boldsymbol{\varepsilon}$ (though note that only one $\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2=\langle\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle$ can be transformed this way, not both), while $\eqref{eq:gloss-2}$ corresponds to converting only part of it. Would it be better to convert all of them? Since there are no experimental results on this, it's hard to say for now. But there is one particularly interesting form: if we only transform the middle part above, the loss function can be written as:
\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] \\[5pt] =&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2 - 2\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\right\rangle + \Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[5pt] =&\,\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{\varepsilon}\Vert^2 + \Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] + \text{const} \\[5pt] \end{aligned}\label{eq:gloss-3}\end{equation}
This is the loss for the student model, i.e. the generator. Now let's compare it against the loss for the denoising model trained on student data, $\eqref{eq:dloss}$:
\begin{equation}\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-1}\end{equation}
Looking at these two expressions together, we can see that the student model is in fact trying to align itself with the teacher model while trying to move away from the denoising model trained on student data — a form reminiscent of LSGAN, where $\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t^{(g)},t)$ plays a role similar to a GAN discriminator. The difference is that in a GAN, the discriminator's loss is usually the sum of two terms while the generator's loss is a single term; in SiD it's the other way around. This actually reflects two different learning philosophies:
1. GAN: at the start, both the forger (generator) and the appraiser (discriminator) are novices. The appraiser continuously compares genuine and fake items to sharpen its appraisal skills, while the forger continuously improves its forgery skills based on feedback from the appraiser.
2. SiD: there is no genuine item at all, but there is an absolutely authoritative master appraiser (the teacher model). The forger (student model) keeps producing fakes while also training its own appraiser (the denoising model trained on student data), and then improves its forgery skills through the exchange between its own appraiser and the master.
Some readers may ask: why doesn't the forger in SiD consult the master directly, instead of relying on its own appraiser for indirect feedback? The reason is that direct communication with the master could lead to a problem where the forger keeps discussing the techniques for just one particular piece of work over and over, ultimately ending up producing only one type of convincing fake (mode collapse). Cultivating its own appraiser can help avoid this to some extent, because the forger's learning strategy becomes "get more approval from the master while minimizing approval from my own appraiser." If the forger keeps producing only one type of fake, then approval from both the master and its own appraiser will keep increasing, which contradicts the forger's learning strategy — forcing the forger to keep developing new products rather than settling into a rut.
Moreover, readers may notice that the entire SiD training process makes no use of any information from the recursive sampling procedure of diffusion models. In other words, it purely exploits the denoising model trained via the denoising objective. This raises a natural question: if our sole goal is to train a single-step generative model — rather than distill an existing diffusion model — would it be better to train a denoising model with just a single fixed noise level? For instance, as in DDE, fix $\bar{\alpha}_t=1$ and $\bar{\beta}_t=\beta=\text{some const}$ to train a denoising model, and then use it to repeat the SiD training procedure — would this simplify the training difficulty and improve training efficiency? This is also a question worth further investigation.
Summary
In this article, we introduced a new scheme for distilling diffusion models into single-step generative models. Its underlying idea traces back to work from a couple of years ago that used denoising autoencoders to train generative models. It requires neither access to the teacher model's real training set nor iterative sampling from the teacher model to generate paired data, and instead introduces GAN-like alternating training, along with a key identity transformation that stabilizes the training process. The overall method has a great deal worth learning from.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.