技術要約：Graph Mamba Operator (GraMO)

問題提起

N体問題、連結体、あるいは変形構造といった相互作用する動的システムのモデリングには、空間的な相互作用と長距離の時系列依存性の両方を捉えることが求められる。既存のアプローチには、主に以下の2つの制限がある：

1. 自己回帰的な誤差の累積： グラフニューラルネットワーク（GNN）は通常、空間的ダイナミクスと時間的ダイナミクスを別々に扱う自己回帰的なロールアウトに依存している。これにより、長いホライゾンにおいて予測誤差が累積する。
2. 限定的なスコープ： 現在の手法は、局所的な相互作用や短い時間コンテキストに焦点を当てることが多く、マルチホップの依存関係やグローバルな構造的ダイナミクスを効果的に捉えることができない。

状態空間モデル（SSM）は、線形コストで長距離の時系列メモリを実現するメカニズムを提供するが、グラフにSSMを適用する場合、人工的な順序付け（ノードのソートやランダムウォークなど）を課す必要があり、これが粒子間の相互作用を支配する固有のトポロジーを破壊してしまうことが多い。逆に、標準的なSSMは入力を非構造的なシーケンスとして扱うため、物理システムに不可欠なグラフ構造を無視してしまう。

手法：Graph Mamba Operator (GraMO)

著者らは、単一のリカレンス（回帰）内で状態空間モデルとグラフベースの相互作用学習を統合した潜在空間シミュレータであるGraMOを提案する。先行研究とは異なり、ノードをシーケンス化したり、空間更新と時間更新を別々のステージで適用したりするのではなく、GraMOはグラフの伝播と時間的な状態進化を同時に結合させる。

コアとなる定式化

GraMOは、潜在状態 $H(t) \in \mathbb{R}^{V \times D \times N}$ （ここで $V$ はノード数、$D$ は特徴量次元、$N$ はメモリ次元）上で動作する。連続時間のダイナミクスは以下のように定義される：
$$\frac{dH(t)}{dt} = A \odot (\tilde{L} H(t) W_s) + B(t) \odot X(t)$$
ここで：

• $\tilde{L}$ は、空間情報を伝播させるためにノード次元に作用する正規化された隣接行列（自己ループを含む）。
• $W_s$ は、状態次元に作用する学習可能な空間混合行列。
• $A$ は、安定した時間的状態の減衰を制御する（HiPPOによって初期化される）。
• $B(t)$ は、現在の観測値 $X(t)$ を潜在状態に注入する入力依存の射影。
• 出力は、縮退 $Y(t) = H(t) \cdot C(t)$ によって得られる。

主要なアーキテクチャの特徴

1. 結合リカレンス（Joint Recurrence）： 各ステップにおける更新は潜在状態 $H(t)$ に対して線形であるが、演算子の係数（$A, B, C, \Delta t$）は入力依存である。これにより、システムは異なるレジーム（接触、緩やかなドリフト、高速な遷移など）に適応でき、潜在状態をリセットすることなく、自己回帰的な予測器に典型的な誤差の蓄積を回避できる。
2. 離散化： モデルは、学習可能なタイムステップ $\Delta t$ を持つゼロ次ホールド（ZOH）離散化を使用する。更新ルールは以下のようになる：
   $$H_{t+1} = \bar{A}_t \odot (\tilde{L} H_t W_s) + \bar{B}_t \odot X_t$$
   この構造は、状態空間モデルとグラフ・ラプラシアン演算子を統合し、連続時間潜在ダイナミクス内での原理に基づいたメッセージパッシングを可能にする。
3. 選択機構（Selectivity Mechanism）： パラメータ $B_t, C_t,$ および $\Delta t$ は、グラフニューラルネットワークを介して入力 $X_t$ から動的に生成される。これにより、コンテンツに応じたゲーティング（どの入力情報が状態に入るかを選択する）と、非定常な適応（現在のグラフ状態に基づいて時間的ダイナミクスを調整する）が可能になる。
4. 双方向処理： 本アーキテクチャは双方向SSM定式化を採用しており、より豊かな軌跡のコンテキストを捉えるために、順方向と逆方向の両方の時間方向にグラフ結合更新を適用する。

理論的特性

