[論文レビュー] JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library

J.E. Campagne, François Lanusse+8 more|arXiv (Cornell University)|Feb 10, 2023

Galaxies: Formation, Evolution, Phenomena被引用数 4

ひとこと要約

jax-cosmo は、JAX に基づいて構築された微分可能で GPU で加速された宇宙論ライブラリであり、宇宙論的尤度のエンドツーエンドの自動微分を可能にする。Fisher 行列の計算と、ハミルトニアン・モンテカルロ や NUTS といった勾配ベースの手法による事後分布推定を高速化し、DES-Y1 3x2pt 分析において、従来の MCMC と比較して GPU 時間あたり 5 倍の有効サンプルサイズを達成した。

ABSTRACT

We present jax-cosmo, a library for automatically differentiable cosmological theory calculations. It uses the JAX library, which has created a new coding ecosystem, especially in probabilistic programming. As well as batch acceleration, just-in-time compilation, and automatic optimization of code for different hardware modalities (CPU, GPU, TPU), JAX exposes an automatic differentiation (autodiff) mechanism. Thanks to autodiff, jax-cosmo gives access to the derivatives of cosmological likelihoods with respect to any of their parameters, and thus enables a range of powerful Bayesian inference algorithms, otherwise impractical in cosmology, such as Hamiltonian Monte Carlo and Variational Inference. In its initial release, jax-cosmo implements background evolution, linear and non-linear power spectra (using halofit or the Eisenstein and Hu transfer function), as well as angular power spectra with the Limber approximation for galaxy and weak lensing probes, all differentiable with respect to the cosmological parameters and their other inputs. We illustrate how autodiff can be a game-changer for common tasks involving Fisher matrix computations, or full posterior inference with gradient-based techniques. In particular, we show how Fisher matrices are now fast, exact, no longer require any fine tuning, and are themselves differentiable. Finally, using a Dark Energy Survey Year 1 3x2pt analysis as a benchmark, we demonstrate how jax-cosmo can be combined with Probabilistic Programming Languages to perform posterior inference with state-of-the-art algorithms including a No U-Turn Sampler, Automatic Differentiation Variational Inference,and Neural Transport HMC. We further demonstrate that Normalizing Flows using Neural Transport are a promising methodology for model validation in the early stages of analysis.

研究の動機と目的

• 宇宙論的パラメータおよびシステムティックなパラメータに関して、宇宙論的尤度の自動微分を可能にする微分可能な宇宙論ライブラリの開発。
• JAX の JIT コンパイル、GPU 加速、バッチ処理を活用して、高次元パラメータ空間における宇宙論的推論を高速化すること。
• 微調整を必要としない安定した Fisher 行列計算と、ハミルトニアン・モンテカルロ や変分ベイズ推定といった勾配ベースの事後分布推定を可能にすること。
• 微分可能な Fisher 行列を用いた調査最適化を可能にし、図の優位性（Figure of Merit）の勾配降下による最適化を可能にすること。
• NumPyro などの確率的プログラミング言語と統合し、宇宙論的分析におけるスケーラブルで高性能な事後分布サンプリングを可能にすること。

提案手法

• JAX の自動微分（autodiff）を活用し、宇宙論的観測量の宇宙論的パラメータおよびシステムティック要因に関する正確な勾配を計算する。
• 背景宇宙の進化、線形および非線形パワースペクトル（Halofit や Eisenstein & Hu を使用）、および Limber 近似を用いた角スペクトル（Cℓ）の微分可能版を実装する。
• JAX の JIT コンパイルと GPU/TPU 加速を活用して、尤度評価とサンプリングアルゴリズムの高速化を実現する。
• NumPyro などの確率的プログラミングフレームワークと統合し、NUTS や ADVI、Neural Transport HMC といった勾配ベースの推論手法を可能にする。
• 複数の GPU や CPU を通じた宇宙論的シミュレーションとパラメータスイープのバッチ処理に vmap を活用する。
• 3x2pt 分析用の微分可能な尤度をサポートし、理論からデータ比較までエンドツーエンドの微分可能性を実現する。

Figure 1: Comparison of the radial comoving distance between CCL and jax-cosmo . More plots are available in the companion notebook \faFileCodeO .

実験結果

リサーチクエスチョン

• RQ1 JAX における自動微分は、宇宙論において微調整を必要としない安定した Fisher 行列計算を可能にするか？
• RQ2 実際の宇宙論的分析において、ハミルトニアン・モンテカルロ や NUTS といった勾配ベースの推論手法は、従来の MCMC と比較してどの程度効率的か？
• RQ3 微分可能な尤度により、図の優位性（Figure of Merit）の勾配降下による直接的な調査設計最適化が可能か？
• RQ4 有効サンプルサイズ/時間単位における JAX ベースのサンプラーと最適化済み古典的 MCMC の性能向上はどの程度か？
• RQ5 正規化フローを用いたニューラルトランスポート HMC は、初期段階の宇宙論的分析におけるモデル妥当性評価を向上させるか？

主な発見

• DES-Y1 3x2pt 分析において、jax-cosmo の NUTS サンプラーは、well-optimized な Cobaya Metropolis-Hastings サンプラーと比較して、GPU 時間あたり約 5 倍の有効サンプルサイズを達成した。
• jax-cosmo で計算された Fisher 行列は正確で安定しており、もはや手動での微調整が不要となり、計算的・実装的オーバーヘッドが顕著に削減された。
• Fisher 行列自体が微分可能であるため、赤方偏移のビン分割や観測戦略といった調査設計パラメータの勾配降下による最適化が可能になった。
• NUTS や ADVI といった勾配ベースの推論手法は、古典的 MCMC と一致する事後分布サンプルを生成しつつ、高次元パラメータ空間へのスケーリングがより効率的である。
• 正規化フローを用いたニューラルトランスポート HMC は、初期段階の宇宙論的分析におけるモデル妥当性評価に有望な結果を示し、高速かつ高精度な事後分布近似を提供した。
• このライブラリにより、宇宙論的尤度のエンドツーエンドの微分可能性が実現され、最小限のコード変更で複雑な推論パイプラインを構築でき、GPU 硬貨上で高いパフォーマンスを発揮した。

Figure 2: Comparison of the non-linear matter power spectrum ( halofit function) between CCL and jax-cosmo . Also shown is the linear power spectrum.