[论文解读] pmwd: A Differentiable Cosmological Particle-Mesh $N$-body Library

Yin Li, Libin Lu+8 more|arXiv (Cornell University)|Nov 18, 2022

Galaxies: Formation, Evolution, Phenomena被引用 14

一句话总结

pmwd 是一个可微分的宇宙学粒子网格（$N$-体）模拟库，利用伴随法（adjoint method）与 JAX 的自动微分功能，实现了内存高效、GPU 加速的大尺度结构形成正向建模。通过时间反演积分来重构模拟历史，其在保持完全可微性的同时，相比标准的反向传播时间（backpropagation-through-time）方法，显著降低了内存开销。

ABSTRACT

The formation of the large-scale structure, the evolution and distribution of galaxies, quasars, and dark matter on cosmological scales, requires numerical simulations. Differentiable simulations provide gradients of the cosmological parameters, that can accelerate the extraction of physical information from statistical analyses of observational data. The deep learning revolution has brought not only myriad powerful neural networks, but also breakthroughs including automatic differentiation (AD) tools and computational accelerators like GPUs, facilitating forward modeling of the Universe with differentiable simulations. Because AD needs to save the whole forward evolution history to backpropagate gradients, current differentiable cosmological simulations are limited by memory. Using the adjoint method, with reverse time integration to reconstruct the evolution history, we develop a differentiable cosmological particle-mesh (PM) simulation library pmwd (particle-mesh with derivatives) with a low memory cost. Based on the powerful AD library JAX, pmwd is fully differentiable, and is highly performant on GPUs.

研究动机与目标

• 开发一个支持基于梯度的推断以约束宇宙学参数的可微分宇宙学模拟库。
• 通过伴随法克服标准反向传播时间方法在可微分模拟中面临的内存瓶颈。
• 实现高分辨率、GPU 加速的模拟，支持场级正向建模的完全可微性。
• 支持与深度学习及半解析模型集成，实现宇宙学参数与初始条件的联合推断。
• 提供可扩展、高性能且开源的工具，服务于下一代宇宙学数据分析。

提出的方法

• 采用伴随法通过时间反演积分计算梯度，避免存储完整的前向演化历史。
• 利用 JAX 实现自动微分与 GPU 加速，支持高性能可微分模拟。
• 采用“先离散化后优化”（discretize-then-optimize）方法，确保反向梯度与前向积分遵循相同的离散轨迹。
• 实现一种具有可逆时间积分的粒子网格（PM）引力求解器，以保证伴随计算的一致性。
• 在设计中引入 C2 对称性以体现时间可逆性，支持内存高效的梯度计算。
• 支持与可微组件（如代理模型和时空优化器）的集成。

实验结果

研究问题

• RQ1 伴随法能否有效应用于宇宙学 $N$-体模拟，以在保持可微性的同时显著降低内存开销？
• RQ2 与存储完整前向状态的现有可微分模拟代码相比，pmwd 的内存效率如何？
• RQ3 pmwd 在可微分宇宙学模拟中，能够多大程度地扩展到更高的时间步数与质量分辨率？
• RQ4 pmwd 是否能够实现宇宙学观测量（包括初始条件与物理参数）的端到端可微分建模？
• RQ5 pmwd 与深度学习及其他可微组件在统一推断流程中的集成效果如何？

主要发现

• 通过使用伴随法而非存储所有前向状态，pmwd 实现了完全可微性，同时显著降低了内存使用。
• 该库与 JAX 完全兼容，并可在 GPU 上高效运行，支持高性能模拟。
• 伴随法使 pmwd 能够扩展到比以往基于自动微分的模拟更多的时间步与更高分辨率。
• pmwd 支持场级正向建模，可实现宇宙学参数与初始条件的联合推断。
• 该库设计为可与可微分深度学习和半解析模型互操作，支持构建复杂的推断流程。
• pmwd 为开源并公开归档，支持可复现性与社区扩展。