論文では、GraMOが入力依存の係数を持つ**グラフ自己回帰移動平均（ARMA）**表現を許容することを確立している。理論的分析により、以下が確認されている：

• 安定性： 正規化されたラプラシアンの有界なスペクトルと、HiPPOで初期化された時間演算器により、漸近的な長距離の情報フローが保証される。
• 潜在空間における線形性： 更新は持ち上げられた潜在状態に対して線形であり、演算子が現在のグラフ状態に適応する時変クープマン・ダイナミクス（time-varying Koopman dynamics）との関連性を示している。

実験結果

GraMOは、N体シミュレーション、ロボティクス、およびモーションキャプチャの3つの異なるドメインで評価された。

ベンチマーク

• N体システム： 電荷を持つ粒子、バネのダイナミクス、重力システム。
• ロボティクス： ロープ制御、ソフトロボットの操作、泳ぐソフトロボット。
• モーションキャプチャ： CMUデータセット（歩行および走行シーケンス）。

パフォーマンス

GraMOは、すべてのベンチマークと指標（平均二乗誤差 - AMSE、および最終二乗誤差 - FMSE）において、ST-GNN、SE(3)-Transformers、Koopmanベースの手法、およびGraphMambaを含む強力なベースラインを上回り、最も低い誤差を達成した。

• N体： 重力システムにおいて、GraMOはFMSEを6.020 (GeoTDM) から3.720へと減少させた。バネのダイナミクスでは、FMSEを0.063 (ESTAG) から0.024へと減少させた。
• ロボティクス： ロープのデータセットにおいて、GraMOはFMSEを0.5380 (GeoTDM) から0.4120へと改善した。
• モーションキャプチャ： GraMOは、走行シーケンスにおける平均FMSEを0.508 (NS-EGNN) から0.161へと減少させ、これは7.3%の相対的な誤差削減に相当する。

アブレーション研究

アブレーション実験により、各コンポーネントの必要性が確認された：

• 時間的モデリングを除去すると、最大のパフォーマンス低下が生じた。
• 双方向SSMを除去すると、特にモーション予測において性能が低下した。
• 入力依存の選択機構（パラメータを固定）を除去すると、誤差が大幅に増加し、適応的なダイナミクスの必要性が浮き彫りになった。
• HiPPO初期化を置き換えると、長距離メモリ能力が低下した。

汎化性能と効率性

• ゼロショット汎化： GraMOは、未知のシステムサイズに対する汎化能力を示した（例：$N \in \{5, 7, 8\}$ の粒子で訓練し、$N=9$ でテスト）。これは組成的なKoopman演算器よりも優れた性能を示した。
• 効率性： より単純な等変モデルは訓練速度は速いが、予測性能は低い。GraMOは、EqMotionやSE(3)-Transformerのような複雑なベースラインと比較して、競争力のある訓練時間を維持しつつ、より低いFMSEを達成するという良好なトレードオフを提供している。

意義と主張

著者らは、GraMOがグラフベースの相互作用学習と状態空間の進化を統合することで、複雑な時空間ダイナミクスを学習するための原理的なフレームワークを提供すると主張している。主な主張は以下の通りである：

1. 結合メカニズム： グラフの伝播と時間的状態の更新を単一のリカレンス内で結合させた初の潜在シミュレータであり、空間的相互作用と時間的メモリが共同で進化することを可能にする。
2. 誤差の軽減： 観測空間で再帰的に予測を行うのではなく、永続的な潜在状態を進化させることで、GraMOは自己回帰的なロールアウトの特徴である誤差の蓄積を回避し、優れた長期間予測を実現する。
3. トポロジーの保存： ノードを直列化する手法（GraphMambaなど）とは異なり、GraMOはラプラシアンベースの伝播を通じて固有の相互作用トポロジーを保持しており、これは相互作用する粒子ダイナミクスのモデリングにおいて極めて重要である。
4. 適応性： 入力依存の係数により、モデルはレジーム切り替えシステム（例：接触 vs 自由飛行）を効果的に扱うことができる。

論文は、GraMOがグラフのサイズに応じてスケールするため、将来的な課題としてスケーラビリティを向上させるためのグラフの疎化（sparsification）の探索が挙げられるものの、相互作用する粒子システムのモデリングにおける重要な進展であると結論付けている。