[论文解读] CONNECT: A neural network based framework for emulating cosmological observables and cosmological parameter inference

Andreas Nygaard, Emil Brinch Holm+2 more|arXiv (Cornell University)|May 30, 2022

Galaxies: Formation, Evolution, Phenomena被引用 5

一句话总结

CONNECT 是一个基于 PyTorch 的神经网络框架，可模拟 CLASS 的玻尔兹曼求解器产生的宇宙学观测量，实现快速且精确的宇宙学参数推断。通过使用迭代优化的拉丁超立方采样生成高效训练数据，其将 CLASS 求解次数减少了两个数量级，并使 MCMC 推断速度提升多个数量级，与完整 CLASS 运行相比，后验分布差异小于 0.1 个标准差。

ABSTRACT

Bayesian parameter inference is an essential tool in modern cosmology, and typically requires the calculation of $10^5$--$10^6$ theoretical models for each inference of model parameters for a given dataset combination. Computing these models by solving the linearised Einstein-Boltzmann system usually takes tens of CPU core-seconds per model, making the entire process very computationally expensive. In this paper we present extsc{connect}, a neural network framework emulating extsc{class} computations as an easy-to-use plug-in for the popular sampler extsc{MontePython}. extsc{connect} uses an iteratively trained neural network which emulates the observables usually computed by extsc{class}. The training data is generated using extsc{class}, but using a novel algorithm for generating favourable points in parameter space for training data, the required number of extsc{class}-evaluations can be reduced by two orders of magnitude compared to a traditional inference run. Once extsc{connect} has been trained for a given model, no additional training is required for different dataset combinations, making extsc{connect} many orders of magnitude faster than extsc{class} (and making the inference process entirely dominated by the speed of the likelihood calculation). For the models investigated in this paper we find that cosmological parameter inference run with extsc{connect} produces posteriors which differ from the posteriors derived using extsc{class} by typically less than $0.01$--$0.1$ standard deviations for all parameters. We also stress that the training data can be produced in parallel, making efficient use of all available compute resources. The extsc{connect} code is publicly available for download at \url{https://github.com/AarhusCosmology}.

研究动机与目标

• 为解决宇宙学中贝叶斯参数推断的高计算成本问题，该问题通常需要 10^5–10^6 次爱因斯坦-玻尔兹曼求解器的调用。
• 通过开发一种新型迭代采样策略，减少训练神经网络模拟器所需的昂贵 CLASS 求解次数。
• 创建一个即插即用的框架，使用户可在不重新训练的情况下，利用已训练的神经网络实现快速、精确的参数推断，适用于新数据集。
• 使模拟器与标准 MCMC 求解器（如 MontePython 和 Cobaya）兼容，确保在现有宇宙学工作流中具有广泛适用性。
• 提供完全可复现的开源解决方案，支持自定义宇宙学输出，并具备 GPU 加速潜力。

提出的方法

• 使用全连接前馈神经网络来模拟由 CLASS 计算的宇宙学观测量（例如，CMB 功率谱、物质功率谱）。
• 采用迭代拉丁超立方采样策略生成高质量训练数据，使所需 CLASS 求解次数减少两个数量级。
• 在宇宙学参数及其对应 CLASS 输出的数据集上训练神经网络，输出变量可通过输入配置文件自定义。
• 通过插件机制与 MontePython 和 Cobaya 无缝集成，将 CLASS 调用替换为快速的神经网络推理。
• 通过模块化代码结构支持用户自定义损失函数和网络架构，便于针对特定宇宙学模型进行定制。
• 支持并行生成训练数据，充分利用高性能计算资源，实现高效的数据生成。

实验结果

研究问题

• RQ1 在仅使用少量 CLASS 求解次数训练的神经网络模拟器，是否能实现对宇宙学观测量的高精度模拟？
• RQ2 迭代采样策略在多大程度上可减少所需 CLASS 求解次数，同时保持模拟器的准确性？
• RQ3 使用模拟器推导的宇宙学参数后验分布与完整 CLASS 计算结果相比，其精度如何？
• RQ4 模拟器是否可在不重新训练的情况下跨不同数据组合重复使用，其对推断速度和精度的影响如何？
• RQ5 该框架在多大程度上可无缝集成到如 MontePython 等现有 MCMC 推理流程中，且仅需极少代码修改？

主要发现

• 与标准推断运行相比，CONNECT 将训练所需的 CLASS 求解次数减少了两个数量级，显著降低了数据生成成本。
• 使用 CONNECT 进行宇宙学参数推断得到的后验分布与完整 CLASS 计算结果相比，所有参数的差异均小于 0.01–0.1 个标准差。
• 该框架使 MCMC 推断在数小时内完成，而非数天或数周，且推断时间主要由似然评估决定，而非模型计算。
• 训练好的模拟器可在不重新训练的情况下跨不同数据组合完全复用，实现在新增数据时快速重新分析。
• 该框架与 MontePython 和 Cobaya 兼容，插件接口无需修改原始代码库。
• 训练数据可并行生成，充分利用可用的 HPC 资源，实现可扩展的数据生成。