[论文解读] 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 加速的宇宙学库，支持宇宙学似然函数的端到端自动微分。它通过基于梯度的方法（如哈密顿蒙特卡洛和 NUTS）加速费舍尔矩阵计算和后验推断，在 DES-Y1 3x2pt 分析中，每 GPU 小时的有效样本数比传统 MCMC 提高了 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 加速和批量处理技术，加速高维参数空间中的宇宙学推断。
• 实现稳健的、无需微调的费舍尔矩阵计算和基于梯度的后验推断（如哈密顿蒙特卡洛和变分推断）。
• 通过可微分的费舍尔矩阵支持探测器优化，实现基于梯度下降的“指标优度”最大化。
• 与概率编程语言（如 NumPyro）集成，实现在宇宙学生分析中可扩展、高性能的后验采样。

提出的方法

• 利用 JAX 的自动微分（autodiff）精确计算宇宙学可观测量相对于宇宙学参数和系统参数的梯度。
• 实现关键宇宙学组件的可微分版本：背景演化、线性和非线性功率谱（使用 Halofit 或 Eisenstein & Hu 模型）以及通过 Limber 近似计算的角功率谱（Cℓ）。
• 利用 JAX 的 JIT 编译和 GPU/TPU 加速，提升似然评估和采样算法的速度。
• 与概率编程框架（如 NumPyro）集成，支持基于梯度的推断技术，如 NUTS、ADVI 和神经流哈密顿蒙特卡洛（HMC）。
• 使用 vmap 实现宇宙学模拟和多 GPU 或 CPU 上的参数扫描的批量处理。
• 支持 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 中的自动微分能否在宇宙学中实现稳健、无需微调的费舍尔矩阵计算？
• RQ2 在真实宇宙学分析中，基于梯度的推断方法（如哈密顿蒙特卡洛和 NUTS）与传统 MCMC 相比，在效率上如何？
• RQ3 可微分似然能否实现通过在“指标优度”上进行梯度下降来直接优化探测器设计？
• RQ4 与优化后的经典 MCMC 相比，使用基于 JAX 的采样器在单位时间内有效样本数的性能提升如何？
• RQ5 带有神经流的神经传输 HMC 是否能在早期宇宙学分析中提升模型验证效果？

主要发现

• 在 DES-Y1 3x2pt 分析中，jax-cosmo 的 JAX-NUTS 采样器每 GPU 小时的有效样本数约为经过良好优化的 Cobaya Metropolis-Hastings 采样器的 5 倍。
• 使用 jax-cosmo 计算的费舍尔矩阵是精确且稳定的，不再需要手动微调，显著降低了计算和实现开销。
• 费舍尔矩阵本身也是可微分的，支持通过梯度下降对探测器设计参数（如红移分箱或观测策略）进行基于梯度的优化。
• 基于梯度的推断技术（如 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.