[论文解读] A Geometric Perspective on Diffusion Models

Defang Chen, Zhenyu Zhou+4 more|arXiv (Cornell University)|May 31, 2023

Model Reduction and Neural Networks被引用 5

一句话总结

本文提出了一套几何框架，用于理解扩散模型中基于常微分方程（ODE）的采样机制，揭示了去噪轨迹——由模型得分函数驱动——构成了一条收敛更快的隐式路径，可生成高质量样本。该研究建立了最优ODE采样与均值漂移（mean-shift）算法之间的理论联系，表明每个欧拉步长近似于当前位置与均值漂移目标的凸组合，且证明去噪轨迹在视觉质量和FID指标上始终优于采样轨迹，从而通过ODE-Jump采样实现更快的推理速度。

ABSTRACT

Recent years have witnessed significant progress in developing effective training and fast sampling techniques for diffusion models. A remarkable advancement is the use of stochastic differential equations (SDEs) and their marginal-preserving ordinary differential equations (ODEs) to describe data perturbation and generative modeling in a unified framework. In this paper, we carefully inspect the ODE-based sampling of a popular variance-exploding SDE and reveal several intriguing structures of its sampling dynamics. We discover that the data distribution and the noise distribution are smoothly connected with a quasi-linear sampling trajectory and another implicit denoising trajectory that even converges faster. Meanwhile, the denoising trajectory governs the curvature of the corresponding sampling trajectory and its finite differences yield various second-order samplers used in practice. Furthermore, we establish a theoretical relationship between the optimal ODE-based sampling and the classic mean-shift (mode-seeking) algorithm, with which we can characterize the asymptotic behavior of diffusion models and identify the empirical score deviation. Code is available at \url{https://github.com/zju-pi/diff-sampler}.

研究动机与目标

• 为基于常微分方程的扩散模型采样动力学，特别是方差爆炸型SDE（stochastic differential equations）提供几何理解。
• 揭示一条与显式采样轨迹不同的、收敛更快的隐式去噪轨迹的存在。
• 建立基于ODE的采样与均值漂移（模式搜索）算法之间的理论联系。
• 解释为何去噪轨迹始终生成比采样轨迹更高质量的图像。
• 提出并验证ODE-Jump采样方法，利用去噪轨迹实现更快、更高效的高质量生成。

提出的方法

• 作者分析了方差爆炸型SDE的基于ODE的采样过程，识别出两条关键轨迹：显式采样轨迹与隐式去噪轨迹。
• 推导出去噪轨迹的导数与采样轨迹的二阶导数的负值对齐，将其与二阶数值求解器联系起来。
• 论文建立了最优ODE采样与退火均值漂移算法之间的理论等价性，表明每个欧拉步长均使样本向当前位置与均值漂移目标的凸组合方向移动。
• 提出ODE-Jump采样方法，利用去噪轨迹在采样过程中实现跳跃式前进，从而在保持高质量的同时减少得分函数评估次数（NFEs）。
• 在CIFAR-10、ImageNet、FFHQ及Stable Diffusion等多个数据集上，通过FID与视觉质量对比，对方法进行了验证，结果与DPM-Solver和EDM等标准采样器进行比较。
• 理论分析表明，在温和条件下，去噪轨迹上的样本似然性高于采样轨迹，从而解释了其优越性能的原因。

Figure 1: The geometric picture of ODE-based sampling in diffusion models. An initial sample (from the noise distribution) starts from a big sphere and converges to its final sample (in the data manifold) along a smooth, quasi-linear sampling trajectory. Meanwhile, its denoising output lays in an im

实验结果

研究问题

• RQ1 基于ODE的扩散模型采样动力学背后存在何种几何结构？
• RQ2 去噪轨迹与显式采样轨迹有何不同？为何其在视觉质量上收敛更快？
• RQ3 基于ODE的采样与均值漂移算法之间存在何种理论关系？
• RQ4 为何去噪轨迹始终生成比采样轨迹更高质量的样本？
• RQ5 是否可利用去噪轨迹设计出更快、更高效的采样方法？

主要发现

• 在Fréchet Inception Distance（FID）指标上，去噪轨迹的收敛速度显著快于采样轨迹，且在更少的得分函数评估次数（NFEs）下即可达到更优的FID表现。
• 在CIFAR-10与ImageNet-64数据集上，使用DPM-Solver-3与NFE=48的ODE-Jump采样方法，其FID收敛速度优于标准EDM采样器。
• 去噪轨迹在似然性与视觉质量上始终优于采样轨迹，且在采样早期阶段FID性能提升最高可达30%。
• 理论分析确认，ODE采样中的每个欧拉步长近似于当前样本与均值漂移目标的凸组合，从而将ODE采样与模式搜索行为联系起来。
• 去噪轨迹的导数对应于采样轨迹二阶导数的相反数，从有限差分的角度解释了二阶采样器的起源。
• 实证结果表明，即使得分函数与最优得分存在轻微偏差，其生成能力仍可保持，且能减少模式崩溃现象，这归因于均值漂移类动力学的内在鲁棒性。

(a) The statistics of magnitude.