Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation

技術要約：ガウス型ニューラル事後分布推定による宇宙論的初期条件の高速サンプリング

問題提起
宇宙論的な初期条件（宇宙論的初期条件、以下IC）を、大規模構造（LSS）の観測結果から再構成することは、宇宙論における根本的な課題である。初期宇宙は線形なガウス揺らぎによって良好に記述されるが、現在の宇宙網（コスミック・ウェブ）は高度に非線形かつ非ガウス的である。この進化を逆転させることは、不良設定問題（ill-posed problem）である。すなわち、初期から後期への写像は単射ではなく（異なるICが同じ後期の場を生み出す可能性がある）、ハロースケールでの構造形成はカオス的であり、測定ノイズを増幅させる。

BORG（Bayesian Origin Reconstruction from Galaxies）のような従来の確率論的手法は、ハミルトニアンモンテカルロ（HMC）を用いて事後分布をサンプリングする。しかし、これらの手法は微分可能なシミュレータ（多くの場合近似スキームに依存する）を必要とし、高次元のパラメータ空間（$\gtrsim 10^6$ パラメータ）で収束させるには、数ヶ月のCPU時間を要するなど計算コストが極めて高い。対照的に、この領域における多くの機械学習アプローチは、点推定器（例：U-Net）に依存しており、不確実性を定量化せずに単一の「最良適合」の再構成を提供する。あるいは、生成モデル（例：スコアベースの拡散モデル）に依存しているが、これらは推論時であっても計算コストが高いままである。

手法
著者らは、**ガウス型ニューラル事後分布推定（Gaussian Neural Posterior Estimation: NPE）**と呼ばれるシミュレーションに基づく推論（SBI）フレームワークを提案している。この核心となる革新は、推論された場 $\mathbf{z}$ における事後分布 $p(\mathbf{z}|\mathbf{x})$ を、フーリエ空間において対角な共分散行列を持つガウス分布としてモデル化することである。

手法は以下の通り進む：

1. 統計的定式化:

   • 事前分布: 事前分布 $p(\mathbf{z})$ は、高赤方偏移（$z=127$）における線形物質パワースペクトルによって定義されるガウス型ランダム場であると仮定される。
   • 尤度: 尤度 $p(\mathbf{x}|\mathbf{z})$ は、観測依存の平均 $\hat{\mathbf{z}}_\theta(\mathbf{x})$ と精度行列 $\mathbf{Q}^L_\theta$ を持つガウス分布としてモデル化される。決定的なのは、$\mathbf{Q}^L_\theta$ が物理の並進不変性を反映して、フーリエ空間において対角であり（かつ近似的に回転不変である）という点である。
   • 事後分布: ガウス事前分布とガウス尤度の組み合わせにより、ガウス型事後分布が得られる。事後分布の精度行列は、事前分布の精度行列と尤度の精度行列の和である（$\mathbf{Q}_\theta = \mathbf{Q}_P + \mathbf{Q}^L_\theta$）。

2. ニューラルアーキテクチャ:

   • MAP推定器: 事後分布の平均 $\hat{\boldsymbol{\mu}}_\theta(\mathbf{x})$ は、3D U-Net アーキテクチャを用いてモデル化される。性能向上のため、著者らは再構成を2つのスケールに分離するカスタム活性化層を導入している。学習可能なスケーリング因子 $\alpha_\theta(\mathbf{k})$ は大規模な線形進化を扱い、U-Netは非線形な小スケールの残差の再構成に焦点を当てる。
   • 精度行列: 尤度の精度行列の対角要素 $D^L_\theta(k)$ が学習される。著者らは2つのパラメータ化をテストした：(1) 各フーリエモードに対して自由パラメータとして扱う方法、(2) 波数 $k$ を入力とする高密度ニューラルネットワークを用いて回転対称性を強制する方法。結果として、「自由」な行列アプローチの方がわずかに安定性が高いことが示された。

3. 学習とデータ:

   • モデルは、Quijote $N$体シミュレーション・スイートからの、初期（$z=127$）と最終（$z=0$）のダークマター過密度のペア2,000組を用いて学習される。
   • フォワードモデルには、完全で非微分的な $N$体シミュレーション（Gadget-III）を使用しており、HMC法で必要とされる微分可能な近似スキームを回避している。
   • 学習は、単一のNvidia A100 GPU上で30エポックにわたって負の対数事後分布（式7）を最小化し、約 1.5時間 で完了する。

主要な貢献

• 効率的なサンプリング: フーリエ空間の対角な精度行列を持つガウス型事後分布を仮定することで、既存のHMCや拡散ベースの手法よりも数桁速い、数秒での数千の事後分布サンプルの生成を可能にした（単一GPUで 1,000サンプルあたり約3秒）。
• 非微分的なシミュレータ: 本手法はフォワードモデルが微分可能であることを要求しないため、近似スキームなしで最先端の完全な $N$体シミュレーションを使用できる。
• 高速サンプリングのための解析的フィット: 学習された精度行列の対角値 $D_\theta(k)$ が、$k \approx 0.2 \, h/\text{Mpc}$ までの範囲で単純な解析関数（ガウス分布＋定数）で正確にフィットできることを著者らは示している。これにより、学習済みの点推定器（MAP）を、フル精度行列の再学習なしに高速な事後分布サンプラーへと変換することが可能になる。
• 不確実性の定量化: 標準的な点推定器とは異なり、本手法は完全な事後分布を提供し、場レベルの再構成における厳密な不確実性の定量化を可能にする。

結果
著者らは、訓練セットに含まれていない標準的なシミュレーションを用いて手法を検証した：

• 視覚的および統計的一致: 生成されたサンプルは真の値を視覚的に模している。要約統計量（パワースペクトル、伝達関数、相互相関）は、真の初期条件と高い一致を示した。パワースペクトルは、ナイキストスケール（$k_{Nyq} \approx 0.4 \, h/\text{Mpc}$）まで、真の値に対して 1–2% 以内の誤差に収まっている。
• 相互相関: 生成されたサンプルと真の値との間の相互相関は、$k \approx 0.3 \, h/\text{Mpc}$ まで高い値（$\gtrsim 70\%$）を維持している。
• ベイズ的一貫性: スケール依存の一貫性テストにより、真のモードと事後分布の平均との正規化された差が、様々な $k$ ビンにおいて標準正規分布（$N(0,1)$）に従うことが確認された。これは、推定された不確実性が適切に較正されていることを検証している。
• 堅牢性: 観測された最終的な場にガウスノイズが加えられた場合でも、手法は性能を維持している。

意義と主張
本論文は、このアプローチが宇宙論における逐次的なSBIのためのシミュレーション効率の高い経路を提供すると主張している。事後分布の複雑さを、対角なフーリエ空間の精度行列を持つガウス形式に落とし込むことで、著者らは、物理的な正確さと計算速度のバランスを、これまで到達不可能であったレベルで実現した。

著者らは、ガウス型フーリエ対角近似は選択されたスケール範囲（$k \lesssim 0.4 \, h/\text{Mpc}$）においては十分であるが、より小さなスケールや、観測効果によって並進不変性が破れる場合には、より洗練された精度行列（例：畳み込みまたはウェーブレットベース）が必要になる可能性があると控えめに述べている。また、現在の研究は固定された宇宙論パラメータと理想的なダークマター観測を前提としているが、本フレームワークは将来的に変化する宇宙論や現実的な銀河サーベイへと拡張できるように設計されていることを強調している。主な意義は、数千の制約付き実現（constrained realizations）を数秒で生成できる能力にあり、これは、これまで計算上不可能であった大規模構造の堅牢な統計分析への道を開くものである。