技術概要：トランスフォーマーによるラグランジュ粒子動力学の統合シミュレーション

問題提起

物理現象の多様性をすべてモデル化できる統合シミュレーションフレームワークの構築は、シミュレーション科学における長年の課題として残されている。従来のアプローチは、各現象ごとに異なる離散化（例えば、流体に対するオイラー法格子、固体に対するラグランジュ法メッシュ）および支配方程式（例えば、ナビエ - ストークス方程式、構成則、原子間ポテンシャル）に依存している。現代のエンジンはいくつかのソルバーバックエンドを組み合わせるが、依然として現象ごと、ソルバーごとのパラダイムで動作するため、現象間の結合が困難であり、単一パイプラインの適用範囲を制限している。

既存の機械学習ベースのシミュレーターは、明示的な偏微分方程式（PDE）を迂回してデータから直接動力学を学習することで、補完的なアプローチを提供する。しかし、現在のニューラルシミュレーターは通常、特定の現象や方程式クラス向けに設計・訓練されている。さらに、低次元化手法は学習された基底を固定された幾何学形状や形状ファミリーに結びつけることが多く、物理や幾何学が変化する際にモデルの再設計を必要とする。ソルバー固有の再設計なしに多様なラグランジュ粒子システムに一般化できる単一のアーキテクチャが必要とされている。

手法

著者らは、布、弾性固体、ニュートン流体および非ニュートン流体、粒状物質、分子動力学のシミュレーションを、単一の固定アーキテクチャ内で統合する「粒子ベースの予測 - 修正トランスフォーマー」を提案する。このモデルは、粒子シミュレーション（例えば、位置ベースダイナミクス）で一般的な予測 - 修正スキームに従い、既知の外力の積分と粒子間相互作用の学習を分離する。

1. 予測 - 修正フレームワーク

シミュレーションステップは2つのフェーズに分割される：

• 明示的予測器： 既知の外力（$\mathbf{F}_t$）を時間ステップ $\Delta t$ にわたって積分し、中間状態（$\tilde{\mathbf{X}}_t, \tilde{\mathbf{V}}_t$）を生成する。このステップは外部加速度を考慮するが、粒子間相互作用は無視する。
• 学習型修正器： ニューラルネットワーク $F_\theta$ が、中間状態、粒子属性（$\mathbf{C}$）、トポロジー（$\mathbf{T}$）、および境界条件（$\mathbf{X}_b, \mathbf{C}_b$）に基づいて、残差更新（$\Delta \mathbf{X}_t, \Delta \mathbf{V}_t$）を予測する。最終状態は $\mathbf{X}_{t+\Delta t} = \tilde{\mathbf{X}}_t + \Delta \mathbf{X}_t$ および $\mathbf{V}_{t+\Delta t} = \tilde{\mathbf{V}}_t + \Delta \mathbf{V}_t$ として更新される。

2. アーキテクチャ構成要素

修正器は、局所的および全球的な相互作用の両方を効率的に処理するように設計された3段階で構成される：

• 粒子トークナイザー： グローバルアテンション前に局所相互作用をエンコードする。学習可能なカーネル（3Dグリッドとしてパラメータ化され、三重線形補間を使用）を用いて3つの近傍からの寄与を集約する：

  • 空間近傍 ($N^s_i$)： 相対変位と速度に基づく粒子 - 粒子相互作用。
  • 境界近傍 ($N^b_i$)： 粒子 - 境界相互作用。
  • トポロジカル近傍 ($N^t_i$)： 布/固体の場合、レスト形状の隣接関係に基づく相互作用。
    出力は、これらの集約された特徴を粒子自身の状態と連結した、粒子ごとのトークン $\mathbf{h}_i$ である。

• スーパートークンエンコーダー： $N$ 個の粒子に対するグローバル自己アテンションの計算コストに対処する。マルチグリッド手法に触発され、このモジュールはグローバル自己アテンションとトークンマージを交互に行う。

  • 自己アテンション： 3Dロータリー位置エンベディング（RoPE）を使用して、現在の解像度におけるトークンを文脈化する。
  • トークンマージ： コサイン類似度が最も高いトークンを貪欲にマッチング・マージすることで、各レイヤーでトークン数を半分にする。これにより、局所特徴、速度、材料情報を集約するコンパクトな「スーパートークン」のセットが作成される。
  • このプロセスは $L_e$ レイヤーにわたって繰り返され、深いレイヤーにおけるシーケンス長と計算コストを漸進的に削減する。

