Using Neural Emulators and Hamiltonian Monte Carlo to constrain the Epoch of Reionization's History with the Ly$α$ Forest Power Spectrum

技術要約：ニューラルエミュレータとハミルトニアンモンテカルロを用いた宇宙再電離史の制約

問題提起
宇宙再電離（EoR）は宇宙史における極めて重要な転換期であるが、その詳細なタイムラインや不均質な構造については、依然として制約が不十分である。宇宙マイクロ波背景放射（CMB）の測定値は、積分された光学的厚さの制約は与えるものの、再電河の具体的な期間やトポロジーには感度を持たない。高赤方偏移クエーサーのスペクトル（例：ダンピング翼、ダークピクセル分率）からの現在の制約は、精度が限定的であり、しばかにモデル依存的な解釈に依存している。赤方偏移 $z \sim 5$ におけるリマン-$\alpha$（Ly$\alpha$）フォレストは、星間物質（IGM）の熱的および電離状態を知るための有望なプローブであるが、詳細な再電離史を抽出するには、Ly$\alpha$ 不透明度の全分布を活用する必要がある。

定量的な推論における大きな障壁は、計算コストである。EoRに関連する小スケールの構造を捉えるのに十分な解像度を持つ流体解析シミュレーションを生成するには、1つのモデルにつき数千時間のGPU時間が必要となる。その結果、これまでの研究は定性的な比較や制限されたパラメータ空間に限定されてきた。さらに、ロバストなパラメータ推定には正確な共分散行列が必要であるが、利用可能な高赤方偏移クエーサーのサンプルは疎であるため、それらを経験的に導出することは困難である。従来のアプローチでは、固定された共分散構造を仮定するか、モデル依存の共分散を計算するために、非常に高価なアンサンブル・シミュレーションに頼る必要があった。

手法
本研究は、微分可能なニューラルエミュタとハミルトニアンモンテカルロ（HMC）を統合した、JAXベースの新しい推論フレームワークを提示する。この手法は、微分可能なLy$\alpha$フォレスト観測量と、フォワードモデリングを通じて進行する。プロセスは以下のステップに従う：

1. シミュレーションデータセットの構築： 著者らは、Nyxコードを用いて、501個の低解像度（$256^3$ ボクセル、$L_{box} = 20 \, h^{-1}$ Mpc）の流体解析シミュレーションからなるデータセットを生成した。これらのシミュレーションは、ユーザー定義のパラメータを電離場にマッピングする半数値的コードであるAMBERによって生成された再電離史によって駆動されている。パラメータ空間には、再電離の中点（$z_{mid}$）、期間（$\Delta z$）、非対称性（$A_z$）、最大加熱温度（$\Delta T_{re}$）、および平均フラックス（$\langle F \rangle$）が含まれる。
2. フォワードモデリング： これらのシミュレーションから、分光器の分解能（$R=30,000$）やノイズ（$SNR=30$）などの観測効果を適用した1次元Ly$\alpha$フラックススキュー（skewers）を生成した。そして、平均1次元パワースペクトル（$\langle P_{Ly\alpha} \rangle$）およびモデル依存の共分散行列（$\Sigma_{P_{Ly\alpha}}$）を各モデルについて算出した。
3. ニューラルエミュレーション： 2つの独立したニューラルエミュタを訓練した：
   • パワースペクトル・エミュレータ： 入力パラメータから $\log_{10}\langle P_{Ly\alpha} \rangle$ を予測する、全結合フィードフォワードネットワーク。
   • 共分散行列エミュレータ： 出力が対称かつ正定値であることを保証するため、著者らはコレスキー分解（$\Sigma = LL^T$）を採用した。ニューラルネットワークは、平坦化された下三角コレスキー因子（$L_{flat}$）を予測し、正定値性を強制するために対角要素に指数関数的な活性化関数を適用している。
4. 推論パイプライン： エミュレータは、Numpyroパッケージを介してHMCと組み合わされた。JAXベースのエミュタの微分可能性により、HMCに必要な勾配の計算が可能となり、標準的なメトロポリス・ヘイスティングス法と比較してパラメータ空間の探索が大幅に加速された。

主な貢献

• 微分可能な共分散エミュレーション： 本研究は、ニューラルネットワークを用いてLy$\alpha$フォレストの共分散行列をエミュレートする初めての試みである。コレスキー因子を予測することで、パラメータ依存の相関を捉えつつ、物理的な妥当性（対称性と正定値性）を確保している。
• エンドツーエンドの微分可能フレームワーク： フォワードモデリング、ニューラルエミュレーション、およびHMCをJAXエコシステム内で統合することで、宇宙論的パラメータ推論のための完全な微分可能パイプラインを実現した。
• 定量的検証： 本フレームワークは100個の模擬観測に対して検証され、真のパラメータを回収できる能力を示し、得られた事後分布の統計的妥当性を評価した。

結果

• エミュレータの精度： パワースペクトル・エミュレータは、テストセットにおいて中央値0.15%の平均絶対パーセント誤差（MAPE）を達成し、95%の誤差が2%未満であった。共分散行列エミュレータは、対角要素の積に対する中央値誤差0.42%を達成し、すべてのスケールにおいて誤差は一般に5%未満であった。両エミュレータとも、フォワードモデリングで想定された期待観測ノイズレベルを十分に下回る精度を示した。
• パラメータ回収： 模擬観測に適用した際、HMCフレームワークは真の入力パラメータを正常に回収した。事後分布は、真の値と視覚的に一致していた。
• 統計的妥当性： 100個の模擬観測に対して行われたカバレッジテストでは、事後分布がわずかに過剰に確信的（推定されたカバレッジ確率が理想的なラインをわずかに下回っている）であることが明らかになった。著者らは、この原因を主に多変量ガウス尤度関数の仮定にあると考えており、Ly$\alpha$フォレストが本質的に非ガウス的なランダム場であることを指摘している。
• データ効率： 著者らは、わずか100個のシミュレーションを用いた縮小された訓練セットでのテストを行った。エミュレーション誤差は増加した（パワースペクトルのMAPEは約1.2%に上昇、共分散誤差も増加）が、得られた事後分布および統計的カバレッジは、フルデータセットから得られたものと同等であった。これは、本手法が限られた訓練データに対しても堅牢であることを示唆しており、将来の高解像度シミュレーションにおいて極めて重要な要素である。

意義と限界
本論文は、このフレームワークがLy$\alpha$フォレストのパワースペクトルを用いて、宇宙再電離の歴史を定量的に制約するための実行可能な道筋を提供すると主張している。主な意義は、ニューラルエミュタが高価なフォワードモデルを効率的に代替できること、およびモデル依存の共分散行列をエミュレートできることを示した点にあり、これにより、より高速で厳密なベイズ推論が可能になる。

しかし、著者らは、使用したシミュレーションの解像度が低く、小スケールのパワーが抑制されているため、現在の結果を実際の観測データに直接適用することはまだできないと明言している。また、再電離モデルは簡略化された熱注入処方箋に依存しており、紫外背景（UVB）の空間的ゆらぎを欠いている。著者らは、将来の研究において、より高解像度のシミュレーション、より物理的に動機付けられた熱的処理、およびゆらぎのあるUVBモデルを組み込む必要があると結論付けている。また、事後分布に見られたわずかな過剰確信は、ガウス尤度仮定に依存しないシミュレーションベースの推論（SBI）手法へと移行することで対処できる可能性があることも示唆している。