[论文解读] An optimal control perspective on diffusion-based generative modeling

Julius Berner, Lorenz Richter+1 more|arXiv (Cornell University)|Nov 2, 2022

Bayesian Methods and Mixture Models被引用 6

一句话总结

本文通过证明扩散模型中的反向时间 SDE 满足 Hamilton–Jacobi–Bellman (HJB) 方程，建立了随机最优控制与基于扩散的生成建模之间的新联系。它利用控制理论中的验证定理推导出证据下界（ELBO），将生成建模重新表述为路径空间 KL 散度最小化问题，并提出一种新型时间反向扩散采样器（DIS），在未归一化密度采样任务中优于现有方法。

ABSTRACT

We establish a connection between stochastic optimal control and generative models based on stochastic differential equations (SDEs), such as recently developed diffusion probabilistic models. In particular, we derive a Hamilton-Jacobi-Bellman equation that governs the evolution of the log-densities of the underlying SDE marginals. This perspective allows to transfer methods from optimal control theory to generative modeling. First, we show that the evidence lower bound is a direct consequence of the well-known verification theorem from control theory. Further, we can formulate diffusion-based generative modeling as a minimization of the Kullback-Leibler divergence between suitable measures in path space. Finally, we develop a novel diffusion-based method for sampling from unnormalized densities -- a problem frequently occurring in statistics and computational sciences. We demonstrate that our time-reversed diffusion sampler (DIS) can outperform other diffusion-based sampling approaches on multiple numerical examples.

研究动机与目标

• 建立随机最优控制与基于扩散的生成模型之间的正式联系。
• 证明扩散模型中的 ELBO 可自然地由控制理论中的验证定理导出。
• 将扩散建模重新解释为在路径空间中最小化 Kullback–Leibler 散度。
• 利用时间反向 SDE 开发一种用于未归一化目标密度的新型采样方法。
• 展示所提出方法（DIS）在实际性能上优于现有基于扩散的采样器。

提出的方法

• 推导出控制反向时间 SDE 的对数密度的 Hamilton–Jacobi–Bellman (HJB) 方程，将其与最优控制联系起来。
• 利用控制理论中的验证定理，正式推导出扩散模型中的 ELBO。
• 将扩散建模重新表述为在连续时间随机过程路径上的测度之间 KL 散度最小化问题。
• 提出一种新型时间反向扩散采样器（DIS），通过学习控制过程实现从未归一化目标密度中采样。
• 使用神经网络参数化反向时间 SDE 的漂移项，初始值基于得分函数和边界条件确定。
• 采用灵活的初始分布和参考 SDE，使先验知识可融入采样过程。

实验结果

研究问题

• RQ1 能否通过 HJB 方程正式建立随机最优控制与扩散模型之间的联系？
• RQ2 控制理论中的验证定理是否能自然导出基于扩散的生成模型中的 ELBO？
• RQ3 路径空间 KL 散度最小化能否为扩散建模提供新的解释？
• RQ4 新型时间反向扩散采样器（DIS）是否能在从未归一化密度中采样时优于现有方法？
• RQ5 DIS 的性能与 ULA 等基于 Langevin 的方法相比如何，特别是在多模态或高势垒分布中？

主要发现

• 扩散过程的反向时间对数密度满足 Hamilton–Jacobi–Bellman (HJB) 方程，从而建立了与最优控制的正式联系。
• 扩散模型中的证据下界（ELBO）是控制理论中验证定理的直接结果。
• 基于扩散的生成建模可被解释为在路径空间上最小化测度之间的 KL 散度。
• 所提出的反向时间扩散采样器（DIS）在多个数值示例中优于其他基于扩散的采样方法，包括多模态和高势垒分布。
• DIS 可在有限时间内实现精确采样，而 ULA 及类似 Langevin 方法由于能量壁垒的存在，其混合时间呈指数增长，这与 Kramer 定律的预测一致。
• DIS 中学习到的漂移项在较大时间步长下更接近最优漂移，且包含得分函数的参数化形式可带来更稳定和精确的结果。