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）。该方法的动机来自于一个有趣的几何观察：扩散模型的生成轨迹几乎位于图像空间的一个二维子空间中！于是，通过直接学习生成轨迹的平均方向，我们可以进一步减小离散误差。此外，该方法可以作为插件运用在各种现有方法上，进一步提升它们的性能。 本文代码已开源，仓库中不仅提供了该方法的代码，还提供了一个采样工具箱，复现了多个现有工作，设立了新的实验基准。
随着理论和实际应用的不断发展和完善，扩散模型（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）。该方法的动机来自于一个有趣的几何观察：扩散模型的生成轨迹几乎位于图像空间的一个二维子空间中！于是，通过直接学习生成轨迹的平均方向，我们可以进一步减小离散误差。此外，该方法可以作为插件运用在各种现有方法上，进一步提升它们的性能。
本文代码已开源，仓库中不仅提供了该方法的代码，还提供了一个采样工具箱，复现了多个现有工作，设立了新的实验基准。
论文地址：
项目地址：
近年来，扩散模型以其出色的生成能力受到越来越多的关注。给定噪声输入，扩散模型利用评分函数（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。方法总览如下。
我们在各个规模的数据集上（图像分辨率从 32x32 到 512x512）验证了本文方法的有效性。主要实验结果如下图所示，其中 NFE（number of function evaluation）为采样中扩散模型的前向传播次数。我们利用到 [3] 中提出的 AFS 技巧，在采样的第一步不使用扩散模型的输出而是直接对第一步的方向做粗略估计，以此可以在几乎不影响样本质量的同时减少一次网络前向传播。
以下是在 Stable Diffusion v1.4 上的定性和定量结果。
总结
在这项工作中，我们提出了名为 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 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：[email protected]
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
·
最新评论
推荐文章
作者最新文章
你可能感兴趣的文章
Copyright Disclaimer: The copyright of contents (including texts, images, videos and audios) posted above belong to the User who shared or the third-party website which the User shared from. If you found your copyright have been infringed, please send a DMCA takedown notice to [email protected]. For more detail of the source, please click on the button "Read Original Post" below. For other communications, please send to [email protected].
版权声明：以上内容为用户推荐收藏至CareerEngine平台，其内容（含文字、图片、视频、音频等）及知识版权均属用户或用户转发自的第三方网站，如涉嫌侵权，请通知[email protected]进行信息删除。如需查看信息来源，请点击“查看原文”。如需洽谈其它事宜，请联系[email protected]。
版权声明：以上内容为用户推荐收藏至CareerEngine平台，其内容（含文字、图片、视频、音频等）及知识版权均属用户或用户转发自的第三方网站，如涉嫌侵权，请通知[email protected]进行信息删除。如需查看信息来源，请点击“查看原文”。如需洽谈其它事宜，请联系[email protected]。