类数值方法PNDM：Stable Diffusion默认加速采样方案

2022-11-04 03:11

©PaperWeekly 原创 · 作者 | 刘路平

单位 | 浙江大学

研究方向 | 扩散模型

扩散模型 [1]（diffusion model）通过逐步迭代的方式建立起高斯分布和数据分布的对应，将深度生成模型提升到了新的高度。扩散模型生成结果的细节准确性和样本丰富度都给人留下了深刻的印象。

然而，扩散模型需要成百上千次的迭代来得到最终的结果，巨大的计算开销阻碍了扩散模型的应用和落地。我们深入分析扩散模型的迭代公式，通过引入数值分析中的高阶数值方法，并与扩散模型中的无随机项迭代法（DDIM）相结合，设计出适合扩散模型的高阶快速迭代方法。

在今年的 ICLR 2022 上，我们提出了全新的适用于扩散模型的类数值方法（PNDM）。这一方法既不需要重新训练模型，也对模型结构和超参数没有额外的限制，仅仅通过修改迭代公式，将扩散模型无任何精度损失加速 20 倍，迭代 50 步就可以达到原来 1000 步的 FID 结果。我们的方法被应用到 Stable Diffusion [2,7] 中，作为其默认加速采样算法。

论文标题：

Pseudo Numerical Methods for Diffusion Models on Manifolds

论文链接：

https://openreview.net/forum?id=PlKWVd2yBkY

代码链接：

https://github.com/luping-liu/PNDM

研究背景

降噪扩散概率模型 [1]（DDPM）通过逐步加噪的方式建立目标数据分布和高斯分布之间的对应，并通过它的逆过程降噪过程生成新的目标样本。

具体来说，是训练数据，加噪迭代过程满足：

这里，控制加噪的速度。进一步我们可以得到与的关系：

这里，, 。根据 DDPM 和 DDIM [3] 的分析，对应的降噪迭代公式满足：

这里，是用神经网络对带噪样本中的噪声的估计，控制着迭代中添加随机噪声的比例。当，在降噪过程中不添加任何新的随机噪声，也就是无随机项迭代法（DDIM）。我们将其记为：

研究思路

扩散模型的迭代可以关于取极限后转化为求解一个常微分方程，对应的常微分方程满足：

进而可以通过使用高阶的数值方法进行加速。然而，通过实验，我们发现这种直接的套用效果并不好。在有些情况下，四阶的数值方法甚至不如一阶的 DDIM。通过对扩散模型的方程深入分析，我们发现这里存在两个问题：

这个常微分方程并不满足数值方法要求的 Lipschitz 连续性。换句话说，他的梯度是会趋于无穷的，这影响了数值方法的对梯度的拟合，也限制了最终的生成效果。
这个常微分方程包含一个神经网络。这意味着方程并不是在整个欧氏空间里有定义，而是只有神经网络被训练过的区域上，才能真实反映数据的梯度和走向。一旦迭代过程中跑出训练数据集中的区域，迭代过程就会出现震荡甚至发散的情况。

为了解决这两个问题，我们的思路是，直接在原本的迭代算法上，设计新的高阶数值方法。而不是像之前的方法，先将原本离散的扩散模型迭代先连续化为微分方程，再反过来通过数值方法离散化求解。这样就解决了第一个问题，因为我们并非使用经典数值方法，也就避免了方程不满足 Lipschitz 连续性的问题。

另一方面，对于具体如何改进原始迭代公式，我们发现了一个重要的性质，如果迭代过程中的噪声项可以被准确拟合，那么通过和我们就可以直接得到精确的。因此，我们只需要拟合出更精确的噪声项就可以了。迭代公式的其他部分并不引入误差，也不需要进一步的修正。明确了优化的目标后，我们就可以通过高阶数值方法对目标进行修正，来达到加速的效果。

