CVPR 2024 | 浙大提出近似平均方向求解器，快速扩散模型数值采样

随着理论和实际应用的不断发展和完善，扩散模型（Diffusion Models）已经成为当今生成式模型的一大支柱，其强大的生成能力给 AIGC 带来了一系列重大突破。然而，扩散模型逐步去噪的特性导致其生成速度极其缓慢，阻碍了它在实践中的应用。 

近年来，已有许多研究利用高阶数值方法极大提高了扩散模型的生成效率。然而，数值方法固有的离散误差会使样本质量在少步数生成下（例如 5 步）出现显著降低。扩散模型在少步数下的生成仍然存在困难。 

本文介绍被 CVPR 接收为 Highlight（2.8%）的工作“Fast ODE-based Sampling for Diffusion Models in Around 5 Steps”。本文提出了近似平均方向求解器（Approximate MEan-Direction Solver，AMED-Solver）。该方法的动机来自于一个有趣的几何观察：扩散模型的生成轨迹几乎位于图像空间的一个二维子空间中！于是，通过直接学习生成轨迹的平均方向，我们可以进一步减小离散误差。此外，该方法可以作为插件运用在各种现有方法上，进一步提升它们的性能。 

本文代码已开源，仓库中不仅提供了该方法的代码，还提供了一个采样工具箱，复现了多个现有工作，设立了新的实验基准。

论文题目：

Fast ODE-based Sampling for Diiffusion Models in Around 5 Steps

论文地址：

https://arxiv.org/abs/2312.00094

项目地址：

https://github.com/zju-pi/diff-sampler

引言

近年来，扩散模型以其出色的生成能力受到越来越多的关注。给定噪声输入，扩散模型利用评分函数（score function）进行迭代去噪实现图像生成、音频生成、视频生成、文生图等任务。这个过程可以解释为离散求解某个随机微分方程（SDE），或是求解其对应的概率流常微分方程（PF-ODE）[1]。与其他生成模型相比，扩散模型在样本质量和训练稳定方面具有优势，但扩散模型缓慢的采样速度对其实际应用带来了阻碍。

在扩散模型加速采样上，现有方法主要分为两大类。一是设计更快的数值求解器，在增加步长同时保持较小的离散误差。这些方法成功地将采样步数从 1000 减少至 20 以下。二是基于知识蒸馏，在数据分布和预先指定的噪声分布之间建立一一映射。如此，训练好的学生模型只需一步就可以实现高质量生成。

然而，基于蒸馏的加速方法常常要求较大的训练开销，或者需要仔细设计训练流程来保证稳定。本文中，我们结合两类方法的长处，在保持较低训练开销的同时，利用蒸馏得到的知识，进一步加速扩散模型采样。本文主要贡献如下：

• 提出 AMED-Solver，一种新的扩散模型 ODE 求解器，其通过学习近似平均方向来最小化离散误差。
• 提出 AMED-Plugin，一个可以应用于各种扩散模型 ODE 求解器的插件，仅引入较小的训练开销和可忽略的采样开销。

本文方法

2.1 核心观察

扩散模型的图像生成可以归结为一个求解微分方程的任务。由于在少步数下离散求解 SDE 难以消除噪声对于解的影响，本文主要考虑求解 PF-ODE。根据[2]中的推导，扩散模型的反向扩散过程（即采样过程）可以由以下常微分方程刻画：

其中，  为预训练的扩散模型。给定一列预先指定的时刻 以及 时刻服从标准高斯分布的噪声图像 ，通过从 时刻到 时刻求解该方程，我们就可以将噪声图像还原成一张逼真的图像。我们称离散求解过程中中间输出构成的集合 为扩散模型的采样轨迹。

我们的核心观察在于，扩散模型的采样轨迹虽然位于一个维度非常高（通常上千甚至上万）的空间，但其可以近似地用两个主成分来刻画，也就是说，这条轨迹几乎处在一个二维子空间上！

