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

解决问题:这篇论文旨在解决宇宙学推断中高维参数空间的问题,通过构建宇宙功率谱的神经仿真器,加速参数估计,以便更有效地探索高维参数空间。此外,论文还试图验证使用JAX框架和GPU加速的梯度优化采样技术在宇宙学研究中的应用。

关键思路:论文的关键思路是使用JAX框架和GPU加速的梯度优化采样技术,通过构建宇宙功率谱的神经仿真器,加速高维参数空间的参数估计。相比当前领域的研究,论文的思路在于使用JAX框架和GPU加速技术,以及神经仿真器的构建,使得高维参数空间的参数估计更加高效。

其他亮点:论文的实验结果表明,CosmoPower-JAX在两个模拟的Stage IV配置中表现出了高精度和高计算效率,其中一个配置包括37个模型参数,另一个配置包括157个模型参数。此外,CosmoPower-JAX是完全使用Python编写的,并已公开发布,以帮助宇宙学界满足下一代调查的精度要求。

关于作者:主要作者D. Piras和A. Spurio Mancini所在的机构分别是意大利国家核物理研究所和欧洲核子研究组织。他们之前的代表作包括Pandas和NumPy等Python数据分析库的开发和推广,以及在高能物理和宇宙学领域的研究。

相关研究:与此论文相关的其他研究包括“Fast cosmological parameter estimation using neural density estimators”(作者:J. M. Wang等,机构:斯坦福大学)、“CosmoFlow: Using Deep Learning to Learn the Universe at Scale”(作者:M. F. Mustafa等,机构:加州大学伯克利分校等)等。

论文摘要:本文介绍了基于JAX实现的CosmoPower-JAX框架,通过构建宇宙功率谱的神经仿真器来加速宇宙学推断。我们展示了如何利用JAX的自动微分、批量计算和即时编译功能,并在图形处理器上运行推断流程,利用先进的基于梯度的采样技术加速参数估计数倍。这些技术可用于高效地探索高维参数空间,例如下一代宇宙学调查所需的空间。我们在两个模拟的第四阶段配置上展示了CosmoPower-JAX的准确性和计算效率。首先,我们考虑一个单一的调查,对37个模型参数进行宇宙剪切分析。我们使用CosmoPower-JAX和哈密顿蒙特卡罗采样器得出的轮廓与使用嵌套采样器和没有仿真器得出的轮廓进行验证,获得了$\mathcal{O}(10^3)$的速度提升。然后,我们考虑了三个第四阶段调查的组合,每个调查都进行联合宇宙剪切和星系聚类(3x2pt)分析,共计157个模型参数。即使在如此高维的参数空间中,CosmoPower-JAX仍然可以在3天内提供收敛的后验轮廓,而标准方法需要估计6年的时间。CosmoPower-JAX完全使用Python编写,并公开提供,以帮助宇宙学界满足下一代调查设定的精度要求。

内容中包含的图片若涉及版权问题,请及时与我们联系删除