专栏名称: 极市平台

极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台，为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯，行业动态，在线分享信息，线下活动等。网站: http://cvmart.net/

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

极市平台 · 公众号 · · 2024-04-22 22:00

正文

↑ 点击蓝字关注极市平台

作者丨zhyzhou

来源丨将门创投

编辑丨极市平台

极市导读

作者提出了名为AMED-Solver的扩散模型ODE求解器，旨在最小化扩散模型在少步数采样下的离散误差。现有扩散模型ODE求解器生成的采样轨迹大致位于一个二维子空间中，因此可以通过中值定理的想法学习一个近似的平均方向。作者将AMED-Solver的思想推广到AMED-Plugin，得到一个可以应用于各种扩散模型ODE求解器的插件，以进一步提高生成样本的质量。 >> 加入极市CV技术交流群，走在计算机视觉的最前沿

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