[論文レビュー] Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes

Jaehyeong Jo, Sung Ju Hwang|arXiv (Cornell University)|Oct 11, 2023

Generative Adversarial Networks and Image SynthesisComputer Science被引用数 3

ひとこと要約

本稿では、一般のリーマン多様体上に確率過程を構築する混合エンドポイント条件付きブリッジ過程としての、Riemannian Diffusion Mixture (RFM) という新しい生成モデルフレームワークを提案する。従来の拡散モデルが抱える時間反転とスコアマッチングの落とし穴を回避する。本手法により、はるかに少ない訓練内シミュレーションステップで効率的かつスケーラブルな学習が可能となり、球面、トーラス、ハイパーボロイド、3Dメッシュを含む多様な多様体で最先端の性能を達成する。

ABSTRACT

Learning the distribution of data on Riemannian manifolds is crucial for modeling data from non-Euclidean space, which is required by many applications in diverse scientific fields. Yet, existing generative models on manifolds suffer from expensive divergence computation or rely on approximations of heat kernel. These limitations restrict their applicability to simple geometries and hinder scalability to high dimensions. In this work, we introduce the Riemannian Diffusion Mixture, a principled framework for building a generative diffusion process on manifolds. Instead of following the denoising approach of previous diffusion models, we construct a diffusion process using a mixture of bridge processes derived on general manifolds without requiring heat kernel estimations. We develop a geometric understanding of the mixture process, deriving the drift as a weighted mean of tangent directions to the data points that guides the process toward the data distribution. We further propose a scalable training objective for learning the mixture process that readily applies to general manifolds. Our method achieves superior performance on diverse manifolds with dramatically reduced number of in-training simulation steps for general manifolds.

研究の動機と目的

• リーマン多様体上での既存の生成モデルが、近似された熱核や計算コストの高い発散計算に依存するという制限を解決すること。
• 高次元および一般の多様体設定において、スコアベースの拡散モデルや発散を多く含むCNFモデルのスケーラビリティと幾何的整合性の欠如を克服すること。
• 事前分布からデータ分布への輸送を、リーマンブリッジ過程の混合として直接モデル化する、原理的かつ幾何学的感覚を持つ生成プロセスの開発。
• 時間反転を回避し、シミュレーションステップを削減する二方向ブリッジマッチング目的関数を導入することで、一般多様体上での効率的学習を可能にすること。
• 特に高次元において、複雑な多様体上での正確なモデリングには確率的プロセスが不可欠であり、決定論的フローだけでは不十分であることを実証すること。

提案手法

• 時間反転によるノイズ除去に依存せず、エンドポイント固定のブリッジ過程の混合として生成プロセスをモデル化するリーマン拡散混合フレームワークを提案する。
• 熱核の近似や扱いにくい発散計算を必要としない、一般多様体上での実用的クラスのリーマンブリッジ過程を導入する。
• 前向きおよび後向きのブリッジ過程をマッチングする二方向ブリッジマッチング目的関数を定義し、任意のリーマン幾何に適用可能で、効率的かつスケーラブルな学習を可能にする。
• 多様体の幾何に基づいて最も確率的なエンドポイントに向かって誘導するドリフトを用いて、事前分布（例：多様体上の一様分布）からデータ分布への確率的写像を構築する。
• 連続正規化フロー（CNF）が、本フレームワークの特別な場合として理論的に同等であることを示し、その一般性を検証する。
• ドリフトベクトル場の効率的ニューラルネットワークパラメータ化を用いてJAXで実装し、高次元推論を可能にする。

実験結果

リサーチクエスチョン

• RQ1 ノイズ付与プロセスの時間反転に依存せず、スコアマッチングや発散計算を回避できる、リーマン多様体上での生成モデルが構築可能か？
• RQ2 一般多様体上での既存の拡散モデルやCNFモデルに比べ、リーマンブリッジ過程の混合がより幾何的直感的かつスケーラブルな代替手段となり得るか？
• RQ3 提案された二方向ブリッジマッチング目的関数により、従来手法と比較して、より少ない訓練内シミュレーションステップで高速かつ高精度な学習が可能になるか？
• RQ4 ハイパボロイドや非コンパクト幾何を含む多様な多様体上での、Riemannian Diffusion Mixtureの性能は、最先端のベースラインと比較してどうか？
• RQ5 複雑な多様体上での正確なモデリングには確率的生成プロセスが必要であり、決定論的フローだけでは不十分であるか？

主な発見

• 提案されたRiemannian Diffusion Mixture (RFM) は、球面、トーラス、ハイパーボロイド、3Dメッシュを含む多様な多様体上での合成および実世界データセットにおいて、最先端のベースラインと同等またはそれを上回る性能を達成する。
• スタンフォード・バニー や スポット・ザ・カウ といった一般閉多様体上では、従来のCNFモデルが要する訓練内シミュレーションステップの5%程度で、優れた性能を達成する。
• 高次元へのスケーリングも効果的に可能であり、発散計算に依存するスコアベースの拡散モデルと比較して、学習時間が著しく短縮される。
• 二方向ブリッジマッチング目的関数により、ボルケーノおよびスタンフォード・バニーのデータセットではわずか15ステップのシミュレーションで正確なモデリングが可能であるが、一方で一方的アプローチでははるかに多くのステップを要する。
• ハイパーボロイド上では、RFMが複雑な多モーダル分布（例：ラップドガウス分布の混合）を正しく学習し、ポアンカレ円板への射影により正確な密度推定を可視化する。
• アブレーションスタディにより、均一な時間サンプリングではなく、時間スケーリングされた訓練が、目的関数の設計選択の重要性を示しており、より優れた性能をもたらすことが確認された。