[論文レビュー] 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がハミルトニアン・ジャコビ・ベルマン（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の対数密度を支配するハミルトニアン・ジャコビ・ベルマン（HJB）方程式を導出し、最適制御と接続する。
• 制御理論における検証定理を用いて、拡散モデルにおけるELBOを形式的に導出する。
• 連続時間確率過程の経路上の測度間のKLダイバージェンス最小化として、拡散モデリングを再定式化する。
• 未正規化ターゲット密度からのサンプリングを学習する制御過程を学習する、新しい時系列反転拡散サンプラー（DIS）を提案する。
• ニューラルネットワークを用いて逆時系列SDEのドリフトをパrametrizeし、スコア関数と境界条件に基づく初期化を実施する。
• 柔軟な初期分布と参照SDEを採用し、サンプリングプロセスに事前知識を統合可能にする。

実験結果

リサーチクエスチョン

• RQ1 確率的最適制御と拡散モデルの間の接続が、HJB方程式を用いて正式に確立可能か？
• RQ2 制御理論における検証定理が、拡散ベース生成モデルにおけるELBOを自然に導くか？
• RQ3 経路空間上のKLダイバージェンス最小化が、拡散モデリングの新たな解釈を提供可能か？
• RQ4 未正規化密度からのサンプリングにおいて、新しい時系列反転拡散サンプラー（DIS）が既存手法を上回るか？
• RQ5 マルチモーダルまたは高エネルギー障壁を持つ分布において、DISの性能はULAなどのラングジュアンベース手法と比べて如何にか？

主な発見

• 拡散過程の逆時系列対数密度は、ハミルトニアン・ジャコビ・ベルマン（HJB）方程式を満たし、最適制御と正式に接続される。
• 拡散モデルにおけるエビデンス下限（ELBO）は、制御理論における検証定理の直接的な結果である。
• 拡散ベース生成モデリングは、経路空間上の測度間のKLダイバージェンス最小化として解釈可能である。
• 提案された時系列反転拡散サンプラー（DIS）は、マルチモーダルおよび高エネルギー障壁を持つ分布を含む複数の数値例において、他の拡散ベースサンプリング手法を上回る性能を示した。
• DISは有限時間で正確なサンプリングを達成するが、Kramerの法則が示唆するように、ULAや類似するラングジュアン手法はエネルギー障壁のため指数的混合時間に苦しむ。
• より大きな時間ステップにおいて、DISで学習されたドリフトは最適ドリフトに近づく。また、スコア関数を組み込んだパラメータ化は、より安定的かつ正確な結果をもたらす。