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分钟。

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