Diffhalos: A Generative Model of Cosmological Lightcones of Dark Matter Halos

技術要約：「Diffhalos: ダークマターハローの宇宙論的ライトコーンの生成モデル」

問題提起
銀河サーベイの現代的な宇宙論的解析には、理論モデルをテストし観測データを解釈するための合成ライトコーン（宇宙時間を通じた銀河のシミュレーションカタログ）が必要である。従来の設計手法は、大規模な重力N体シミュレーションを実行するか、出力スナップショット間を補間してライトコーンを生成することに依存しているが、これらのアプローチは計算コストが高く、微分可能性を欠くことが多い。さらに、物理的に現実的な合成集団を生成するには、合体や降着を通じてハローがどのように成長するかを記述する質量集積履歴（MAH）を組み込む必要がある。拡張プレス・シェレ（Press-Schechter）形式や標準的なフィッティング関数などの既存のMAHモデルは、多様な集積経路を捉えるための柔軟性や、勾配ベースの推論に必要な計算効率に欠ける場合が多い。したがって、ダークマターのハロー、サブハロー、およびそのMAHを含むライトコーンを統計的な正確さで生成でき、かつパラメータの効率的な最適化のための自動微分をサポートする生成フレームワークが求められている。

手法
著者らは、自動微分とGPU加速を活用するためにJAXで実装された生成モデルであるDiffhalosを紹介する。このフレームワークは、宇宙論的ライトコーンの生成を、以下の3つの明確かつ微分可能なステージに分解する：

1. ホストハロー生成：
   モデルは、ハロー質量関数（HMF）$n_{\text{halo}}(m_{\text{halo}}, z)$からサンプリングすることで、ライトコーン体内のホストハロー集団を生成する。著者らは2つの微分可能な実装を検討している：

   • Haloxベース： 解析的なハローモデル（Tinker et al. 2008のフィッティング公式を使用）を用いたJAX実装であり、線形物質パワースペクトルを用いてHMFを計算する。これにより、宇宙論パラメータに対する勾配を直接計算することが可能になる。
   • エミュレータベース： 解析的アプローチにおける数値積分に伴うメモリオーバーヘッドを削減するため、累積HMF $N_{\text{halo}}(> m_{\text{halo}}, z)$をエミュレートする多層パーセプトロン（MLP）を訓練する。微分HMFはjax.gradを介して導出される。
     サンプリングは、赤方偏移と質量のグリッド上での逆変換サンプリングを用いて行われる。また、集団を近似するために重み付きグリッドを用いる**準モンテカルロ（QMC）**ジェネレータも実装しており、統計的一貫性を維持しながらメモリ使用量を大幅に削減している。

2. サブハロー生成：
   ホストハローが確立された後、モデルは微分可能な**条件付きサブハロー質量関数（CSHMF）**を用いて、それらにサブハローを配置する。CSHMFは、$\mu = M_{\text{sub}}/M_{\text{host}}$における累積分布 $N_{\text{sub}}(> \mu | M_{\text{host}})$としてパラメータ化される。著者らは、Discoveryシミュレーション（HACCコード）およびJiang & van den Bosch (2016) のフィッティング関数に対して較正された「sig-slope」カーネル（べき乗則にシグモイド変化する傾斜を組み合わせたもの）を利用している。ホスト生成と同様に、モンテカルロ法とQMC法の両方がサブハローのサンプリングに使用可能である。

3. 質量集積履歴（MAH）の割り当て：
   最終ステージでは、すべてのホストおよびサブハローにMAHを割り当てる。MAHは、パラメータ集合 $\theta_{\text{MAH}}$ でパラメータ化されたDiffmah関数形式を用いて近似される。これらのパラメータをシミュレーション集団にマッピングするために、著者らは、Small MultiDark Planck (SMDPL) N体シミュレーションのマーガースリー（合体樹）を用いて訓練された正規化フローアーキテクチャであるDiffmahNetを導入した。DiffmahNetは、観測時点におけるハローの質量と赤方偏移に条件付けられた確率分布 $P(\theta_{\text{MAH}} | z, M_{\text{halo}}(z))$ を学習する。これにより、多様で物理的に妥当な集積履歴の生成が可能となる。

主な貢献

• Diffhalosフレームワーク： ダークマターハロー、サブハロー、およびそのMAHの宇宙論的ライトコーンを生成するための、完全に微分可能なJAXベースのライブラリ。
• 微分可能なHMFおよびCSHMF： 微分可能なハロー質量関数（Haloxおよびカスタムエミュレータ経由）および条件付きサブハロー質量関数の実装により、宇宙論的および天体物理学的パラメータに関する勾配ベースの最適化を可能にする。
• DiffmahNet： 以前のパラメータ的MAHモデル（Hearin et al. 2021）を一般化し、全集団のMAHを生成する正規化フロー。これは、平均的な成長トレンドと集団の分散の両方を捉える。
• 効率化ツール： ホストおよびサブハロー集団の両方に準モンテカルロ（QMC）ジェネレータを導入し、標準的なモンテカルロ実現と比較して、メモリ効率において数桁の向上を実現している。
• 公開コード： 銀河・ハロー結合モデルの開発と較正を支援するため、Diffhalosのコードベースを公開する。

結果
著者らは、N体シミュレーションのデータに対してDiffhalosフレームワークの検証を行っている：

• HMFの精度： Halox実装は、異なる宇宙論において、広範な質量と赤方偏移（$z=0$から$5$）にわたってColossusライブラリとパーセントレベルの一致を示した。MLPベースのエミュレータは$\le 10\%$の精度を達成しており、これはTinker et al. (2008) のフィッティング関数自体の固有の精度に匹敵する。
• サブハロー統計： サブハロージェネレータは、ホスト質量 $10^{10} < M_{\text{host}}/M_\odot < 10^{14}$ および $0 < z < 4$ の範囲において、Discoveryシミュレーションからのホストおよびサブハローの総質量関数を妥当な忠実度で再現している。
• MAHの忠実度： DiffmahNetは、異なるエポック（$t_{\text{obs}} = 6.42$ および $13.8$ Gyr）で特定されたハローにおける、SMDPLシミュレーションに見られるMAHの平均的な質量成長と分散を正確に再現している。
• 勾配計算： フレームワークは、宇宙論パラメータ（$\Omega_m, \sigma_8$）に対するHMFの勾配を正常に計算しており、微分可能なアプローチの有用性を実証している。

意義と主張
本論文は、Diffhalosを主に銀河・ハロー結合モデルの開発と較正を支援するために設計された「概念実証（プルーフ・オブ・コンセプト）」ツールとして位置づけている。著者らは、このフレームワークの微分可能な性質が、ハミルトニアンモンテカルロ法や勾配降下法のような、セミ解析的モデル（SAM）や経験的モデルを観測データに対して最適化するために不可欠な、効率的な勾配ベースのアプリケーションを可能にすることを強調している。

著者らは、現在の実装は概念実証であることを明示的に述べており、サブハロー質量関数およびMAHの宇宙論依存性が、まだ（現在はSMDPL宇宙論のみに訓練されているため）一連の宇宙論に対して十分に較正されていないことを謙虚に指摘している。今後の課題として、最先端のエミュレータの統合、サブハローおよびMAHの宇宙論依存性の較正、および空間的なクラスタリング（空の座標）やバリオン効果を含む拡張が必要であると特定している。主要な価値提案は、合成的なハロー集団を生成するための、N体シミュレーションに代わる計算効率の高い微分可能な代替手段を提供することであり、これにより銀河形成モデルの最適化を加速させることにある。