引言
今天我们分享的是一篇名为《Score Identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation》的新论文。该论文探讨了如何更快更好地蒸馏扩散模型。
友情链接：ACEJoy
即便没有做过蒸馏，大家可能也能猜到蒸馏的常规步骤：随机采样大量输入，然后用扩散模型生成相应结果作为输出，用这些输入输出作为训练数据对，来监督训练一个新模型。然而，众所周知，作为教师的原始扩散模型通常需要多步（比如1000步）迭代才能生成高质量输出，所以且不论中间训练细节如何，该方案的一个显著缺点是生成训练数据太费时费力。此外，蒸馏之后的学生模型通常或多或少都有效果损失。
有没有方法能一次性解决这两个缺点呢？这就是上述论文试图要解决的问题。
思路简介
论文将所提方案称为“Score Identity Distillation（SiD）”，基于几个恒等式来设计和推导了整个框架。实际上，它的设计思想与几个恒等式并没有直接联系，其次几个恒等式都是已知的公式而不是新的，所以这个名字显得相当随意。
接下来笔者用自己的思路去介绍SiD。
初级形式
假设我们有一个在目标数据集训练好的教师扩散模型 ( \epsilon_{\phi^}(x_t, t) )，它需要多步采样才能生成高质量图片。我们的目标是训练一个单步采样的学生模型 ( x = g_{\theta}(z) )，即一个类似GAN的生成器，输入指定噪声 ( z ) 就可以直接生成符合要求的图像。如果我们有很多的 ( (z, x) ) 对，那么直接监督训练就可以了（当然损失函数和其他细节还需要进一步确定，读者可以自行参考相关工作），但如果没有呢？肯定不是不能训，因为就算没有 ( \epsilon_{\phi^}(x_t, t) ) 也能训，比如GAN，所以关键是怎么借助已经训练好的扩散模型提供更好的信号。
SiD及前作DDE使用了一个看上去很绕但是也很聪明的思路：
如果 ( g_{\theta}(z) ) 产生的数据分布跟目标分布很相似，那么拿 ( g_{\theta}(z) ) 生成的数据集去训练一个扩散模型 ( \epsilon_{\psi^}(x_t, t) ) 的话，它也应该跟 ( \epsilon_{\phi^}(x_t, t) ) 很相似？
这个思路的聪明之处在于，它绕开了对教师模型生成样本的需求，也不需要训练教师模型的真实样本，因为“拿 ( g_{\theta}(z) ) 生成的数据集去训练一个扩散模型”只需要学生模型 ( g_{\theta}(z) ) 生成的数据（简称“学生数据”），而 ( g_{\theta}(z) ) 是一个单步模型，用它来生成数据时间上比较友好。
当然，这还只是思路，将其转换为实际可行的训练方案还有一段路要走。
方法与公式
扩散模型回顾
我们采用《生成扩散模型漫谈（三）：DDPM = 贝叶斯 + 去噪》的形式，对输入 ( x_0 ) 进行加噪：
[ x_t = \bar{\alpha}_t x_0 + \bar{\beta}_t \epsilon, \quad \epsilon \sim N(0, I) ]
训练 ( \epsilon_{\phi^}(x_t, t) ) 的方式则是去噪： [ \phi^ = \arg\min_{\phi} E_{x_0 \sim \tilde{p}(x_0), \epsilon \sim N(0, I)} \left[ | \epsilon_{\phi}(\bar{\alpha}_t x_0 + \bar{\beta}_t \epsilon, t) – \epsilon |^2 \right] ]
同样地，如果我们想用 ( g_{\theta}(z) ) 的学生数据训练一个扩散模型，那么训练目标是：
[ \psi^* = \arg\min_{\psi} E_{z, \epsilon \sim N(0, I)} \left[ | \epsilon_{\psi}(x_t^{(g)}, t) – \epsilon |^2 \right] ]
其中 ( x_t^{(g)} = \bar{\alpha}t g{\theta}(z) + \bar{\beta}t \epsilon )，是由学生数据加噪后的样本，其分布记为 ( p{\theta}(x_t^{(g)}) )。
学生模型的学习
我们可以通过最小化教师模型和学生模型生成的数据分布差异来学习学生模型：
[ \theta^* = \arg\min_{\theta} E_{z, \epsilon \sim N(0, I)} \left[ | \epsilon_{\phi^}(x_t^{(g)}, t) – \epsilon_{\psi^}(x_t^{(g)}, t) |^2 \right] ]
注意这个优化依赖于 ( \theta )，所以当 ( \theta ) 通过上式发生改变时，( \psi^* ) 的值也随之改变，因此需要交替优化，类似GAN一样。
点睛之笔
上述方法存在理论与实践之间的gap，主要体现在两个问题：
- 理论上要求先求出上式的最优解，然后才去优化，但实际上从训练成本考虑，我们并没有将它训练到最优就去优化了；
- 理论上 ( \psi^* ) 随 ( \theta ) 而变，即应该写成 ( \psi^(\theta) )，从而在优化时应该多出一项 ( \psi^(\theta) ) 对 ( \theta ) 的梯度，但实际上在优化时我们都只当 ( \psi^* ) 是常数。
SiD的核心贡献是通过恒等变换，尽量消除优化目标对 ( \psi^* ) 的依赖，从而有效缓解上述两个问题。
恒等变换
我们具体来看做了什么恒等变换。我们先来看去噪目标：
[ E_{x_0 \sim \tilde{p}(x_0), x_t \sim p(x_t | x_0)} \left[ | \epsilon_{\phi}(x_t, t) + \bar{\beta}t \nabla{x_t} \log p(x_t | x_0) |^2 \right] ]
根据得分匹配相关结果，上述目标的最优解是 ( \epsilon_{\phi^}(x_t, t) = -\bar{\beta}t \nabla{x_t} \log p(x_t) )。同理，学生模型的训练目标的最优解是 ( \epsilon_{\psi^}(x_t^{(g)}, t) = -\bar{\beta}t \nabla{x_t^{(g)}} \log p_{\theta}(x_t^{(g)}) )。
此时我们有：
[ E_{z, \epsilon \sim N(0, I)} \left[ | \epsilon_{\phi^}(x_t^{(g)}, t) – \epsilon_{\psi^}(x_t^{(g)}, t) |^2 \right] ]
通过恒等变换，我们可以将上式化简为：
[ E_{z, \epsilon \sim N(0, I)} \left[ \langle \epsilon_{\phi^}(x_t^{(g)}, t) – \epsilon_{\psi^}(x_t^{(g)}, t), \epsilon_{\phi^*}(x_t^{(g)}, t) – \epsilon \rangle \right] ]
这就是SiD的核心结果，能够高效地实现蒸馏。
其他细节
- 论文的推导默认了 ( \bar{\alpha}_t = 1 )。
- 论文的结果是以 ( \mu(\bar{x}_t) = x_t – \bar{\beta}_t \epsilon(x_t, t) / \bar{\alpha}_t ) 为标准给出的，这与扩散模型常见的表述方式不同。
- SiD最终取了上式的相反数作为额外的损失函数，加权到改进的损失函数上，以取得更优的蒸馏效果。
延伸思考
对于式(3)和式(4)的交替优化，有不少读者可能会想到，但SiD的精彩之处是提出了恒等变换，使得训练更加稳定高效。
文章小结
在这篇文章中，我们介绍了一种新的将扩散模型蒸馏为单步生成模型的方案，其思想可以追溯到前两年的利用去噪自编码器训练生成模型的工作，它不需要获得教师模型的真实训练集，也不需要迭代教师模型来生成样本对，而引入了类似GAN的交替训练，同时提出了关键的恒等变换来稳定训练过程，整个方法有颇多值得学习之处。