Emulating Cosmic Structure Formation with a Lagrangian Neural Cellular Automaton

技術要約：ラグランジュ型ニューラル・セル・オートマトンによる宇宙構造形成のエミュレーション

問題提起
次世代の銀河レッドシフト・サーベイ（DESK、Euclid、LSSTなど）は、標準的なパワースペクトルのような要約統計量を超えて、宇宙論的マップから完全な物理的情報内容を抽出するための、フィールドレベルの推論手法を求めている。この逆問題には、線形初期条件と非線形な後期構造とを結びつけるフォワードモデルが必要である。そのようなモデルは、以下の3つの競合する制約を満たさなければならない：

1. 物理的正確性： 従来のラグランジュ摂動論（LPT）が破綻する領域である、シェルクロッシングやハロー形成を含む非線形ダイナミクスを正しくモデル化すること。
2. 計算効率： 反復的なループが必要なベイズ推論パイプラインにおいて、標準的な$N$体シミュレーションのように計算コストが高すぎて使用できないことがないよう、十分に高速であること。
3. 微分可能性： 初期条件および宇宙論パラメータの再構成のための勾配ベースの最適化を可能にするため、完全に微分可能であること。

既存の解決策はトレードオフに直面している：$N$体シミュレーションは正確だが遅く、LPTや有効場理論（EFT）モデルは高速だが非線形領域で破綻し、標準的なオイラー型畳み込みニューラルネットワーク（CNN）は、広範囲にわたる質量の移流（advection）に苦しみ、ガリレイ不変性を欠き、ライトコーン構築に必要な時間的に連続した軌道を生成できないことが多い。

手法：ラグランジュ型ニューラル・セル・オートマトン（LNCA）
著者らは、宇宙構造形成を、共動格子上の局所的かつ反復的な動力学的プロセスとしてエミュレートするハイブリッド深層学習フレームワークである、ラグランジュ型ニューラル・セル・オートマトン（LNCA）を導入する。

• ラグランジュ・フレームと固定トポロジー： 固定された空間グリッド上で動作するオイラー型CNNとは異なり、LNCAはラグランジュ・フレームで動作する。これは、初期の近傍点との接続関係が固定された、ラグランジュ座標（$q$）の規則的なデカルト格子を定義する。ノード（流体要素を表す）は更新則に従って変位するが、近傍グラフのトポロジーは不変である。これにより、グラフベースのネットワークが持つ連続的な粒子追跡を維持しつつ、高度に最適化された3D畳み込みの使用が可能になる。
• 残差学習： 大規模な一貫性と効率を確保するため、ネットワークは全重力ポテンシャルを学習するのではなく、ゼルドビッチ近似（ZA）に対する残差変位補正のみを学習する。全変位は $\Psi(t) = D(t)\Psi_{IC} + \Psi_{res}(t)$ と分解され、ネットワークは $\Psi_{res}$ を予測する。
• 等変アーキテクチャ： 本モデルは、構成上、厳密な$E(3)$（並進および回転）等変性を強制する。生の状態ベクトルではなく、変位場の導関数（ラプラシアン、ヘッセ行列、ベクトル積）から導出された幾何学的不変量の基底に基づいて動作する。
  • スカラープール ($I_i$): 勾配と内積から構成される、$E(3)$不変スカラーの集合。
  • ベクトルメニュー ($E_i$): 等変ベクトル（変位、速度、潜在ベクトル）の基底。
  • 更新則は、スカラープールを係数へと写像し、それらがベクトルメニューを線形に変調する。
• コルモゴロフ・アーノルド・ネットワーク（KAN）： コアとなる更新関数 $f_\theta$ は、標準的なReLUではなく学習可能な多項式である活性化関数を用いる、制約付きKANとして実装されている。これにより、解析的な連続性、解釈可能性、および学習パラメータの大幅な削減が保証される。
• 時間的連続性： モデルは変位と潜在場の時間微分を予測し、連続的な時間ステップを用いて微分可能なスキームで積分される。これにより、単発のマッピングではなく、完全な粒子軌道が生成され、中間赤方偏移やライトコーンの構築が可能になる。
• 学習： モデルは、Quijoteシミュレーションスイートを用いて、各粒子の変位、運動量、実空間密度、およびハロー集団モーメントのエラーを罰する多目的損失関数によって学習される。また、すべてのノードにブロードキャストされる宇宙論パラメータ（$\Omega_m, \Omega_b, h, n_s, \sigma_8$）によって条件付けられる。

主な結果
学習されたLNCAは、広範な赤方偏移（$z \in [0, 3]$）および宇宙論（保持されたラチン・ハイパーキューブの点を含む）に対して、$N$体シミュレーションと比較検証された：

• パワースペクトルと相互相関： モデルは、非線形領域（具体的には $k \lesssim 0.5 \, h \text{Mpc}^{-1}$）において、物質パワースペクトル ($P(k)$) および相互相関 ($r(k)$) に関してパーセントレベルの精度を達成した。$k \sim 1 \, h \text{Mpc}^{-1}$ においても、偏差は5%未満に留まっている。
• 変位の正確性： 中央値の変位誤差は約 $0.2 \, \text{Mpc}/h$ であり、95%の粒子が $2 \, \text{Mpc}/h$ 未満の誤差を持つ。モデルは、シェルクロッシング・コースティクスやバブル状のボイドといった複雑な非線形特徴を正常に再構成している。
• ハロー形成： モデルは、大規模なハローの存在量と空間分布を正確に回復する。低質量ハローについては、（シェルクロッシング後の拡散成分による）わずかな不足が見られるものの、高質量端についてはポアソン散乱の範囲内で$N$体真値と一致している。
• 汎用性： モデルは、学習分布から離れた宇宙論に対しても頑健に汎化し、再学習なしに高い忠実度を維持する。

意義と主張
論文は、LNCAが精度、効率、および物理的制約を同時に解決することで、宇宙論的エミュレーションにおける重要な進歩を実現したと主張している：

1. パラメータ効率： LNCAは、約 $2 \times 10^7$ 個の重みを必要とする最先端のU-Netエミュレータと同等の精度を、わずか $\sim 4 \times 10^3$ 個の学習パラメータで達成している。この4桁の減少は、単一のコンパクトで物理的に制約された更新則を再帰的に適用したことによるものである。
2. 局所性とスケーラビリティ： 全域的なフーリエ変換（$N$体コードにおけるポアソンソルバーに必要）に依存せず、局所的な更新に依存することで、LNCAはグローバルな同期障壁を回避する。これにより、通信量は体積ではなく表面積に比例してスケールするため、大規模なボリュームに対して優れたスケーラビリティを提供する。
3. 微分可能性とライトコーン： スナップショット型のエミュレータとは異なり、LNCAは連続的で微分可能な軌道を生成する。これは、追加の計算パスなしにライトコーン出力（過去のライトコーンに沿った観測のシミュレーション）を構築することをネイティブにサポートしており、サーベイ解析における重要な機能である。
4. 解釈可能性： KANと等変基底の使用により、内部ダイナミクスを物理的不変量上の明示的な代数式として検査することが可能となり、「ブラックボックス」的な深層学習には欠けているレベルの解釈可能性を提供している。

著者らは、LNCAがFastPMやpmwdのような特化したソルバーと比較して、生の実行時間（ウォールクロックタイム）において最速のフォワードモデルではないかもしれないが、局所性、再帰性、微分可能性、および物理的制約の独自の組み合わせにより、次世代のフィールドレベルの宇宙論的推論のための信頼できるスケーラブルなツールとなることを結論付けている。