[论文解读] Capse.jl: efficient and auto-differentiable CMB power spectra emulation

Marco Bonici, F. Bianchini+1 more|arXiv (Cornell University)|Jul 26, 2023

Cosmology and Gravitation TheoriesPhysics and Astronomy参考文献 67被引用 3

一句话总结

Capse.jl 是一个基于神经网络的模拟器，可在微秒级别内以低于 0.1σ 的误差预测 CMB 温度、极化与引力透镜功率谱，覆盖 CMB-S4 相关的尺度。通过结合可微分架构与基于梯度的采样器（如 NUTS 和 Pathfinder），该方法将贝叶斯推断速度提升至传统方法的 6 个数量级，实现与普朗克（Planck）相当的宇宙学约束，且耗时不足一秒。

ABSTRACT

We present Capse.jl, a novel neural network-based emulator designed for rapid and accurate prediction of Cosmic Microwave Background (CMB) temperature, polarization, and lensing angular power spectra. The emulator computes predictions in just a few microseconds with emulation errors below $0.1σ$ for all the scales relevant for the upcoming CMB-S4 survey. Capse.jl can also be trained in an hour's time on a 8-cores CPU. We test Capse.jl on Planck 2018, ACT DR4, and 2018 SPT-3G data and demonstrate its capability to derive cosmological constraints comparable to those obtained by traditional methods, but with a computational efficiency that is three to six orders of magnitude higher. We take advantage of the differentiability of our emulators to use gradient-based methods, such as Pathfinder and Hamiltonian Monte Carlo (HMC), which speed up the convergence and increase sampling efficiency. Together, these features make Capse.jl a powerful tool for studying the CMB and its implications for cosmology. When using the fastest combination of our likelihoods, emulators, and analysis algorithm, we are able to perform a Plancky TT + TE + EE analysis in less than a second. To ensure full reproducibility, we provide open access to the codes and data required to reproduce all the results of this work.

研究动机与目标

• 解决使用传统 Boltzmann 求解器（如 CAMB）进行贝叶斯宇宙学推断时 CMB 似然评估的高计算成本问题。
• 开发一种高效、精确且可微分的模拟器，用于温度、极化与引力透镜的 CMB 角功率谱。
• 利用基于梯度的推断方法（如哈密顿蒙特卡洛与 Pathfinder）实现快速采样与优化。
• 与使用 CAMB 和 Cobaya 的标准 MCMC 流程相比，实现 3 至 6 个数量级的计算加速。
• 确保代码与训练数据的完全可复现性与开源，以促进社区采纳与验证。

提出的方法

• 在广泛的宇宙学参数空间内训练深度神经网络，以模拟 CMB 角功率谱（TT、TE、EE、BB 与引力透镜）。
• 利用 Julia 中的自动微分技术，使模拟器完全可微分，从而支持基于梯度的推断算法。
• 应用切比雪夫多项式分解以降低输出维度，加速似然评估，提升计算效率。
• 采用多尺度训练策略，结合合成参数组合，确保在 CMB-S4 相关全多极数范围内的准确性。
• 将模拟器集成至完整的贝叶斯推断流程中，使用 Cobaya 框架替换似然评估中的 CAMB。
• 使用基于梯度的采样器（NUTS、微正则 HMC 与 Pathfinder）进行性能基准测试，比较收敛性与采样效率。

Figure 1 : Comparison of the relative differences between Capse.jl and high-accuracy predictions from CAMB for the four 2-point statistics emulated in this study. The differences are normalized to the forecasted CMB-S4 statistical uncertainties ( $f_{\mathrm{sky}}=0.4$ ) and measured in units of cos

实验结果

研究问题

• RQ1 可微分神经网络模拟器是否能在未来巡天相关的所有尺度上实现低于 0.1σ 的 CMB 功率谱预测精度？
• RQ2 当与可微分模拟器结合使用时，NUTS 与 Pathfinder 等基于梯度的采样器在多大程度上能加速宇宙学推断？
• RQ3 在普朗克（Planck）、ACT DR4 与 SPT-3G 数据集上，Capse.jl 的计算性能与基于 CAMB 的传统流程相比，其实际运行时间如何？
• RQ4 切比雪夫多项式分解是否能在不损失精度的前提下显著提升似然评估效率？
• RQ5 该模拟器是否足够稳健与精确，足以在真实数据上重现与 CAMB 和 Cobaya 相当的宇宙学约束？

主要发现

• Capse.jl 每次调用计算 CMB 功率谱的时间少于 10 微秒，且在 CMB-S4 相关的所有尺度上，模拟误差均低于 0.1σ。
• 在普朗克 2018、ACT DR4 与 SPT-3G 2018 数据上，该模拟器实现了与 CAMB 基准流程相当的完整宇宙学约束，后验分布差异小于 0.1σ。
• 使用 Capse.jl、最快似然组合与 NUTS 采样进行的普朗克 TT+TE+EE 分析，耗时不足一秒，而使用 CAMB 与 Cobaya 则需约 100 CPU 小时。
• 使用 NUTS 与 Pathfinder 等基于梯度的采样器，可将分析时间缩短 2 至 6 个数量级：普朗克为 0.5 CPU 小时，ACT 为 1 小时，SPT-3G 为 14 小时。
• 微正则 HMC 在效率上比标准 NUTS 提升 2 至 3 倍，同时产生几乎相同的后验分布。
• 切比雪夫多项式分解使普朗克分析效率提升近一个数量级，并支持可扩展、高精度的模拟，未来可拓展至星系大尺度结构分析。

Figure 2 : Triangle plot, showing the standard Planck lite chains and the one obtained using Capse.jl in combination with Turing.jl . The lower triangular part of the plot focuses on the comparison between CAMB and Capse.jl , in order to ensure the precision of our emulators and likelihoods. The upp