[論文レビュー] Equivariant Graph Neural Operator for Modeling 3D Dynamics

Minkai Xu, Jiaqi Han+7 more|arXiv (Cornell University)|Jan 19, 2024

Machine Learning in Materials ScienceMaterials Science被引用数 3

ひとこと要約

本稿では、3次元物理的ダイナミクスをSE(3)-等長なニューラルオペレーターとフーリエに基づく時系列畳み込みを用いて連続的軌道としてモデル化する、新規手法Equivariant Graph Neural Operator（EGNO）を提案する。時間関数としてダイナミクスを学習することにより、粒子シミュレーション、人間のモーショングラブキャプチャ、分子動力学的分野で優れた性能を達成し、EGNNなどのベースラインと比較して平均で最大52%の向上を達成する。

ABSTRACT

Modeling the complex three-dimensional (3D) dynamics of relational systems is an important problem in the natural sciences, with applications ranging from molecular simulations to particle mechanics. Machine learning methods have achieved good success by learning graph neural networks to model spatial interactions. However, these approaches do not faithfully capture temporal correlations since they only model next-step predictions. In this work, we propose Equivariant Graph Neural Operator (EGNO), a novel and principled method that directly models dynamics as trajectories instead of just next-step prediction. Different from existing methods, EGNO explicitly learns the temporal evolution of 3D dynamics where we formulate the dynamics as a function over time and learn neural operators to approximate it. To capture the temporal correlations while keeping the intrinsic SE(3)-equivariance, we develop equivariant temporal convolutions parameterized in the Fourier space and build EGNO by stacking the Fourier layers over equivariant networks. EGNO is the first operator learning framework that is capable of modeling solution dynamics functions over time while retaining 3D equivariance. Comprehensive experiments in multiple domains, including particle simulations, human motion capture, and molecular dynamics, demonstrate the significantly superior performance of EGNO against existing methods, thanks to the equivariant temporal modeling. Our code is available at https://github.com/MinkaiXu/egno.

研究の動機と目的

• 3次元物理的ダイナミクスにおける長距離時系列相関を捉えることに限界を示す従来のグラフニューラルネットワークの課題に対処すること。
• 次ステップ予測の制限を克服し、時間全体の解軌道を連続関数としてモデル化すること。
• 物理的一般化に不可欠なSE(3)-等長性を維持しつつ、3次元空間的ダイナミクスにおける時系列依存関係を学習すること。
• 再トレーニングなしに任意の時間離散化に対して効率的かつ並列な推論を可能にし、デプロイの柔軟性を高めること。
• 多様な物理的ダイナミクス応用に適した、任意のEGNNアーキテクチャと統合可能な汎用的フレームワークを開発すること。

提案手法

• 初期状態から全軌道への関数写像として3次元ダイナミクスを定式化し、ニューラルオペレーターを用いた連続時刻モデル化を可能にする。
• 時系列モデリング中にSE(3)-等長性を保つために、周波数ドメインでパrameter化された等長時系列畳み込み層を導入する。
• 周波数空間における回転および並進不変性を維持するように、カーネル積分オペレーターを設計することで、フーリエ変換および逆フーリエ変換の等長性を活用する。
• フーリエに基づく時系列層と標準的なEGNN部品をスタックしてEGNOを構築し、幾何学的および時系列的ダイナミクスのエンドツーエンド学習を可能にする。
• 時間埋め込みを用いて、トレーニング時の時間ステップ間の補間により、任意の時間分解能でのゼロショット推論を可能にする。
• 幾何的特徴（ノード特徴h、位置x、速度v）を時系列畳み込みに統合し、豊かな動的パターンを捉える。

Figure 1: Illustration of EGNO . EGNO blocks (green) can be built with any EGNN layers (blue) and the proposed equivariant temporal convolution layers (yellow). Consider discretizing the time window $\Delta T$ into $P$ points $\{\Delta t_{1},\dots,\Delta t_{P}\}$ . Given a current state ${\mathcal{G

実験結果

リサーチクエスチョン

• RQ1 ニューラルオペレーター枠組みは、連続的時刻3次元ダイナミクスをモデル化する際、SE(3)-等長性を保持できるか？
• RQ2 周波数ドメインにおける時系列畳み込みを組み込むことで、自己回帰的次ステップ予測と比較して、長距離時系列依存関係のモデル化がどの程度向上するか？
• RQ3 再トレーニングなしに、EGNOが未観測の時間離散化ステップにどの程度一般化できるか？
• RQ4 異なる幾何的特徴（h, x, v）が、EGNOにおける時系列モデリング性能に果たす寄与度はどの程度か？
• RQ5 EGNOは、多様な物理的系にわたる軌道精度と一般化性能において、従来のGNNベースのモデルと比較してどのように差をつけるか？

主な発見

• アスピリン分子動力学ベンチマークにおいて、EGNOはEGNN比でF-MSEで36%の相対的改善を達成した。
• 人間のモーショングラブキャプチャ（Mocap-Run）データセットにおいて、EGNOはEGNNと比較して予測精度で平均52%の向上を達成した。
• アブレーションスタディの結果、時系列畳み込みに速度（v）を組み込むことで性能が顕著に向上し、N体シミュレーションにおいて全EGNOバージョン（h, x, vを併用）はアブレーションバージョンを最大15%上回った。
• EGNOは再トレーニングなしに時間分解能を2倍（2P）に一般化するゼロショット推論を可能にし、低分解能時間ステップ間の補間により滑らかで正確な軌道を生成した。
• フーリエモード数（I）は性能に非単調な影響を及ぼし、最適性能はI=2またはI=3で達成された。それ以上になると過学習の可能性が生じる。
• EGNOは、自己回帰的モデルが逐次推論を要するのに対し、一度のフォワードパスで任意の時間窓における将来状態の並列復元を可能にした。

Figure 2: Ablation studies on the number of modes $I$ on N-body simulation and Mocap- Run datasets.