在生成模型领域，扩散模型（Diffusion Models, DMs）因其卓越的生成质量而成为最新的技术趋势。但这些模型的一个关键缺点是它们的采样速度较慢，需要通过大型神经网络进行多次顺序函数评估。扩散模型通过一个称为采样计划的离散噪声水平集来解决微分方程。尽管过去的研究主要集中在开发高效的求解器上，但很少有人关注寻找最优的采样计划。大多数现有工作都依赖于手工制定的启发式计划，如简单多项式和余弦函数。本文提出了一种新的框架，名为“Align Your Steps”（AYS），用于优化扩散模型中的采样计划，显著提高了输出质量，尤其是在仅进行少量步骤合成时。

AYS框架

AYS框架利用随机微积分方法，为不同的求解器、训练有素的DMs和数据集找到最优的采样计划。该框架基于所有随机SDE求解器都可以重新解释为在短间隔内精确求解近似线性化SDE的观察结果。通过将问题框架化为采样计划上的优化问题，使用随机微积分技术最小化近似线性SDE和真实生成SDE之间的不匹配。

AYS框架的核心在于它能够根据数据集的特性、所使用的模型以及选定的求解器，定制出最优的采样计划。这种定制化的方法是传统启发式采样计划所不具备的，后者通常依赖于简单的多项式衰减或余弦退火等规则。通过优化采样计划，AYS框架能够确保在给定的计算预算内，生成的样本尽可能地接近真实数据分布。

为了找到最优的采样计划，AYS框架采用了一种基于随机微积分的方法。它通过最小化真实生成的SDE（随机微分方程）与近似线性化SDE之间的Kullback-Leibler (KL) 散度上界（KLUB），来优化采样计划。这种方法允许研究者在不同的时间点调整采样计划，以确保生成过程的每一步都能够有效地逼近真实的数据生成过程。

AYS框架的另一个显著优势是其快速合成能力。在需要迅速生成高质量样本的应用场景中，AYS框架通过优化采样步骤，减少了生成过程中所需的前向评估次数，从而加快了整个合成过程。AYS框架还具有很好的通用性。它不仅限于特定的数据类型或模型架构，而是可以泛化到多种不同的扩散模型和求解器。无论研究者或开发者使用的是图像、文本还是视频数据，都可以利用AYS框架来优化他们的采样计划。

在实现上，AYS框架采用了迭代优化的方法。首先使用启发式方法初始化采样计划，然后通过迭代过程调整各个时间点，以最小化KL散度上界。这个过程可以高度并行化，因为各个时间点的优化可以独立进行。为了确保优化过程的稳定性和有效性，AYS框架还引入了早停机制。

优化采样计划

在进行少步骤合成时，现有的采样计划可能导致显著的离散化误差，从而影响最终的输出质量。为了证明优化采样计划的必要性，研究者以一个简单的高斯数据分布为例，展示了如何通过解析方法得到最优采样计划，并发现该计划与文献中常用的启发式采样计划有显著不同。

研究者使用Girsanov定理来分析离散化误差，该定理提供了两个共享扩散项的随机微分方程（SDEs）输出差异的上界。通过这个上界（KL-divergence Upper Bound, KLUB），研究者将寻找最优采样计划的问题转化为一个优化问题，就是最小化实际求解SDE/ODE时使用的离散采样计划与理论上无需离散化的逆向生成SDE之间的差异。

为了解决上述优化问题，研究者提出了一个迭代方法。该方法首先选择一个采样计划的索引，然后在其周围选择多个候选点，计算每个候选点的KLUB，并选择使KLUB最小的候选点作为优化后的采样点。这个过程可以高度并行化，特别是对于非邻近的索引。

为了减少方差并提高KLUB估计的准确性，研究者使用了针对时间t的重要性采样。通过假设高斯数据分布，并解析计算所有积分项，然后从与这些计算值匹配的概率密度函数中采样t，研究者显著降低了KLUB估计的方差。

在实践中，采样计划的优化是分层进行的。首先，使用启发式计划初始化一个10步的采样计划，然后迭代优化所有中间点，并使用早停机制避免过度优化。接下来，进行两轮细分和进一步微调，以获得40步的采样计划。在每次细分后，只优化新添加的中间点，而保持其他点固定。这使得采样计划的一般“形状”变得固定，从而在这些后期阶段不需要早停。

为了获得与[10, 20, 40]不同的步数的计划，研究者将40步计划视为分段对数线性函数，并对其进行插值以匹配所需的步数。

实验

实验部分旨在展示优化采样计划（Optimizing Sampling Schedules）对扩散模型性能的显著影响。研究者采用了FID得分作为主要的定量评估指标，并通过用户研究提供了定性评估。通过一系列实验验证了AYS（Align Your Steps）框架在不同数据集和模型上的有效性，包括2D玩具数据、标准图像数据集以及文本到图像和视频生成模型。

研究者首先在2D玩具数据集上展示了优化采样计划的优势。这些数据集具有已知的真实分布，使得研究者能够直观地比较不同采样计划生成样本的质量和原始分布的接近程度。实验结果表明，使用优化采样计划的模型能够生成更接近原始分布的样本，并且减少了异常值的出现。

接下来，研究者在CIFAR10、FFHQ和ImageNet等标准图像数据集上进行了实验。他们使用了预训练的连续时间扩散模型，并采用了不同的随机和确定性求解器。实验结果显示，优化的采样计划在低前向评估（NFE）次数下显著提高了生成图像的质量，减少了FID（Fréchet Inception Distance）得分，这是评估生成图像质量的常用指标。

研究者还将AYS框架应用于流行的开源文本到图像模型，如Stable Diffusion 1.5和SDXL。这些模型依赖于分类器自由引导（classifier-free guidance），研究者发现，即使使用默认引导值优化的采样计划，也能在合理的引导值范围内有效。通过用户研究，研究者进一步验证了优化采样计划在图像质量和文本对齐方面的优势。

对于视频生成，研究者使用Stable Video Diffusion（SVD）模型进行了实验。他们发现，使用优化的采样计划能够改善视频生成中的时间颜色一致性，并解决了视频帧过度饱和的问题。通过用户研究，研究者证明了优化采样计划在视频生成中的实际优势。

实验结果一致表明，通过优化采样计划，可以在保持相同计算预算的情况下，显著提高扩散模型生成样本的质量。研究者还讨论了优化采样计划对于模型多样性和实时应用的潜在影响。

实验结果表明，无论是在图像、视频还是2D玩具数据上，AYS都能显著超越传统的启发式采样计划，实现更快、更准确的数据生成。AYS的通用性和自动化特性，为广泛的应用领域和不同需求的用户提供了一个强大的工具。随着扩散模型在各个领域的不断进步，AYS框架有望成为推动生成模型性能提升的关键技术之一，为未来的研究和实践开辟新的道路。

论文链接：https://arxiv.org/abs/2404.14507

项目地址：https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/