Riemannian Diffusion Models on General Manifolds via Physics-Informed Neural Networks

技術要約：物理情報ニューラルネットワークを用いた一般多様体上のリーマン拡散モデル

1. 問題提起

リーマン拡散モデルは、リーマン・ブラウン運動を通じてフォワード拡散過程を定義することにより、滑らかな多様体上に存在するデータに対してスコアベースの生成モデリングを一般化するものである。これらのモデルの学習とサンプリングにおける決定的なボトルネックは、多様体の熱核（heat kernel）（ブラウン運動の遷移密度）とその勾配へのアクセスが必要であることである。具体的には、学習目的関数（デノイジング・スコア・マッチング）は条件付きスコア $\nabla_x \log p_{t|0}(x|x_0)$ の評価を必要とし、フォワードサンプリングには熱核からのサンプリングが必要となる。

熱核は、高度に対称的な多様体（ユークリッド空間、球面、トーラスなど）では閉形式で知られているが、一般的なリーマン多様体においては解析的に得られることは稀である。既存の近似手法（Varadhanの公式などの短時間漸近展開やスペクトル展開）は、特定の領域（微小な時間 $t$）に制限されているか、あるいは強い構造的仮定（コンパクトなリー群など）を必要とするため、多様な幾何学的設定における実用的な適用可能性が制限されている。

2. 手法

著者らは、**物理情報ニューラルネットワーク（PINNs）**を用いて熱方程式を直接解くことにより、熱核を近似する一般的なフレームワークを提案している。この手法は、主に以下の4つの段階で構成される。

A. 座標系の選択と埋め込み
局所的な座標チャートのパッチを構成する代わりに、本手法ではターゲットとなる多様体 $M$ を、より大きな周囲の多様体 $\tilde{M}$（通常は $\mathbb{R}^N$）へ、部分多様体または商多様体として埋め込む。

• 部分多様体: 制約式 $f_i(x)=0$ によって定義される。計量および微分演算子は、周囲の空間から $M$ の接空間へと写像する射影行列 $P(x)$ を用いて導出される。
• 商多様体: 群作用（置換対称性など）の下で点を同一視することで定義される。本手法では、軌道の代表点（canonical representatives）を用いることで、軌道の点を同一の状態として扱う。

B. 熱方程式の導出
周囲の座標と射影行列を用いて、ラプラス・ベルトラミ演算子 $\Delta_M$ と対応するフォッカー・プランク（熱）方程式を明示的に導出する：
$$\partial_t p_{t|0}(x|x_0) = \Delta_M p_{t|0}(x|x_0)$$
この方程式は、周囲の計量と射影行列のみを用いて表現されるため、固有の座標チャートを必要とせずに座標ベースの計算が可能となる。

C. 初期条件の近似
ディラックのデルタ初期条件 $p_{0|0} = \delta_{x_0}$ は数値的に課すことができないため、著者らは原理的な短時間漸近近似を用いて、微小な時間 $t_0 > 0$ におけるPINNを初期化する。具体的には、パラメトリックス法（Minakshisundaram–Pleijel再帰）を利用して、熱核の高次展開を得る：
$$p_{t_0}(x|x_0) \approx G_t^{(n)}(x, x_0) = G_t(x, x_0) \left( u_0 + u_1 t + \dots + u_{n-1} t^{n-1} \right)$$
ここで、$G_t$ はガウス型項であり、$u_k$ は幾何学的特性（体積の歪みや曲率など）から導かれる係数である。

D. PINNの学習
ニューラルネットワーク $\phi_\theta(t, x)$ が、対数熱核 $\log p_{t|0}$ を近似するように学習される。学習目的関数は、以下の組み合わせを最小化する：

1. 初期/境界損失: ネットワークの出力と $t=t_0$ における漸近的初期化関数の差。
2. PDE残差損失: 対数密度に関する熱方程式 $\partial_t \phi - \Delta_M \phi - \|\nabla \phi\|^2 = 0$ の残差。

安定性を確保するため、適応的な損失重み付けとカリキュラム学習（時間間隔を段階的に拡大する手法）を採用している。

E. 拡散モデルへの統合
学習されたPINNは熱核のサロゲート（代理モデル）として機能し、以下の2つの主要な操作を可能にする：