• スーパートークンデコーダー： $L_d$ レイヤーを通じて、粗いコンテキストを完全な粒子解像度へと引き上げる。

  • 粒子トークンはクロスアテンションを介してスーパートークンをクエリし、スーパートークンにエンコードされたグローバル情報にアクセスできるようにする。
  • 粒子トークンはその後、自己アテンションおよびフィードフォワードネットワークを通じて洗練される。
  • 最終的なMLPは、洗練されたトークンを残差位置および速度補正にマッピングする。

3. 訓練

モデルは、$W$ 状態のウィンドウにわたる自己回帰的ロールアウト損失を用いて訓練される。損失関数は、位置と速度の平均二乗誤差（MSE）と、オプションの物理情報に基づく正則化子（例えば、非圧縮性流体のための発散フリー制約）を組み合わせる。アーキテクチャはすべての動力学カテゴリにわたって固定されたままとなり、変化する唯一の要素は訓練データの分布である。

主要な貢献

1. 統合アーキテクチャ： ネットワーク構造を変更することなく、6つの異なるラグランジュ動力学カテゴリ（流体、布、粒状物質、弾性固体、非ニュートン流体、および分子動力学）をシミュレートできる単一のトランスフォーマーベースモデル。
2. 新規トークン機構： 局所近傍相互作用のための粒子トークナイザーと、プログレッシブなトークンマージを通じてグローバル通信を可能にするスーパートークンエンコーダー・デコーダーの組み合わせ。これにより、表現力と計算効率のバランスが取れている。
3. 一般化と応用： 未見の材料、境界構成、初期条件、および外力に対する一般化を実証。このフレームワークは、インタラクティブ制御、微分可能なロールアウトによる逆設計、および実世界の操作データからの学習を含むダウンストリームタスクをサポートする。

実験結果

モデルは、確立されたソルバー（DFSPH、GIPC、MPM、VBD、OpenMM）によって生成されたデータを用いて、6つのカテゴリで評価された。

• 精度： 提案モデルは、流体、布、砂のタスクにおいて、最先端のベースライン（GNS、DeepLagrangian、Neural Operators）と比較して、位置および速度のロールアウト平均二乗誤差（MSE）が低い。例えば、流体タスクでは、位置誤差は（スケーリング後）$5.17 \times 10^{-1}$ であり、長期的なロールアウトで浮遊粒子や相互貫入などのアーティファクトを示したベースラインを大幅に上回った。
• 安定性： ベースラインがしばしば発散または崩壊する長期的な視野（例えば、流体で800フレーム、布で250フレーム）にわたって、モデルは安定したシミュレーションを維持する。
• 一般化：
  • 未見の条件： 未見の境界幾何学（例えば、虎の形をした砂の山）および初期条件を成功裡にシミュレートする。
  • 粒子密度： 訓練セットとは異なる粒子サンプリング密度（例えば、約7k粒子で訓練し、8k〜9kでテスト）に一般化する。
  • 長期的ロールアウト： 訓練視野（200フレーム）を超えて250フレームでテストされた場合でも、安定性を維持する。
• アブレーション研究： 局所トークナイザーまたはグローバルスーパートークン通信のいずれかを除去すると性能が低下し、局所的な詳細と全球的な結合の両方の必要性が確認された。発散フリー正則化子を追加すると、流体の物理的一貫性がさらに向上する。

意義と主張

本論文は、この研究が、単一の微分可能なアーキテクチャが複数の特殊化されたソルバー（FEM、SPH、MPM）の必要性を置き換える統合シミュレーションフレームワークへの一歩であると主張している。各現象を、共有の入力 - 出力フォーマット上の異なる訓練分布として扱うことで、モデルは現象ごとのソルバー工学の必要性を低減する。

著者らは、このアーキテクチャによって可能となる3つの具体的なダウンストリーム応用を強調している：

1. インタラクティブ制御： 任意の外力を積分する予測器の能力を活用し、ユーザー指定の外力下での変形物体のリアルタイム操作。
2. 逆設計： ニューラルロールアウトを通じた逆伝播による物理パラメータ（例えば、摩擦係数）の勾配ベース最適化。
3. 実世界データからの学習： 古典的シミュレーターに依存することなく、実世界の観測から抽出された粒子軌跡（3Dガウススプラッティングなどの技術を使用）を直接訓練データとして利用する能力。

著者らは、主に第一エンコーダーレイヤーにおけるフル解像度自己アテンションに起因するメモリ制約により、現在の実装は単一の80GB GPU上で約50,000個の粒子に限定されているという限界を認めている。スパースアテンションのバリアントが、この容量を拡張する可能性があるとしている。さらに、アーキテクチャは統合されているが、現在の実装では動力学カテゴリごとに個別のパラメータセットを訓練している。将来の目標は、すべてのカテゴリで共有重みを持つ単一モデルを実現することである。