[論文レビュー] Equivariant flow matching

Leon Klein, Andreas Krämer+1 more|arXiv (Cornell University)|Jun 26, 2023

Topic Modeling被引用数 6

ひとこと要約

本稿では、等長的連続正規化フロー（CNF）のための新たな学習目的、等長的フローマッチングを提案する。この手法は最適輸送の原則を活用し、シミュレーションを伴わない効率的な学習を可能にする。回転や置換不変性といった物理的対称性を活用することで、より短い統合経路、より優れたサンプリング効率、および高いスケーラビリティを達成する。LJ13 やアラニンジペプチドといった系で、特化した内部座標に依存せずに検証された。

ABSTRACT

Normalizing flows are a class of deep generative models that are especially interesting for modeling probability distributions in physics, where the exact likelihood of flows allows reweighting to known target energy functions and computing unbiased observables. For instance, Boltzmann generators tackle the long-standing sampling problem in statistical physics by training flows to produce equilibrium samples of many-body systems such as small molecules and proteins. To build effective models for such systems, it is crucial to incorporate the symmetries of the target energy into the model, which can be achieved by equivariant continuous normalizing flows (CNFs). However, CNFs can be computationally expensive to train and generate samples from, which has hampered their scalability and practical application. In this paper, we introduce equivariant flow matching, a new training objective for equivariant CNFs that is based on the recently proposed optimal transport flow matching. Equivariant flow matching exploits the physical symmetries of the target energy for efficient, simulation-free training of equivariant CNFs. We demonstrate the effectiveness of flow matching on rotation and permutation invariant many-particle systems and a small molecule, alanine dipeptide, where for the first time we obtain a Boltzmann generator with significant sampling efficiency without relying on tailored internal coordinate featurization. Our results show that the equivariant flow matching objective yields flows with shorter integration paths, improved sampling efficiency, and higher scalability compared to existing methods.

研究の動機と目的

• 対称的な多体系における等長的連続正規化フロー（CNF）の学習および推論の高い計算コストを低減すること。
• 最適輸送（OT）マップを活用しつつ、回転や置換不変性といった物理的対称性を尊重する、シミュレーションを伴わない学習目的の開発。
• 統計力学のアプリケーションにおけるボルツマンジェネレータモデルのサンプリング効率を向上させるとともに、統合経路長を短縮すること。
• システム固有の特徴抽出に依存せずに、LJ55 やアラニンジペプチドといった大規模で対称性を持つ系へのスケーラブルな学習を可能にすること。
• 対称分布に対する尤度ベース学習や標準的なOTフローマッチングの新しい効率的代替手法を提供すること。

提案手法

• 事前分布とターゲット分布の間で対称性に配慮した輸送経路を強制する、等長的最適輸送（OT）フローマッチング目的関数を提案する。
• ターゲットエネルギー関数の対称性を尊重する群に等長的なニューラルネットワークを組み込むことで、不変密度に適応したOTフローマッチングフレームワークを拡張する。
• バッチ再順序付けと最適輸送ペアの並列事前生成を用いて、学習中の計算ボトル neck を軽減する。
• ボルツマン分布に基づく重み付きサンプリング戦略を採用し、低確率領域（例：アラニンジペプチドのφ二重角遷移）における密度推定を改善する。
• 輸送コストを最小化するとともに、関連する対称性群（例：SO(D) および S_N）の下でフローが等長的であることを保証する、対称性に配慮した損失関数を導入する。
• 残留ブロックと等長的メッセージパッシングを備えた連続正規化フローのアーキテクチャを用いて、不変分布に対するCNFの学習を実施する。

実験結果

リサーチクエスチョン

• RQ1 不変分布において、物理的対称性を尊重しつつ、近似的に最適な輸送経路を達成するフローマッチング目的関数を設計可能か？
• RQ2 標準的なOTフローマッチングや尤度ベース学習と比較して、等長的フローマッチングは統合経路長を短縮し、サンプリング効率を向上させるか？
• RQ3 内部座標特徴抽出に依存せずに、LJ55 やアラニンジペプチドといった大規模で対称性を持つ系に効果的にスケーリング可能か？
• RQ4 輸送コストと経路精度の観点から、標準的なOTフローマッチングと比較して等長的フローマッチングの性能はどのように異なるか？
• RQ5 遅い遷移を示す系（例：アラニンジペプチドのφ二重角）において、訓練データのバイアス化が自由エネルギー推定を改善できるか？

主な発見

• 等長的フローマッチングは、標準的なOTフローマッチングと比較して顕著に短い統合経路を達成し、LJ13 やアラニンジペプチドでは弧長を最大30％短縮した。
• LJ13 では学習時間を0.5時間、アラニンジペプチドでは25時間まで短縮した。これは尤度ベース学習（それぞれ13時間および17時間）と比較して顕著に高速である。
• LJ55 系では、バッチサイズ256で17時間の学習時間を達成し、大規模な系へのスケーラビリティを示した。
• アラニンジペプチドでは、200エポックで有効サンプルサイズ1×10⁵を達成し、安定した収束と改善されたサンプリング効率を示した。
• 等長的OTフローマッチングの統合経路長は、特に高対称性系において、理論的最適輸送経路に常に標準的なOTフローマッチングよりも近づいている。
• 等長的OTフローマッチングのバッチ再順序付けは、CPU上で1バッチあたりたった22.4秒で実行可能であり、効率的な並列化と大規模系へのスケーラビリティを可能にした。