1. フォワードサンプリング: PINNによって評価された密度を用い、マルコフ連鎖モンテカルロ（MCMC）法（測地線ランダムウォークなど）を通じてサンプルを生成する。
2. スコア評価: PINNを微分することで条件付きスコア $\nabla_x \log p_{t|0}$ を計算し、これをデノイジング・スコア・マッチング（DSM）を通じて逆時間拡散モデルの学習に使用する。

3. 主な貢献

• 汎用的なPDEベースのサロゲート: 熱方程式を数値的に解くことで、任意の多様体上の熱核を近似する手法を導入し、閉形式の解や制限的なスペクトル仮定を回避した。
• 座標フリーの導出: ユーークリッド空間の部分多様体または商として定義される多様体に対して、熱方程式を体系的に導出するフレームワークを提供しており、これにより $S^2, SO(3), SPD(n)$, および置換商を用いた点群のような複雑な幾何学への適用が可能となった。
• 安定性解析: 著者らは、近似的な初期条件から熱方程式を解くことが、すべての $t \ge t_0$ において誤差が有界な解を与えることを示す理論的根拠（定理4.1）を提示し、PINNサロゲートの信頼性を保証している。
• 置換不変拡散: 置換商多様体 $\mathbb{R}^{k \times n}/S_n$ に本手法を適用し、原子の置換を明示的に排除する拡散プロセスを実現した。これは、標準的なユークリッド拡散では構造的対称性を維持できない設定である。

4. 実験結果

本手法は、多様な多様体とデータセットを用いて評価された：

• 球面データ ($S^2$): 地球物理学的イベント（火山、地震、洪水、森林火災）の密度推定において、Varadhan近似やスペクトル法に依存するリーマン・スコアベース生成モデル（RSGM）と比較して、競争力のある、あるいは優れたテスト対数尤度を達成した。
• 回転データ ($SO(3)$): 合成混合モデルにおいて、RSGMベースラインよりも改善された対数尤度を示した。
• 正定値行列 ($SPD(n)$):
  • 交通分析 ($SPD(10)$): ニューヨークのタクシーデータからの交通流行列の条件付き生成。本手法は、SPD-DDPMと比較して、真値に対するフロベニウス距離を減少させた。
  • 脳結合性 ($SPD(n)$): EEG結合行列の条件付き生成。本手法は、DiffeoCFM（ユークリッド空間への微分同相写像を用いる手法）と競合する高い精度とF1スコアを達成し、LowTriDDPMを上回った。
• 分子生成 ($\mathbb{R}^{k \times n}/S_n$): 置換商空間における分子の生成。本手法は、分散爆発型（VE）EDMベースラインよりも高い妥当性（validity）と安定性を持つ分子を生成したが、新規性（novelty）については低くなっており、既知の妥当性と新規性のトレードオフを反映している。

検証: 著者らは、学習された対数密度とスコアを、利用可能な場合は参照近似と比較し、さらに正規化されたPDE残差と境界誤差を報告することで、PINNサロゲートが熱方程式を満たしていることを確認し、その妥当性を検証した。

5. 意義と主張

本論文は、PDEベースの熱核近似が、リーマン多様体上の拡散における既存手法に対する実用的かつ汎用的な代替案となることを主張している。熱核を既知のPDEの解として扱うことで、このアプローチは、解析的な熱核が存在しない多様体へとドリフトフリーのリーマン拡散の適用範囲を広げている。

著者らは、本研究を「熱方程式を解く」という視点によるものであり、単純な幾何学を超えて多様体拡散の適用可能性を広げるものであると位置づけている。ただし、以下の限界についても謙虚に認めている：

• 計算コスト: 高次の自動微分とコロケーション評価が必要なため、PINNの学習には多大なリソースを要する。
• 安定性: 高次元や複雑な幾何学においては学習が不安定になる可能性があり、より強力な帰納バイアスが必要となる場合がある。
• 複雑性: 熱核の推定を全く行わないFlow Matchingなどの代替手法よりも、パイプラインが複雑である。

本研究はFlow Matchingに取って代わることを目的としているのではなく、むしろ、測地線補間や指数写像への依存が不十分な場合や、熱核の特定の構造が必要とされる場面において、リーマン多様体上のスコアベースモデルに不可欠なバックエンドを提供するものである。