ICML 2024 | 论扩散模型采样轨迹的规律性及快速采样算法

扩散式生成模型近年来在理论与应用方面均取得了巨大的进展。随机微分方程及其数值方法的引入使得我们能够对扩散模型在连续时间中分析，离散步数中采样；概率流常微分方程及其数值方法的广泛使用则是使得扩散模型的少步数快速采样成为可能。

本文进一步揭示了扩散模型采样轨迹蕴含的规律性，即，高维空间中的采样轨迹总是呈现出低维的回旋镖结构。该采样轨迹可以由一个方向向量及其正交补空间中的两个主成分有效表达。该现象与采样轨迹的初始噪声、合成内容均无关，且可以借助于变尺度核密度估计进行理论分析。

本文刻画的几何图像既能有效解释多种在实际中取得成功的启发式策略，又能从原理上导出简单、几乎无代价的步长规划策略，从而进一步增强扩散模型在少步数情况下的采样质量。

论文标题： On the Trajectory Regularity of ODE-based Diffusion Sampling 论文地址： https://arxiv.org/abs/2405.11326 代码地址： https://github.com/zju-pi/diff-sampler

一、预备知识

扩散模型通过正向加噪过程将数据转化成噪声，再通过反向去噪过程从噪声合成数据。该想法可以被随机微分方程这一数学工具形式化。特别地，上述过程也可以被确定性的概率流常微分方程所描述。这两种框架能够在保持边缘概率密度函数不变的情况下进行等价转换。

本文的讨论主要基于一种特定的线性扩散过程（很容易证明能够推广到其他线性扩散过程），其正向加噪过程由下述随机微分方程所刻画

初值条件为 。其中， 是关于时间 单调递增的噪声函数。

该扩散过程对应的概率流常微分方程（PF-ODE）为

根据经验贝叶斯（Empirical Bayes）的结论[ ]，我们可以通过训练一个去噪自编码器 (noise-dependent Denoising AutoEncoder)来估计 ，并将上述PF-ODE写为如下empirical PF-ODE的形式

扩散模型的采样过程：首先从噪声分布中随机采样初始噪声，即 ，再基于预先指定的时间步 和数值方法求解empirical PF-ODE从而得到合成的数据 。我们记 为采样轨迹，记 为去噪轨迹。

二、扩散模型采样轨迹的规律性

由于我们难以直接在原始高维数据空间中可视化采样轨迹，本文提出了低维轨迹投影技术。

2.1 一维轨迹投影

对于给定的采样轨迹，我们将轨迹初始点 和轨迹终点 相连，并计算采样轨迹的所有中间点 到该直线的垂直距离。该距离如下图红线所示（从右往左看）。下图蓝线则表示采样轨迹的所有中间点距离轨迹终点 的距离。该图的统计结果由5千条随机初始化的采样轨迹计算得到。

该实验结果显示了采样轨迹上的点会略微偏离初始点 和终点 构成的直线，但是轨迹的最大偏移量相对来说仍是比较小的（30/8868 0.0034）。即，所有采样轨迹均会共享一个近似线性结构。

另外，轨迹偏离量也等价于将原始高维采样轨迹用一维向量 表达所带来的重构误差。该实验结果也体现了一维轨迹投影未能精准地刻画原始采样轨迹的形状。

2.2 多维轨迹投影

为了进一步减少投影操作带来的误差，从而使得低维投影轨迹能够尽可能地保留原始高维轨迹的形状，我们除了将 的方向作为基向量以外，还进一步利用该向量正交补空间中的主成分作为基向量。

具体来说，如下图(a)所示，我们首先将原始 维空间中的采样轨迹投影到一维向量 的正交补空间（ 维），再将该正交补空间中的top-k主成分向量作为轨迹重构的基向量。下图(b)-(d)的结果显示，我们仅需使用一维向量 再加上两个主成分向量，就足以捕获原始 维采样轨迹表达的视觉信息，也能够大幅度降低轨迹重构的误差。

因此，我们能够很安全地在三维子空间中重构并可视化原始维采样轨迹的形状，并观察其中蕴含的规律性。

上图的投影结果显示了扩散模型的采样轨迹总是呈现出低维回旋镖结构，并且该结构与初始噪声、合成内容均无关。

本文对该结构出现的原因进行了理论分析，详见论文第四章"Understanding the Trajectory Regularity"。

三、几何启发的步长规划策略

基于本文揭示的扩散模型采样结构，我们给出了一种基于动态规划的步长分配策略。核心思想：既然所有采样轨迹均共享类似的形状，那么我们就可以采用少量的样本估计出更优的步长分配策略。其原则是，在轨迹弯曲程度较大的地方使用较小的采样步长；而在轨迹弯曲程度较小的地方使用较大的采样步长。该算法的具体细节见论文第五章“Geometry-Inspired Time Scheduling”。

对于任一数据集，本文提出的步长规划策略仅需运行一次，就能得到扩散模型在各个NFE情况下的最优步长分配策略。该方法的运行速度极快，对于CIFAR-10，FFHQ，ImageNet 64 64等小分辨率数据集单卡仅需1分钟左右，对于Stable Diffusion使用的LAION数据集也仅需15分钟。

实验结果如下所示。基于对扩散模型采样轨迹性质更加深入的理解，我们很自然地就以几乎无代价的方式取得了目前最好的少步数数值采样结果。

四、结论

本文揭示了扩散模型采样轨迹蕴含的规律性。该规律性可以通过刻画和分析扩散模型采样过程中形成的去噪轨迹进行分析。本文展示的几何图像也能够启发我们进一步提高少步数扩散模型数值采样的效果。

llustration From IconScout By Kawalan Studio

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（

www.techbeat.net

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

melodybai@thejiangmen.com

或添加工作人员微信（yellowsubbj）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