⚛️ general relativity

HIcosmo: a differentiable JAX-based framework for cosmology inference

本文介绍了 HIcosmo，这是一个基于 JAX 的开源宇宙学推断框架，它利用自动微分和 GPU 加速来实现比 Cobaya 等现有工具显著更快的采样速度，同时支持广泛的第四阶段巡天似然函数和宇宙学模型。

原作者： Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

发布于 2026-06-29

📖 1 分钟阅读🧠 深度阅读

原作者： Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，宇宙就像一个巨大而复杂的拼图，科学家们正试图解开它。几十年来，他们一直使用一套标准的规则（称为 $\Lambda$ CDM 模型）来研究宇宙是如何开始的、它是如何膨胀的，以及它是由什么组成的。但最近，拼图的碎片似乎无法完美契合。来自早期宇宙的测量结果（如宇宙微波背景）与来自近处宇宙的测量结果（如被称为超新星的爆炸恒星）存在分歧。这被称为“哈勃张力”（Hubble Tension）。

为了解决这个问题，科学家们正在建造新的大型望远镜，它们将收集数十亿个数据点。问题在于？科学家用来解决这些谜题的计算机程序太慢且太笨重，无法处理如此庞大的信息流。它们就像是试图用一个缓慢的计算器去数沙滩上的每一粒沙子。

这篇论文介绍了一个全新的、超快速的计算机工具——HIcosmo，旨在高效地解决这些宇宙谜题。以下是它的工作原理，通过简单的解释呈现：

1. “智能计算器”（JAX 与可微性）

传统的工具就像一位厨师，必须手动测量每一种配料，写下结果，然后再手动计算下一步。如果他们想知道改变某一种配料会对最终口味产生什么影响，他们只能通过猜测和尝试，这非常耗时。

HIcosmo 则不同。它基于一种被称为 JAX 的技术，JAX 就像是一个能同时理解整个食谱的“智能计算器”。

自动微分（Automatic Differentiation）： 它不再靠猜测，而是精确知道如果你微调一个数字，结果会如何变化。这就像拥有一个 GPS，它能立即告诉你最佳路线，而不需要通过驾驶每一条可能的道路来查看哪条路最快。
即时编译（Just-in-Time Compilation）： 当你第一次要求 HIcosmo 进行计算时，它需要花费极短的时间进行“预热”，并将指令转化为经过高度优化的机器码。在此之后，它会以闪电般的速度运行，绕过了会让旧工具变慢的“翻译器”（Python 解释器）。

2. 三层蛋糕

作者将 HIcosmo 构建得像一个三层蛋糕，每一层都有特定的职责：

底层（模型）： 这是“理论”部分。它根据宇宙的规则（例如暗物质或暗能量有多少）来计算宇宙应该呈现的样子。
中间层（似然）： 这是“比较”部分。它获取来自底层层的理论，并将其与来自望远ل镜的真实数据（如 Pantheon+ 超新星数据或 DESI 星系图）进行对比。它会计算该理论到底“错”在哪里。
顶层（推断）： 这是“求解器”部分。它利用来自中间层的信息来寻找最佳答案。它使用了一种称为 NUTS（无转向采样器）的聪明方法，利用“智能计算器”提供的梯度直接冲向正确答案，而不是像旧工具那样漫无目的地徘徊。

3. 它能做什么（以及不能做什么）

它能做的：
HIcosmo 目前是**背景宇宙学（background cosmology）**的大师。可以把这理解为测量宇宙的“大局”：它如何膨胀、物体距离有多远，以及空间的整体形状。它可以处理来自以下方面的数据：

爆炸的恒星（超新星）。
星系聚集（BAO）。
宇宙大爆炸的余晖（CMB 距离先验）。
引力透镜的时间延迟。

它已经针对目前的金标准软件（称为 Cobaya）进行了测试，结果显示其给出的答案在统计学上是完全一致的，但速度快得多。

在标准计算机处理器（CPU）上，它比旧工具快约 8.7 倍。
在强大的图形处理器（GPU）上（这种处理器非常擅长同时处理海量数据），随着数据量的增大，它的速度最高可提升 20 倍。

它不能做的：
论文非常明确地指出，HIcosmo 目前并不处理宇宙的“精细细节”或“扰动”，例如宇宙网中复杂的褶皱或早期宇宙详细的功率谱。这些需要另一个非常沉重的引擎（称为 Boltzmann 求解器）来处理，HIcosmo 将其留给了其他工具。你可以把 HIcosmo 想象成研究气球形状的专家，而研究橡胶纹理则需要其他工具。

