[論文レビュー] Diffusion Models for Constrained Domains

Nic Fishman, Leo Klarner+3 more|arXiv (Cornell University)|Apr 11, 2023

Reservoir Engineering and Simulation Methods被引用数 5

ひとこと要約

本稿では、測地線が定義されていないため従来のリーマン多様体拡散手法が失敗する不等式制約によって定義される制約付き多様体のための2つの新しい拡散モデルを紹介する。本稿では対数バリアに基づくノイズプロセスと反射ブラウン運動アプローチを提案し、凸多面体や対称正定値行列上でのスコアベース生成モデリングを可能にした。実験的に、タンパク質骨格およびロボットアームのコンformationalモデリングにおいて成功を収めた。

ABSTRACT

Denoising diffusion models are a novel class of generative algorithms that achieve state-of-the-art performance across a range of domains, including image generation and text-to-image tasks. Building on this success, diffusion models have recently been extended to the Riemannian manifold setting, broadening their applicability to a range of problems from the natural and engineering sciences. However, these Riemannian diffusion models are built on the assumption that their forward and backward processes are well-defined for all times, preventing them from being applied to an important set of tasks that consider manifolds defined via a set of inequality constraints. In this work, we introduce a principled framework to bridge this gap. We present two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints. As existing diffusion model techniques cannot be applied in this setting, we derive new tools to define such models in our framework. We then demonstrate the practical utility of our methods on a number of synthetic and real-world tasks, including applications from robotics and protein design.

研究の動機と目的

• 既存のリーマン拡散モデルの制限を解決する。これらは全空間に定義された測地線を必要とし、不等式制約によって定義される多様体では失敗する。
• 凸多面体、対称正定値行列、およびロボット工学やタンパク質設計における多様体を含む制約付きドメインにおける生成モデリングを可能にする。
• 時間反転およびスコアマッチングフレームワークに適合する、原理的で整合性のあるノイズプロセス（対数バリアおよび反射ブラウン運動）を開発する。
• 非測地線的完全でない多様体への拡散モデルの理論的・実用的枠組みを拡張し、実世界の科学的分野への応用可能性を広げる。
• 合成的および実世界のタスクを通じて、高次元の制約付き多様体における提案モデルのスケーラビリティと柔軟性を実証する。

提案手法

• 不等式制約を尊重し、定義された拡散経路を保証するリーマン構造を導く対数バリアメトリクスに基づく前向きノイズプロセスを設計する。
• 制約付きドメインの境界で軌道を反射する反射ブラウン運動に基づくノイズプロセスを導入し、多様体の幾何構造を保持する。
• 非ユークリッド的かつ制約付き設定へのスコアマッチングおよびスティンスコア推定を拡張することで、両方のノイズプロセスの時間反転SDEを導出する。
• 制約付き多様体上でのノイズ除去目的を用いて、ニューラルネットワークで後退ドリフト（スティンスコア）を近似することで、スコアベース生成モデルを構築する。
• 温度パラメータを用いた低温サンプリングにより、生成分布における忠実度と多様性のトレードオフを制御する。
• 実データと生成データの分布類似度を評価するために、最大平均差分（MMD）のU統計量近似を用いる。

Figure 1 : The behaviour of different types of noising processes considered in this work defined on the unit interval. $\mathbf{B}_{t}$ : Euclidean (unconstrained) Brownian motion. $\hat{\mathbf{B}}_{t}$ : Log barrier forward noising process. $\bar{\mathbf{B}}_{t}$ : Reflected Brownian motion. All s

実験結果

リサーチクエスチョン

• RQ1 測地線がすべての時間で定義されていない、不等式制約によって定義される多様体に対しても、拡散モデルを拡張可能か？
• RQ2 幾何構造を保持し、時間反転を可能にする、制約付き多様体のための原理的ノイズプロセスをどのように構築できるか？
• RQ3 対数バリアメトリクスおよび反射ブラウン運動は、標準的な測地線ベースノイズと比較して、拡散モデリングの代替手段として有効か？
• RQ4 Birkhoff多面体やタンパク質コンformational空間などの高次元制約付きドメインにおいて、提案モデルの性能はいかがなものか？
• RQ5 ロボット工学および計算生物学の実世界応用に、モデルはどの程度一般化可能か？

主な発見

• 提案された対数バリアおよび反射ブラウン運動ノイズプロセスは、標準的リーマン拡散が失敗する制約付き多様体上での拡散モデリングに成功した。
• 超立方体、単体、Birkhoff多面体における実験的結果から、MMD値が低く抑えられ、実データと生成データの分布が良好に一致していることが示された。
• サイズ $ n=4 $ のBirkhoff多面体において、モデルは高い忠実度で有効な二重確率行列を生成したことが、ペairwiseおよび周辺分布の可視化により確認された。
• 本手法は高次元設定にも一般化可能であり、$ n=10 $ のタンパク質骨格コンformationやロボットアーム構成においてもスケーラビリティを示した。
• 温度パラメータ $ heta_0 \in \{1,2,3,5,10\} $ を用いた低温サンプリングにより、サンプルの多様性と忠実度の制御が可能であり、$ \theta_0 $ が大きいほど鋭く集中したサンプルが得られた。
• スコアネットワークが予測するドリフトが時間経過とともに適切に変化することが、単体上での可視化で確認され、安定した学習と意味のあるスコア関数の学習が行われたことが裏付けられた。

Figure 2 : A convex polytope defined by six constraints $\{f_{i}\}_{i\in\mathcal{I}}$ , along with the log barrier potential, and ‘straight trajectories’ under the log-barrier metric and under the Euclidean metric with and without reflection at the boundary.