[论文解读] CosmicRIM : Reconstructing Early Universe by Combining Differentiable Simulations with Recurrent Inference Machines

Chirag Modi, François Lanusse+3 more|arXiv (Cornell University)|Apr 26, 2021

Computational Physics and Python Applications参考文献 17被引用 5

一句话总结

CosmicRIM 提出了一种可微分的循环推理机（RIM）框架，该框架将可微分的N体模拟（FlowPM）与学习优化相结合，用于重建宇宙学中的原始初始条件。该方法在收敛速度上比 ADAM 和带退火的 LBFGS 快 40 倍，同时在大尺度上保持或提升了重建质量，通过隐式学习一种受物理启发的退火调度方案实现。

ABSTRACT

Reconstructing the Gaussian initial conditions at the beginning of the Universe from the survey data in a forward modeling framework is a major challenge in cosmology. This requires solving a high dimensional inverse problem with an expensive, non-linear forward model: a cosmological N-body simulation. While intractable until recently, we propose to solve this inference problem using an automatically differentiable N-body solver, combined with a recurrent networks to learn the inference scheme and obtain the maximum-a-posteriori (MAP) estimate of the initial conditions of the Universe. We demonstrate using realistic cosmological observables that learnt inference is 40 times faster than traditional algorithms such as ADAM and LBFGS, which require specialized annealing schemes, and obtains solution of higher quality.

研究动机与目标

• 解决从宇宙学巡天数据中重建原始密度场这一高维、非线性逆问题。
• 克服传统优化方法（如 ADAM 和带退火的 LBFGS）在计算上不可行的问题，这些方法需要数百次昂贵的 N 体模拟。
• 学习一种高效的推理方案，隐式复制基于物理的退火过程，而无需手动调参。
• 为下一代需要重复重建的宇宙学巡天提供快速、可扩展的推理能力。

提出的方法

• 该方法使用可微分的 N 体模拟器（FlowPM）计算前向模型相对于初始条件的梯度。
• 它采用一种具有循环更新规则的循环推理机（RIM），结合隐藏状态、当前重建结果和后验对数梯度。
• RIM 经过训练以最小化负对数后验，学习一种隐式模仿物理退火的优化路径。
• 该架构包含一个记忆状态以及用于重建和隐藏状态的独立更新函数，支持迭代优化。
• 框架在初始条件上使用高斯先验，并通过添加高斯噪声的似然模型定义后验分布。
• 训练在具有固定基准功率谱的模拟数据上进行，支持端到端优化推理网络。

实验结果

研究问题

• RQ1 学习的优化方案是否能在重建原始密度场方面超越传统的基于梯度的方法（如 ADAM 和 LBFGS）？
• RQ2 RIM 框架是否在未显式设计的情况下隐式学习到一种受物理启发的退火调度？
• RQ3 与专门设计的退火方案相比，RIM 的性能在收敛速度和重建精度方面如何？
• RQ4 RIM 框架是否能够被训练以泛化到不同的初始条件和巡天配置？

主要发现

• CosmicRIM 在重建质量上与 ADAM 和带退火的 LBFGS 相当或更优，尤其在大尺度上显著改善了转移函数。
• 该方法仅需 10 次推理步骤，相比需要 400 步的退火 ADAM 和 LBFGS 实现了 40 倍的速度提升。
• 当以标准重建结果作为初始化时，CosmicRIM 进一步提升了性能，表明其对初始化具有鲁棒性。
• RIM 隐式学习到一条先重建大尺度结构的优化路径，与基于物理的退火方案行为一致。
• 从 RIM 最终输出开始，再使用 ADAM 或 LBFGS 进行进一步优化无法带来性能提升，表明 RIM 已达到高质量的 MAP 估计。
• 该框架支持可扩展的推理，可通过分阶段训练子网络实现，从而降低内存开销。