4. 现实世界的成果

作者不仅制造了它，还用它解决了实际问题：

联合约束（Joint Constraints）： 他们结合了超新星、星系调查和 CMB 的数据，以惊人的精度锁定了宇宙的膨胀率（ $H_0$ ）和物质密度（ $\Omega_m$ ）。
未来预测（Future Forecasting）： 他们使用 HIcosmo 预测了未来的射电望远镜（如 SKA、MeerKAT 和 CHIME）能够多好地测量暗能量的“状态方程”（这是一个专业术语，意为询问：“暗能量是恒定的，还是随时间变化的？”）。他们发现，这些新望远镜可以极其精确地测量这一指标。

核心总结

HIcosmo 是一个现代化的、高速的宇宙学引擎。它用一个流线型的、自动化的系统取代了过去缓慢、手动的模式，能够应对来自下一代望远镜的海量数据洪流。它并没有改变物理定律，但它改变了我们测试这些定律并将其应用于宇宙的速度与效率。

技术摘要：HIcosmo —— 一个用于宇宙学推断的可微 JAX 框架

1. 问题陈述

即将到来的第四阶段（Stage IV）宇宙学巡天（如 Euclid、LSST、DESI、SKA）将产生前所未有的海量观测数据。虽然背景层级的宇宙学推断（涉及距离、距离比和距离模数）通常在低维参数空间（ $d \sim 5\text{--}7$ ）中运行，但由于马尔可夫链蒙特卡洛（MCMC）方法所需的似然评估次数极多，计算成本仍然是一个瓶颈。

传统的流水线（如 CosmoMC、Cobaya、MontePython）依赖于随机游走梅特罗波利斯-黑斯廷斯（MH）采样器。这些方法存在混合时间随维度（ $d$ ）线性增长的问题，并且需要有限差分近似来计算梯度，这在计算上非常昂贵（需要额外的 $O(d)$ 次评估），并且对步长选择非常敏感。此外，现有的成熟代码并非以可微程序的形式实现，这阻碍了使用基于梯度的采样器（如 No-U-Turn Sampler, NUTS）以及用于 Fisher 预报的自动微分（AD）。虽然已经出现了基于 JAX 的扰动级分析工具（如 jax-cosmo、CosmoPower），但目前还没有一个经过系统验证的、端到端的、用于背景层级宇宙学的可微框架，能够将似然函数、NUTS 采样和基于 AD 的 Fisher 预报集成在单一架构内。

2. 方法论与架构

本文介绍了 HIcosmo（高效率宇宙学推断），这是一个完全基于 JAX 构建的开源框架。该框架的设计理念围绕端到端可微性展开，确保从宇宙学参数到标量对数似然的所有操作都能被 JAX 的自动微分系统追踪。

2.1 分层架构

HIcosmo 由三个主要层组成：

模型层（Model Layer）： 计算理论预测。它实现了四种背景模型： $\Lambda$ CDM、 $w$ CDM、 $w_0w_a$ CDM（CPL 参数化）以及一种相互作用暗能量模型（ $I\Lambda$ CDM）。至关重要的是，它使用固定网格梯形求积法在稠密的红移网格上取代了自适应积分器（后者是不可微的）。这使得通过链式法则进行精确梯度计算成为可能，而无需使用有限差分。
似然层（Likelihood Layer）： 将模型预测与观测数据进行比较。它支持：
- Ia 型超新星： Pantheon+、DES-SN5YR 和 Union3。
- 重子声学振荡（BAO）： DESI DR1/DR2、SDSS、BOSS 和 6dFGS。
- CMB： Planck 2018 距离先验（压缩为 $R$ 、 $l_A$ 、 $\omega_b$ ）。
- 局部 $H_0$ ： SH0ES 距离阶梯。
- 强引力透镜： H0LiCOW 时间延迟。
- 所有似然函数均使用纯 JAX 实现，支持即时编译（JIT）和向量化。
推断层（Inference Layer）： 消耗对数似然及其导数。它提供：
- NUTS 采样： 通过 NumPyro 实现，直接利用来自 JAX 的梯度来改善混合时间（其缩放比例为 $d^{1/4}$ ）。
- Fisher 预报： 通过嵌套自动微分计算 Hessian 矩阵，消除了对步长选择的需求以及有限差分误差。

2.2 实现细节