这也很大程度上解决了第二个问题，因为不需要整体改变迭代公式，这保证了迭代生成的轨迹和原始 DDIM 迭代公式相似。而 DDIM 的迭代公式恰恰可以保证生成数据和训练数据符合相同的分布。因此迭代过程可以大概率落在训练数据集中的高密度区域，不会跑出原始数据分布。先连续化再离散的策略则破坏了这种一致性，这也是 DDIM 能好于直接套用四阶数值方法的原因。

我们的方法融合了 DDIM 和高阶数值方法方法的优点，既保证了拟合分布，也提升了收敛速度。实验结果如下图所示：

具体来说，我们的方法是在原始的 DDIM 的方法上：

对其中的项使用线性多步法做四阶修正，这一方法称作类线性多步法（PLMS）：

线性多步法可以在没有额外开销的情况下，提升收敛阶数。但是，这一方法是不能自启动的，需要额外的方法提供开始的三步的数据。这里我们使用 Runge-Kutta 方法计算前三步，这一方法称作类 Runge-Kutta 法（PRK）：

这虽然会在一开始造成额外的开销，但是可以有效防止一开始造成较大误差并在后面的迭代中不断累积放大，这也保证了方法整体的收敛阶数。这两种子方法合在一起被称作类数值方法（PNDM），即：

方法扩展

上面的思路和方法是我们论文中的主要方法，但在具体落地应用中，又有了一些有趣的变化和拓展。

3.1 快速启动

使用 Runge-Kutta 法启动方法会有额外的开销，在总的迭代步数相对较少时，这个开销是不可忽视的。在这种情况下，可以选择依次使用一二三阶的类线性多步法 [4] 启动迭代过程，这样每一步都没有了额外的开销，反而能在一些情况下取得更好的效果。

3.2 剪裁

我们的方法并没有破坏 DDIM 的迭代逻辑，只是增加了额外的一项修正。因此，其他可以提升 DDIM 迭代效果的方法也可以直接使用到我们的方法上。其中的重要代表就是数据剪裁 [5]，因为迭代过程还是存在中间过程跑出数据分布的情况，通过剪裁就可以进一步避免这种情况的发生。我们的方法可以直接将：

替换为对应的：

这样就将剪裁集成到了我们的方法中，对于更一般的动态剪裁，方法也是类似的。这里用到了这一项代表了对的估计。

3.3 条件生成

我们的方法不仅仅可以用来做无条件生成，还可以用来加速有条件生成。以无分类器指导的条件生成 [6] 为例，这种条件生成通过修改来实现。即将迭代公式中的修改为：

这里，为添加的条件，控制条件指导的强度。这样简单替换就可将我们的方法用于加速条件生成。由此可以看到，我们的方法和 DDIM 保持了一致，都可以非常方便地进一步对包括和在内的中间项进行调整和优化，从而可以被应用到广泛的落地场景中。

可以说，我们的类数值方法 PNDM 为扩散模型的落地扫除了计算速度上的重要障碍。目前，PNDM 被 AI 内容生成独角兽公司 Stability AI 的 AI 绘图平台 stable diffusion [2,7] 用作默认加速采样算法，并被包括 AI 开源平台 huggingface 的 diffusers [8] 和浙大语音合成项目 DiffSinger [9] 等多个开源库收录和使用。欢迎大家都来多多关注和使用我们的开源项目。

参考文献

[1] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33 (2020): 6840-6851.

[2] Rombach R, Blattmann A, Lorenz D, et al. High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 10684-10695.

[3] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. International Conference on Learning Representations. 2021.

[4] Zhang Q, Chen Y. Fast Sampling of Diffusion Models with Exponential Integrator. arXiv preprint arXiv:2204.13902, 2022.

[5] Saharia C, Chan W, Saxena S, et al. Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding. arXiv preprint arXiv:2205.11487, 2022.

[6] Ho J, Salimans T. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.

[7] https://github.com/CompVis/stable-diffusion

[8] https://github.com/huggingface/diffusers

[9] https://github.com/MoonInTheRiver/DiffSinger

更多阅读