[论文解读] SA-Solver: Stochastic Adams Solver for Fast Sampling of Diffusion Models

Shuchen Xue, Mingyang Yi+5 more|arXiv (Cornell University)|Sep 10, 2023

Bayesian Methods and Mixture Models被引用 6

一句话总结

本文提出 SA-Solver，一种用于扩散概率模型中快速且高质量采样的新型随机 Adams 方法。通过引入具有方差控制特性的扩散 SDE 家族，并利用高效的线性多步 SDE 求解器，SA-Solver 在较少函数评估次数（NFE）下，于多个基准测试中实现了最先进（SOTA）的 FID 分数，优于所有确定性和现有随机采样器，尤其在低 NFE 设置及部分训练模型下表现更优。

ABSTRACT

Diffusion Probabilistic Models (DPMs) have achieved considerable success in generation tasks. As sampling from DPMs is equivalent to solving diffusion SDE or ODE which is time-consuming, numerous fast sampling methods built upon improved differential equation solvers are proposed. The majority of such techniques consider solving the diffusion ODE due to its superior efficiency. However, stochastic sampling could offer additional advantages in generating diverse and high-quality data. In this work, we engage in a comprehensive analysis of stochastic sampling from two aspects: variance-controlled diffusion SDE and linear multi-step SDE solver. Based on our analysis, we propose extit{SA-Solver}, which is an improved efficient stochastic Adams method for solving diffusion SDE to generate data with high quality. Our experiments show that extit{SA-Solver} achieves: 1) improved or comparable performance compared with the existing state-of-the-art (SOTA) sampling methods for few-step sampling; 2) SOTA FID on substantial benchmark datasets under a suitable number of function evaluations (NFEs). Code is available at https://github.com/scxue/SA-Solver.

研究动机与目标

• 解决现有无需训练的扩散模型采样方法在效率和质量上的局限性。
• 探索随机采样在生成比确定性求解器更高质量、更多样化数据方面的潜力。
• 通过设计具有可调噪声尺度的方差控制型扩散 SDE 新家族，提升采样效率。
• 开发一种高效的线性多步 SDE 求解器，使其性能达到或超过现有 ODE 求解器。
• 在无需额外模型训练的前提下，实现在低 NFE 设置下的最先进生成质量。

提出的方法

• 提出一类与原始扩散 SDE 具有相同边缘分布但可通过可调参数 τ 控制噪声尺度（方差）的 SDE 家族。
• 将 [29] 中的随机 Adams 方法改进以求解所得 SDE，通过引入适用于扩散过程的指数加权积分。
• 为所提出的随机 Adams 求解器提供理论收敛阶分析，以确保数值稳定性和准确性。
• 采用线性多步框架以提升效率，减少采样所需的函数评估次数（NFE）。
• 通过校准噪声尺度 τ 平衡生成样本的多样性与质量，实现生成质量的可控权衡。
• 采用与标准 DPM 相同的反向过程，但通过所提出的 SA-Solver 实现更优的数值积分。

实验结果

研究问题

• RQ1 与确定性求解器相比，通过控制噪声方差的随机采样是否能提升扩散模型的生成质量？
• RQ2 如何设计一种高效的线性多步 SDE 求解器，使其在 NFE 效率方面达到或超过现有 ODE 求解器的性能？
• RQ3 在不同 NFE 水平下，噪声尺度（τ）与采样质量（以 FID 衡量）之间的最优权衡是什么？
• RQ4 当扩散模型未完全训练时，所提出方法是否仍保持优越性能，表明其对模型容量的鲁棒性？
• RQ5 无需训练的随机求解器是否能在极少数函数评估下实现最先进 FID 分数？

主要发现

• 在所有基准数据集（CIFAR-10、LSUN-Bedroom、ImageNet）上，SA-Solver 在合适的函数评估次数（NFE）下实现了最先进 FID 分数，优于所有确定性和随机基线方法。
• 在极低 NFE 条件下（如 5–10 次），SA-Solver 的 FID 分数与现有最先进方法（如 DDIM、DPM-Solver 和 UniPC）相当或更优。
• 在 LSUN-Bedroom 256x256 数据集上，SA-Solver 在 NFE=60 时达到最低 FID 3.06，优于所有基线方法，包括 EDM(SDE) 和 UniPC。
• 对于使用分类器自由指导的 ImageNet 512x512 DiT 模型，SA-Solver 在 NFE=60 时生成的样本在多样性和细节上均优于 DDIM 和 UniPC。
• 当模型未完全训练时，SA-Solver 仍保持优于确定性采样器的性能，表明其对模型容量具有鲁棒性。
• 消融实验表明，增加噪声尺度 τ（最高至 τ=1.0）可提升低 NFE 条件下的样本质量，且在大多数基准上，τ=0.8–1.0 时性能最优。