如下图所示，我们利用 EDM 提供的 Heun Solver[2] 进行 40 步离散采样得到 1000 条轨迹，对每条轨迹做主成分分析并提取部分主成分重构轨迹，得到的相对重构误差非常小，重构轨迹也能解释原轨迹 99% 以上的方差。

2.2 利用近似平均方向减小离散误差

注意到前文常微分方程的解是

结合上述观察，我们想到利用中值定理的思想，尝试寻找一个最优的中间时刻 以及一个伸缩系数 ，来得到一个近似平均方向，满足

从而有如下的离散求解形式：

本文称该离散求解器为 Approximate Mean-Direction Solver（AMED-Solver）。

基于类似的思想，我们可以将这个做法推广到各种求解器上，成为一个可训练的插件。例如，为了进行 8 步采样，我们可以初始化一个 4 步的采样时刻表，然后在相邻两步之间插入可学习的时刻以及相应的伸缩系数 ，以此来减小离散误差。我们称这个推广为 AMED-Plugin。更多细节请参考原文。

2.3 训练流程

对于每一对需要寻找的参数 和 ，我们额外训练一个非常小的神经网络（仅有共 9000 个可训练参数）来预测它们，称这个网络为 AMED predictor。方法总览如下。

AMED predictor 的训练利用到蒸馏的思想。以训练 AMED-Solver 为例（AMED-Plugin 类似），我们先生成一个基本采样时刻表 ，AMED predictor 负责在相邻两个时刻之间插入一个可学习的时刻 （以及伸缩系数 ），采样出一条步数较少的学生轨迹。

对于教师轨迹，我们在相邻两个时刻之间插入多个新的时刻，得到一条更为精确的轨迹，进而从 时刻到 时刻依次计算学生样本和教师样本之间的 L2 范数并进行反向传播，对 AMED predictor 进行更新。利用一张 NVIDIA A100 GPU，整个训练流程在小数据集（如 CIFAR10）上仅仅只需要几分钟。

我们在各个规模的数据集上（图像分辨率从 32x32 到 512x512）验证了本文方法的有效性。主要实验结果如下图所示，其中 NFE（number of function evaluation）为采样中扩散模型的前向传播次数。我们利用到 [3] 中提出的 AFS 技巧，在采样的第一步不使用扩散模型的输出而是直接对第一步的方向做粗略估计，以此可以在几乎不影响样本质量的同时减少一次网络前向传播。

以下是在 Stable Diffusion v1.4 上的定性和定量结果。

以下是在 FFHQ64 和 Imagenet64 上，AMED-Solver 与其他采样器图像质量的对比。

总结

在这项工作中，我们提出了名为 AMED-Solver 的扩散模型 ODE 求解器，旨在最小化扩散模型在少步数采样下的离散误差。我们的主要观察是，现有扩散模型 ODE 求解器生成的采样轨迹大致位于一个二维子空间中，因此我们可以通过中值定理的想法学习一个近似的平均方向。

我们将 AMED-Solver 的思想推广到 AMED-Plugin，得到一个可以应用于各种扩散模型 ODE 求解器的插件，以进一步提高生成样本的质量。实验表明，我们的方法在极小的离散步数下（5 步左右）取得了先进的效果。希望我们的尝试能激发未来的研究，进一步释放快速求解器在扩散模型采样上的潜力。

参考文献

[1] Song Y, Sohl-Dickstein J, Kingma D P, et al. Score-based generative modeling through stochastic differential equations[J]. arXiv preprint arXiv:2011.13456, 2020.

[2] Karras T, Aittala M, Aila T, et al. Elucidating the design space of diffusion-based generative models[J]. Advances in Neural Information Processing Systems, 2022, 35: 26565-26577.

[3] Dockhorn T, Vahdat A, Kreis K. Genie: Higher-order denoising diffusion solvers[J]. Advances in Neural Information Processing Systems, 2022, 35: 30150-30166.

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·