[論文レビュー] Scaling Riemannian Diffusion Models

Aaron Lou, Minkai Xu+1 more|arXiv (Cornell University)|Oct 30, 2023

Machine Learning in Healthcare被引用数 4

ひとこと要約

本稿では、リーマン対称空間を活用することで、熱核量の高精度な計算を可能にするスケーラブルなリーマン多様体拡散モデルを提案する。これにより、スコアマッチィングにおける従来の扱いが困難な近似を克服する。熱核を正確に計算し、確率フロー常微分方程式を用いた正確なサンプリングを可能にすることで、$SU(3)$格子（128次元）や超球的対照埋め込み（127次元）といった高次元多様体において、最先端の性能を達成する。これは、数百年にわたる次元の実世界タスクへの微分方程式ベースの生成モデルの初の成功した応用を可能にする。

ABSTRACT

Riemannian diffusion models draw inspiration from standard Euclidean space diffusion models to learn distributions on general manifolds. Unfortunately, the additional geometric complexity renders the diffusion transition term inexpressible in closed form, so prior methods resort to imprecise approximations of the score matching training objective that degrade performance and preclude applications in high dimensions. In this work, we reexamine these approximations and propose several practical improvements. Our key observation is that most relevant manifolds are symmetric spaces, which are much more amenable to computation. By leveraging and combining various ansätze, we can quickly compute relevant quantities to high precision. On low dimensional datasets, our correction produces a noticeable improvement, allowing diffusion to compete with other methods. Additionally, we show that our method enables us to scale to high dimensional tasks on nontrivial manifolds. In particular, we model QCD densities on $SU(n)$ lattices and contrastively learned embeddings on high dimensional hyperspheres.

研究の動機と目的

• リーマン多様体拡散モデルにおけるノイズ除去スコアマッチィングの非可解性が、複雑な幾何的構造に起因することを解決すること。
• 物理学および表現学習に特に関連する高次元多様体上での、スケーラブルかつ高精度なリーマン多様体拡散モデルの訓練を可能にすること。
• 高次元へのスケーラビリティを制限する、従来のスコアマッチィングにおける近似の欠陥を克服すること。
• 微分方程式ベースの生成モデルが、$SU(n)$格子や超球的埋め込みといった実世界の高次元多様体に実際に適用可能であることを示すこと。
• 真の超球的分布をモデル化することで、対照学習における解釈可能性と性能を向上させること、ペンルティメートレイヤー特徴量に依存しないこと。

提案手法

• リーマン対称空間を活用し、ノイズ除去スコアマッチィングに不可欠な熱核の効率的かつ高精度な計算を可能にする。
• 既知の熱核技術をリーマンスコアマッチィングの文脈に適応させ、高速かつ安定した訓練を可能にする。
• 最大トーラス上での標準的な常微分方程式ソルバーを用いて解ける確率フロー常微分方程式を用いて、熱核からの正確なサンプリング手法を開発する。
• 最大トーラスを用いることで計算複雑度を低減し、対称空間上でのシミュレーションフリーな効率的サンプリングを実現する。
• 正確な熱核推定器をノイズ除去スコアマッチィングの目的関数に統合し、訓練の安定性と性能を向上させる。
• $4\times4$ $SU(3)$格子におけるウィルソン作用の学習や、$S^{127}$上での対照学習埋め込みのモデル化といった高次元タスクに本手法を適用する。

(a) $S^{2}$

実験結果

リサーチクエスチョン

• RQ1 リーマン多様体拡散モデルは、$SU(n)$格子や高次元超球といった高次元多様体にスケーリング可能か？
• RQ2 対称空間上での正確な熱核計算により、スコアマッチィングにおける不正確な近似の必要性が排除可能か？
• RQ3 確率フロー常微分方程式による正確なサンプリングは、多様体上でのモデル性能と訓練安定性を向上させるか？
• RQ4 リーマン多様体拡散モデルは、対照的に学習された超球的埋め込みから忠実な分布を回復でき、解釈可能性とOOD検出性能を向上させるか？
• RQ5 微分方程式ベースのアプローチを用いて、非自明な多様体上での高次元生成モデルで最先端の性能を達成可能か？

主な発見

• $4\times4$ $SU(3)$格子上でウィルソン作用を学習する際、有効サンプルサイズが0.62に達し、128次元という高次元多様体を対象としている。
• 127次元超球的対照学習埋め込みにおいて、本リーマン多様体拡散モデルは、ペンルティメートレイヤー特徴量を用いた最先端手法と同等またはそれ以上のOOD検出性能を達成した。
• 平坦なトーラス上での複雑なチェッカーボード分布の学習において、本モデルはRFMやフローマッチングを上回り、複雑度の増加に伴う性能劣化を回避した。
• 正確な熱核推定器がなければ、ノイズ除去スコアマッチィングは劣悪な結果を生じるため、正確な幾何的計算の必要性が明確に示された。
• 本手法により、数百年にわたる次元の実世界タスクに、微分方程式ベースの生成モデルの初の成功した応用が可能になった。
• 本モデルは、対照学習における崩壊したプロジェクションヘッド表現を回復させ、真の超球的分布をモデル化することで解釈可能性を向上させた。

(b) $S^{127}$