渐进蒸馏和v-prediction

渐进蒸馏和v-prediction

TL;DR：比较早期的用蒸馏的思想来做扩散模型采样加速的方法，通过渐进地对预训练的扩散模型进行蒸馏，学生模型一步学习教师模型两步的去噪结果，不断降低采样步数。并提出一种新的参数化形式 $\mathbf{v}$-prediction 来解决渐进蒸馏过程中信噪比太低时误差影响较大的问题。

渐进蒸馏

在一开始，我们有一个预训练的原始扩散模型作为初始教师模型。我们首先将学生模型初始化为一个结构、参数都与教师模型一模一样的扩散模型。然后，不断采样干净图像数据，加噪声，训练学生模型的去噪能力。由于我们要进行蒸馏，所以这里学生模型的预测目标不是干净的图片 $\mathbf{x}$，而是要学生模型单步（DDIM）预测出教师模型两步（DDIM）的去噪结果 $\tilde{\mathbf{x}}$。

具体来说，我们这里考虑的是连续时间步 $t\in [0,1]$，目标步数（即学生模型的步数）为 $N$，从而步长是 $1/N$，在时刻 $t$ 是要去噪从 ${\mathbf{z}}_t$ 到 ${\mathbf{z}}_{t-1/N}$。这样教师模型的步数是 $2N$，每一步是从 ${\mathbf{z}}_t$ 到 ${\mathbf{z}}_{t-0.5/N}$。我们这里连续运行教师模型两步，即从 ${\mathbf{z}}_t$ 到 ${\mathbf{z}}_{t-0.5/N}$ 再到 ${\mathbf{z}}_{t-1/N}$，我们的学生模型训练目标就是要一步直接从 ${\mathbf{z}}_t$ 预测出教师模型的两步去噪的结果 ${\mathbf{z}}_{t-1/N}$。

在收敛之后，我们将当前的学生模型作为下一轮的噪声模型，再将自身进行拷贝重新初始化一个新的学生模型，重复上述步骤。循环往复，即可通过渐进蒸馏不断降低模型的采样步数。

下面是渐进蒸馏的算法流程，对比了标准的扩散模型训练流程，主要就是将模型的预测目标从上一步的加噪结果改换成了教师模型的两步去噪结果，并渐进式地迭代这一过程。

参数化形式和训练损失

自从 DDPM 以来，扩散模型的参数化形式一般都是 $ϵ$-prediction，即预测噪声，再根据噪声计算出数据 $\mathbf{x}$。相当于间接地预测 $\mathbf{x}$：${\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t)=\frac{1}{{\alpha }_t}({\mathbf{z}}_t-{\sigma }_t{\widehat{ϵ}}_{\theta }({\mathbf{z}}_t))$。

在常规的扩散模型训练以及渐进蒸馏训练的早期（步数还比较多时），噪声预测的参数化形式工作得很好。因为这时信噪比 ${\alpha }_t^2/{\sigma }_t^2$ 在一个比较宽的范围内。当随着渐进蒸馏的进行，步数越来越少，信噪比越来越低以至于接近于 0，此时 ${\alpha }_t$ 接近于 0。根据上式，${\alpha }_t$ 在间接预测 ${\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t)$ 公式的分母上，因此此时网络输出预测噪声 ${\widehat{ϵ}}_{\theta }({\mathbf{z}}_t)$ 都会噪声 $\mathbf{x}$ 的巨大变化，从而导致训练不稳定。并且渐进蒸馏后期步数较少，无法通过后面的步数进行修正。

最终，如果我们将模型蒸馏到只剩下一个采样步，那么模型的输入就只是纯噪声 $ϵ$，此时信噪比为零，即 ${\alpha }_t=0,{\sigma }_t=1$。在这种极端情况下，$ϵ$ 预测和 $\mathbf{x}$ 预测之间的联系完全中断：观测数据 $z_t=ϵ$ 不再包含 $\mathbf{x}$ 的信息，并且 $ϵ$ 的预测 ${\widehat{ϵ}}_{\theta }({\mathbf{z}}_t)$ 也无法再间接地预测 $\mathbf{x}$。在损失函数中，加权函数 $w({\lambda }_t)$ 在此时的权重也成了 0。

为了解决这一问题，作者尝试了直接预测 $\mathbf{x}$、同时分别预测 $\mathbf{x}$ 和 $ϵ$ 后合并出 $\hat{\mathbf{x}}$，还提出了一种新的参数化形式 $\mathbf{v}$-prediction：
$$\mathbf{v}\equiv {\alpha }_tϵ-{\sigma }_t\mathbf{x}$$
从而：
$$\hat{\mathbf{x}}={\alpha }_t{\mathbf{z}}_t-{\sigma }_t{\hat{\mathbf{v}}}_{\theta }({\mathbf{z}}_t)$$
实验显示，这三种方式在渐进蒸馏训练中都表现得不错，并在在常规扩散模型的训练中效果也很好。

下面对作者设计的 $\mathbf{v}$-prediction 进行推导：

DDPM 的加噪公式：
$${\mathbf{z}}_t={\alpha }_t\mathbf{x}+{\sigma }_tϵ$$
令 ${\varphi }_t=\arctan ({\sigma }_t/{\alpha }_t)$，则有 ${\alpha }_t=\cos (\varphi ),{\sigma }_t=\sin (\varphi )$，从而：
$${\mathbf{z}}_{\varphi }=\cos (\varphi )\mathbf{x}+\sin (\varphi )ϵ$$
定义 $z_{\varphi }$ 的 “速度” 为其关于 $\varphi$ 的导数：
$${\mathbf{v}}_{\varphi }\equiv \frac{d{\mathbf{z}}_{\varphi }}{d\varphi }=\frac{d\cos (\varphi )}{d\varphi }\mathbf{x}+\frac{d\sin \varphi }{d\varphi }ϵ=\sin (\varphi )\mathbf{x}-\cos (\varphi )ϵ$$
这里就是上面 $\mathbf{v}$ 的定义 $\mathbf{v}\equiv {\alpha }_tϵ-{\sigma }_t\mathbf{x}$。稍微进行变换，有：
$$\begin{array}{rl}\sin (\varphi )& =\cos (\varphi )ϵ-{\mathbf{v}}_{\varphi }\\ & =\frac{\cos (\varphi )}{\sin (\varphi )}(\mathbf{z}-\cos (\varphi )\mathbf{x})-{\mathbf{v}}_{\varphi }\\ {\sin }^2(\varphi )\mathbf{x}& =\cos (\varphi )\mathbf{z}-{\cos }^2(\varphi )\mathbf{x}-\sin (\varphi ){\mathbf{v}}_{\varphi }\\ {\sin }^2(\varphi )\mathbf{x}+{\cos }^2(\varphi )\mathbf{x}& =\cos (\varphi )\mathbf{z}-\sin (\varphi ){\mathbf{v}}_{\varphi }\\ \mathbf{x}& =\cos (\varphi )\mathbf{z}-\sin (\varphi ){\mathbf{v}}_{\varphi }\end{array}$$
这里就是上面的第二个公式 $\hat{\mathbf{x}}={\alpha }_t{\mathbf{z}}_t-{\sigma }_t{\hat{\mathbf{v}}}_{\theta }({\mathbf{z}}_t)$。这个推导过程可以参考下图来理解。

总结

早期提出的渐进蒸馏是一种比较直觉的扩散模型步数蒸馏方法，其提出的 v-prediction 在后来也有广泛的应用。