[論文レビュー] Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics

Liming Wu, Zhichao Hou+3 more|arXiv (Cornell University)|May 21, 2024

Advanced Graph Neural Networks被引用数 4

ひとこと要約

本稿では、過去の軌道履歴を活用して非マルコフ型の物理的ダイナミクスをモデル化する、等長性を保つスパatio-時系列グラフネットワークであるESTAGを提案する。周期的パターン抽出に等長離散フーリエ変換（EDFT）を統合し、等長空間モジュール（ESM）と、前方アテンションおよび等長プーリングを備えた等長時系列モジュール（ETM）を組み合わせることで、分子、タンパク質、マクロスケールのダイナミクスベンチマークで最先端の性能を達成した。

ABSTRACT

Learning to represent and simulate the dynamics of physical systems is a crucial yet challenging task. Existing equivariant Graph Neural Network (GNN) based methods have encapsulated the symmetry of physics, \emph{e.g.}, translations, rotations, etc, leading to better generalization ability. Nevertheless, their frame-to-frame formulation of the task overlooks the non-Markov property mainly incurred by unobserved dynamics in the environment. In this paper, we reformulate dynamics simulation as a spatio-temporal prediction task, by employing the trajectory in the past period to recover the Non-Markovian interactions. We propose Equivariant Spatio-Temporal Attentive Graph Networks (ESTAG), an equivariant version of spatio-temporal GNNs, to fulfill our purpose. At its core, we design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic patterns from the history frames, and then construct an Equivariant Spatial Module (ESM) to accomplish spatial message passing, and an Equivariant Temporal Module (ETM) with the forward attention and equivariant pooling mechanisms to aggregate temporal message. We evaluate our model on three real datasets corresponding to the molecular-, protein- and macro-level. Experimental results verify the effectiveness of ESTAG compared to typical spatio-temporal GNNs and equivariant GNNs.

研究の動機と目的

• 未観測の環境的相互作用を捉えられない、従来の物理的ダイナミクスシミュレーション手法におけるマルコフ仮定の制限を解消すること。
• 歴史的軌道データを用いてダイナミクスシミュレーションをスパatio-時系列予測タスクに再定式化し、非マルコフ型依存関係をモデル化すること。
• 時系列的および空間的依存関係をモデル化しつつ、ユークリッド対称性（回転、並進）を保持する等長アーキテクチャを設計すること。
• 新規の等長離散フーリエ変換（EDFT）を用いて周期的パターンモデリングを実現し、物理系における繰り返し現れるダイナミクスを捉えること。
• 時系列記憶と等長性を組み合わせたアプローチが、標準のスパatio-時系列GNNや非等長モデルに比べて優れていることを実証すること。

提案手法

• E(3)の等長性を保ちながら、歴史的ノード軌道からの周波数ドメイン特徴量を抽出できる、新規の等長離散フーリエ変換（EDFT）を提案する。
• ノード座標と特徴量の等長更新を用いて、空間グラフ上でメッセージパッシングを実行する等長空間モジュール（ESM）を設計する。
• 前方アテンションと等長プーリングを備えた等長時系列モジュール（ETM）を導入し、フレーム間のメッセージを集約しながら幾何的対称性を維持する。
• フレーム間仮定を超えて長距離時系列依存関係をモデル化するため、歴史的システム状態からスパatio-時系列グラフを構築する。
• メッセージパッシングおよび集約レイヤーの設計を慎重に行い、すべての学習済み表現と更新が並進、回転、反転に関して等長性を保つようにする。
• ETMで学習可能なアテンション機構を用い、過去のフレームの重要性を動的に重みづけすることで、適応的時系列モデリングを可能にする。

実験結果

リサーチクエスチョン

• RQ1 フレームごとの予測ではなく、スパatio-時系列予測タスクとして物理的ダイナミクスをモデル化することで、未観測の環境的相互作用に起因する非マルコフ型行動をよりよく捉えられるか？
• RQ2 E(3)対称性を保ちつつ、長距離時系列依存関係を効果的にモデル化できる等長スパatio-時系列GNNをどのように設計できるか？
• RQ3 EDFTによる周期的パターンモデリングが、物理的ダイナミクスシミュレーションの性能に与える寄与度はどの程度か？
• RQ4 時系列記憶と幾何的等長性を組み合わせることで、標準のSTGNNや等長性を考慮しないGNNと比較して、一般化性能と精度がどの程度向上するか？
• RQ5 EDFT、アテンション、等長性、時系列プーリングといった個々のモジュールが、全体のモデル性能にどの程度寄与しているか？

主な発見

• ESTAGはMD17、タンパク質、マクロスケールのデータセットで最先端の性能を達成し、アスピリンでは平均MSEが0.063、ベンゼンでは0.003を記錪し、非等長および非時系列ベースラインを著しく上回った。
• アブレーションスタディの結果、EDFTを除去すると平均でMSEが25%上昇し、空間的周期性のモデル化におけるEDFTの重要性が示された。
• ETMにおけるアテンション機構を除去すると、平均でMSEが38%上昇し、動的時系列相互作用を捉えるうえでアテンションの重要性が裏付けられた。
• 非等長バージョンのESTAGでは、平均でMSEが10倍に増加し、幾何的等長性が物理的シミュレーションの堅牢性にとって不可欠であることが証明された。
• 等長時系列プーリングを学習可能な係数に置き換えると、MSEが35%上昇し、構造的時系列集約の利点が確認された。
• L=2層の設定では、エタノールデータセットにおいて、精度と効率の最適なトレードオフが達成され、MSEは0.990×10⁻⁴を記錪した。