大规模可微天体物理：在GPU上求解和微分ODE系综
Differentiable astrophysics at scale: solving and differentiating ODE ensembles on the GPU
浏览论文内容
中文总结 AI 辅助
本文提出GRADSOLVE，一个基于JAX的GPU库，用于高效求解和微分大规模ODE系综，在百万轨迹规模下比CPU快8.5至1500倍，比现有JAX库快11至15倍，使梯度分析成为常规。
中文摘要 AI 辅助
天文学家越来越多地使用基于梯度的方法（如哈密顿蒙特卡洛）来拟合模型，这些方法需要模型关于其参数的导数。在许多分析中，一次预测需要为数千到数百万个参数集求解一个小型常微分方程组（ODE），而这一系综的积分往往决定了分析的成本。我们向天文学界介绍GRADSOLVE，一个JAX库，它将这一计算迁移到图形处理单元（GPU）上：它在每个GPU线程中独立积分系综的每个成员，使用各自的步长，并在同一次遍历中返回关于参数的导数。我们在来自天文学不同领域的三个示例中测量了其加速效果：银河系势中的恒星轨道、暗能量宇宙学的膨胀历史以及双黑洞的自旋进动，每个代码都满足相同的精度要求。对于一百万个轨迹，GRADSOLVE的运行速度比每个示例在CPU所有128个核上运行的串行CPU代码快8.5到1500倍，比在单核上快几个数量级。在同一GPU上，它也比JAX中最先进的ODE库DIFFRAX快11到15倍，在全部三个示例中均如此。使用GRADSOLVE，一百万个这样的积分在一个GPU上只需几秒或更短时间，因此需要如此规模系综的基于梯度的分析变得常规。代码公开于https://github.com/ECLIPSE-AI4Science/gradsolve。
英文摘要
Astronomers increasingly fit their models with gradient-based methods, such as Hamiltonian Monte Carlo, which need the derivatives of the model with respect to its parameters. In many analyses a prediction requires solving a small system of ordinary differential equations (ODEs) for thousands to millions of parameter sets, and this ensemble of integrations often sets the cost of the analysis. We introduce to the astronomical community GRADSOLVE, a JAX library that moves this computation to graphics processing units (GPUs): it integrates each member of the ensemble in its own GPU thread with its own step size and returns the derivatives with respect to the parameters in the same pass. We measure the speed-up it provides in three examples from different fields of astronomy, stellar orbits in the Galactic potential, the expansion history of a dark-energy cosmology and the spin precession of binary black holes, with every code held to the same accuracy requirement. For a million trajectories GRADSOLVE runs 8.5 to 1500 times faster than the serial CPU code of each example running on all 128 cores of a CPU, and several orders of magnitude faster than on one core. On the same GPU it is also 11 to 15 times faster than DIFFRAX, the state-of-the-art ODE library in JAX, in all three examples. With GRADSOLVE a million such integrations take seconds or less on one GPU, so gradient-based analyses that need ensembles of this size become routine. The code is publicly available at https://github.com/ECLIPSE-AI4Science/gradsolve.
发表机构
- Royal Holloway, University of London（伦敦大学皇家霍洛威学院）
机构由 AI 辅助整理，请以论文原文为准。