可微性： 通过使用固定步长积分并预计算与参数无关的量（如逆协方差矩阵），该框架确保可以直接通过反向模式 AD 获取梯度和 Hessian 矩阵。
硬件加速： 框架利用 JAX 统一的 CPU/GPU/TPU 执行环境。JIT 编译将似效函数优化为机器码，从而绕过了 MCMC 运行期间的 Python 解释器开销。
校准： 为确保物理准确性，框架针对 Boltzmann 求解器 CAMB 校准了导出量（如声速标度 $r_d$ ），将残差降低至 $\sim 0.01\%$ 。

3. 关键结果

3.1 验证与精度

HIcosmo 在 $\Lambda$ CDM、 $w$ CDM 和 $w_0w_a$ CDM 模型下，针对参考实现（Cobaya）和官方似然代码进行了系统验证。

似然精度： $\chi^2$ 值与参考代码的绝对差异在 $10^{-6}$ 到 $10^{-2}$ 之间。
参数约束： HIcosmo 和 Cobaya 的边缘化约束在所有测试分析中差异均小于 0.2 $\sigma$ 。对于联合 SN+BAO+CMB 分析，一致性在 0.04 $\sigma$ 以内。
一致性： 该框架重现了已建立的约束，例如联合分析中的 $H_0 = 68.24 \pm 0.38$ km/s/Mpc 和 $\Omega_m = 0.307 \pm 0.005$ ，证实了可微近似并未引入显著偏差。

3.2 性能基准

相对于 Scipy 的加速： 对于距离计算，HIcosmo 在 JIT 编译后比 scipy 自适应求积快 140 $\times$ 。对于完整的 Pantheon+ 似然流水线，在 CPU 上快了 16 $\times$ 。
采样吞吐量： 与使用 MH 采样器的 Cobaya（在 8 核 CPU 上）相比，HIcosmo（使用 NUTS）实现了更高的有效样本量（ESS）吞吐量（400 ESS/s 对比 45.8 ESS/s）。这既得益于更快的似然评估，也得益于 NUTS 更好的混合性能（ $\tau \approx 1.2$ 对比 $\tau \approx 11.4$ ）。
GPU 加速： 虽然在小数据集（Pantheon+）上的单次评估由于开销原因 GPU 增益有限，但 GPU 加速会随数据集规模增大而提升。对于包含 $N_{SN} = 20,000$ 的合成数据集，GPU 比 CPU 快了 20.8 $\times$ 。
Fisher 矩阵： 基于 AD 的 Fisher 矩阵计算稳定，且不需要步长调优，避免了在小步长下困扰有限差分方法的浮点舍入误差。

3.3 应用

论文通过以下方式展示了该框架的实用性：

多探测器联合约束： 结合 Pantheon+、DESI BAO 和 Planck 先验，收紧对 $H_0$ 和 $\Omega_m$ 的约束。
暗能量约束： 推导出与 DESI 2024 结果一致的 $w_0$ 和 $w_a$ 约束。
21 cm 强度映射预报： 对六个巡天项目（SKA1、MeerKAT、BINGO、Tianlai、CHIME）进行 Fisher 预报。该框架预测 SKA1 Band 1 可以将暗能量状态方程约束在 $\sigma(w) \approx 0.046$ 。

4. 意义与主张

作者声称，HIcosmo 通过提供第一个经过系统验证的、端到端可微的 JAX 原生背景宇宙学流水线，填补了宇宙学推断领域的关键空白。

效率： 它证明了基于梯度的采样（NUTS）结合 JIT 编译的似然函数，在计算速度和统计效率（ESS）方面都能显著优于传统的随机游走 MCMC。
可扩展性： 该框架通过利用 GPU 加速（加速可达 20 $\times$ CPU 执行），能够独特地应对未来巡天（Stage IV）的海量数据集。
模块化与可扩展性： 通过将模型、似然和推断层分离，该框架允许在不重写导数代码的情况下轻松添加新模型或数据集。
面向扰动级推断的基础： 虽然目前仅限于背景宇宙学（不包括完整的 Boltzmann 求解器），但论文指出，在背景层级成功完成端到端验证，为更高维度的 Stage IV 问题奠定了稳健的基础。该架构旨在通过将模型层替换为可微 Boltzmann 模拟器（如 CosmoPower-JAX），从而扩展到扰动层级的分析。

总之，HIcosmo 提供了一个可重复、高性能的现有工具替代方案，证明了自动微分和现代硬件加速可以有效地应用于标准的宇宙学推断任务。