作者提出了名为AMED-Solver的扩散模型ODE求解器，旨在最小化扩散模型在少步数采样下的离散误差。现有扩散模型ODE求解器生成的采样轨迹大致位于一个二维子空间中，因此可以通过中值定理的想法学习一个近似的平均方向。作者将AMED-Solver的思想推广到AMED-Plugin，得到一个可以应用于各种扩散模型ODE求解器的插件，以进一步提高生成样本的质量。随着理论和实际应用的不断发展和完善，扩散模型（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]中的推导，扩散模型的反向扩散过程（即采样过程）可以由以下常微分方程刻画：
2.2 利用近似平均方向减小离散误差
注意到前文常微分方程的解是
2.3 训练流程
三、实验效果
我们在各个规模的数据集上（图像分辨率从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步左右）取得了先进的效果。希望我们的尝试能激发未来的研究，进一步释放快速求解器在扩散模型采样上的潜力。
