CVPR 2024 Highlight | AMED-Solver：快速扩散模型数值采样

随着理论和实际应用的不断发展和完善，扩散模型（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条轨迹，对每条轨迹做主成分分析并提取部分主成分重构轨迹，得到的相对重构误差非常小，重构轨迹也能解释原轨迹 以上的方差。

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.

llustration From IconScout By TanahAir Studio

公众号后台回复“数据集”获取100+深度学习各方向资源整理