[论文解读] The probability flow ODE is provably fast

Sitan Chen, Sinho Chewi+4 more|arXiv (Cornell University)|May 19, 2023

Markov Chains and Monte Carlo Methods被引用 5

一句话总结

该论文通过引入基于非阻尼朗之万扩散的新型校正步骤，首次为基于分数的生成建模中的概率流ODE提供了多项式时间收敛保证。其维度依赖性达到$O(\sqrt{d})$，优于先前DDPM分析中的$O(d)$界，并通过Wasserstein–到–TV正则化与随机校正，建立了ODE框架在经验上速度优势的理论基础。

ABSTRACT

We provide the first polynomial-time convergence guarantees for the probability flow ODE implementation (together with a corrector step) of score-based generative modeling. Our analysis is carried out in the wake of recent results obtaining such guarantees for the SDE-based implementation (i.e., denoising diffusion probabilistic modeling or DDPM), but requires the development of novel techniques for studying deterministic dynamics without contractivity. Through the use of a specially chosen corrector step based on the underdamped Langevin diffusion, we obtain better dimension dependence than prior works on DDPM ($O(\sqrt{d})$ vs. $O(d)$, assuming smoothness of the data distribution), highlighting potential advantages of the ODE framework.

研究动机与目标

• 在最小假设下，为基于分数的生成建模中的概率流ODE提供首个严格的收敛保证。
• 克服基于ODE的生成模型中非收缩确定性动力学导致的误差累积挑战。
• 证明在相似条件下，ODE框架的维度依赖性优于基于SDE的DDPMs。
• 基于Wasserstein–到–TV正则化，构建一个新理论框架，用于分析具有随机校正的确定性动力学。
• 阐明校正步骤是否对收敛性至关重要，以及其是否与经验性能一致。

提出的方法

• 提出一种混合预测器–校正器算法：预测器步骤遵循概率流ODE，校正器步骤则使用估计分数的非阻尼朗之万扩散。
• 设计一种基于非阻尼朗之万过程的特殊校正步骤，以引入随机性，从而在非收缩动力学中实现误差控制。
• 开发一种Wasserstein–到–TV正则化技术，通过Wasserstein距离界定向生成分布与目标分布之间的总变差距离。
• 采用两阶段分析：第一阶段为步长几何递减的预测器阶段；第二阶段为最终校正阶段，以精炼输出分布。
• 应用分数扰动引理，并利用Ornstein–Uhlenbeck过程的收缩性质来控制误差传播。
• 在最小假设下建立收敛性：无需对数凹性或$L^\infty$分数误差界，仅需$L^2$有界分数估计误差。

实验结果

研究问题

• RQ1 在最小假设下，能否证明基于分数的生成建模中概率流ODE的多项式时间收敛性？
• RQ2 尽管缺乏随机性，为何ODE框架在速度上经验上优于基于SDE的DDPMs？
• RQ3 校正步骤是否对收敛性必不可少？能否用于实现优于先前DDPM分析的维度依赖性？
• RQ4 基于确定性ODE的方法能否实现$O(\sqrt{d})$的维度依赖性，与SDE的最佳已知结果相匹配？
• RQ5 能否构建一个新理论框架，用于分析生成建模中非收缩的确定性动力学？

主要发现

• 采用非阻尼朗之万校正步骤的概率流ODE实现了$O(\sqrt{d})$的维度依赖性，优于先前DDPM分析中的$O(d)$依赖性。
• 该方法在最小假设下提供了多项式时间收敛保证：无需对数凹性或$L^\infty$分数误差界。
• 生成分布与目标分布之间的总变差距离随时间呈指数衰减，衰减速率取决于$\sqrt{d} + \mathfrak{m}_2$，其中$\mathfrak{m}_2$为目标分布的二阶矩。
• 校正步骤对证明至关重要，因为它使随机正则化得以应用，从而控制非收缩动力学中的误差累积。
• 数值实验表明，该方法即使从低概率成分出发，也能成功恢复复杂且非对数凹的目标分布，如高斯混合分布。
• 基于Wasserstein–到–TV正则化的理论框架本身具有独立兴趣，可能适用于其他确定性生成模型。