CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators

要約

CosmoPower-JAX は、CosmoPower フレームワークの JAX ベースの実装であり、宇宙論的パワー スペクトルのニューラル エミュレーターを構築することで宇宙論的推論を加速します。
JAX の自動微分、バッチ評価、ジャストインタイム コンパイル機能を使用し、グラフィックス プロセッシング ユニット (GPU) で推論パイプラインを実行することで、高度な勾配ベースのサンプリングによりパラメーター推定を桁違いに高速化する方法を示します。
テクニック。
これらを使用すると、次世代の宇宙論調査の解析に必要な高次元パラメータ空間などを効率的に探索できます。
2 つのシミュレートされた Stage IV 構成での CosmoPower-JAX の精度と計算効率を紹介します。
まず、合計 37 のモデル パラメーターを含む宇宙せん断解析を実行する 1 つの調査について考えます。
CosmoPower-JAX とハミルトニアン モンテカルロ サンプラーで導出された等高線を、ネストされたサンプラーでエミュレーターなしで導出された等高線と比較して検証し、$\mathcal{O}(10^3)​​$ の高速化係数を取得します。
次に、合計 157 個のモデル パラメーターに対して、宇宙シアーと銀河クラスタリング (3x2pt) の共同解析を実行する 3 つのステージ IV 調査の組み合わせを検討します。
このような高次元パラメータ空間であっても、標準的な方法では推定 6 年かかるのに対し、CosmoPower-JAX は 3 日で収束した後部輪郭を提供します。
CosmoPower-JAX は完全に Python で書かれており、宇宙論コミュニティが次世代測量で設定された精度要件を満たすのを支援するために、私たちはこれを一般公開しています。

要約(オリジナル)

We present CosmoPower-JAX, a JAX-based implementation of the CosmoPower framework, which accelerates cosmological inference by building neural emulators of cosmological power spectra. We show how, using the automatic differentiation, batch evaluation and just-in-time compilation features of JAX, and running the inference pipeline on graphics processing units (GPUs), parameter estimation can be accelerated by orders of magnitude with advanced gradient-based sampling techniques. These can be used to efficiently explore high-dimensional parameter spaces, such as those needed for the analysis of next-generation cosmological surveys. We showcase the accuracy and computational efficiency of CosmoPower-JAX on two simulated Stage IV configurations. We first consider a single survey performing a cosmic shear analysis totalling 37 model parameters. We validate the contours derived with CosmoPower-JAX and a Hamiltonian Monte Carlo sampler against those derived with a nested sampler and without emulators, obtaining a speed-up factor of $\mathcal{O}(10^3)$. We then consider a combination of three Stage IV surveys, each performing a joint cosmic shear and galaxy clustering (3x2pt) analysis, for a total of 157 model parameters. Even with such a high-dimensional parameter space, CosmoPower-JAX provides converged posterior contours in 3 days, as opposed to the estimated 6 years required by standard methods. CosmoPower-JAX is fully written in Python, and we make it publicly available to help the cosmological community meet the accuracy requirements set by next-generation surveys.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google