现代天文学已经进入了一个科学家不仅是在观测天空，还在以惊人的精度模拟天空的时代，以此来理解宇宙是如何运作的。为了验证他们的想法，研究人员构建了复杂的数学模型，这些模型描述了从我们银河系中恒星的旋转路径到空间本身膨胀的一切现象。这些模型通常依赖于求解一种被称为常微分方程的特定类型的数学难题。可以将这些方程想象成一套指令，告诉计算机一个系统如何随时间变化，例如一颗恒星在引力作用下如何运动，或者宇宙如何随着年龄的增长而拉伸。为了找到最理想的模型，天文学家必须将这些指令运行数百万次，每次都微调数值，以观察哪个版本与现实世界最为吻契合。这一过程是现代科学发现的支柱，但它遇到了一个瓶颈：所需的计算量如此巨大，以至于即使是最快的超级计算机也难以在合理的时间内完成任务。
现在，一位研究人员通过将这些计算从标准计算机处理器转移到图形处理器（最初为渲染视频游戏而设计的强大芯片）上，突破了这一瓶颈。他们开发了一种名为 gradsolve 的新软件工具，该工具将数百万次模拟中的每一次都视为一个独立的“工人”，并为每一个模拟在计算机芯片上分配一条专属的通道。在传统方法中，计算机尝试一个接一个地或以僵化的分组方式解决这些问题，这往往迫使快速完成的模拟闲置等待，直到最慢的模拟追赶上来。这种新方法允许每一次模拟都以自己的速度运行，在准备就绪时即刻完成，而不会拖累其他人。其结果是实现了速度的飞跃，将曾经需要数天才能完成的任务缩短到了仅需数秒。
研究人员在天文学领域内三个截然不同的问题上测试了这种方法，以证明其通用性。首先，他们观察了从球状星团中被剥离出来、形成横跨银河系的长而细的星流的恒星路径。通过同时模拟一百万颗恒星的运动，他们的新工具比在拥有 128 个处理器的标准计算机上运行的最优现有软件快了 1500 倍。其次，他们研究了宇宙膨胀的历史，特别是测试关于暗能量（驱动宇宙不断扩张的神秘力量）的理论。在这里，该工具计算了一百万种不同版本的宇宙距离，同样比单台计算机核心快了 1500 倍。最后，他们检查了两颗黑洞在碰撞前相互螺旋靠近时的扭转运动（或称进动）。对于一百万个双黑洞群体，该新软件比天文学家目前使用的标准代码快了 760 倍。
这项成就之所以特别显著，是因为这个新工具不仅运行得更快，而且在同一次处理中就能计算出必要的导数（即变化率）。在过去，寻找这些变化率需要进行另一次同样昂贵的单独计算，这会使所需时间翻倍。通过同时处理模拟及其对变化的敏感度，研究人员消除了进行复杂统计分析的一个主要障碍。这种速度增益是如此巨大，以至于以前因计算成本过高而无法进行的分析，现在正变得常规化。例如，一项单台计算机核心需要近一年才能完成的任务，现在可以在单个图形卡上仅用数秒即可完成。
研究人员非常谨慎地确保了这种巨大的速度提升并不会以牺牲准确性为代价。他们为结果与已知的、高精度的参考解之间的接近程度设定了严格的标准。在每一次测试中，新软件都达到了这些严苛的精度要求，证明了这种速度提升源于更聪明的组织工作方式，而非通过偷工减料实现的。该研究明确表明，对于涉及数百万个独立计算且不需要复杂的“刚性”数学步骤的问题，这种新方法优于目前最先进的方法。它表明，通过重新思考计算机如何处理大规模模拟，天文学家现在可以以前所未有的细节和速度探索宇宙，从而为研究从我们的银河系结构到暗能量基本性质的更精确模型打开了大门。
技术摘要：大规模可微天体物理学
问题陈述
现代天体物理推断日益依赖于基于梯度的算法，例如哈密顿蒙特卡洛（HMC）和费舍尔预报（Fisher forecasts），这些方法需要计算模型预测相对于其参数的导数。这类分析中的一个常见计算瓶颈是需要求解常微分方程（ODE）系综——其规模从数千个到数百万个轨迹不等，每个轨迹对应于不同的参数集或初始条件。
目前的方法面临两个主要局限性：
- 串行 CPU 执行： 在中央处理器（CPU）上求解这些系综在计算上非常昂贵，往往限制了分析的规模。
- 向量化 GPU 局限性： JAX 中最先进的可微 ODE 库（如
diffrax）通常将系综视为单个数组进行向量化。在自适应求解器中，步长根据局部误差估计而变化，这种“同步锁步”（lock-step）执行强制所有轨迹在完成一步时必须等待最慢的一个成员。因此，计算成本由需要最多步数的轨迹决定，而非平均步数，这导致了 GPU 上的显著效率低下。
方法论
作者引入了 gradsolve，这是一个旨在通过“每线程一个轨迹”架构将 ODE 系综的积分转移到 GPU 上的 JAX 库。
- 每线程集成： 与向量化方法不同，
gradsolve 生成一个 GPU 核函数（kernel），其中每个线程独立集成单个轨迹。每个线程维护自己的状态、中间评估值和灵敏度于私有内存中。这使得轨迹可以以各自的自适应步长进行推进，而无需等待其他轨迹，仅在 NVIDIA GPU 的 32 个线程组（warps）之间进行同步。
- 前向灵敏度： 该库在同一个核函数内，与 ODE 同步集成前向灵敏度方程（状态对参数的导数）。这种前向模式微分对于参数较少而输出较多的问题非常高效，避免了反向模式（伴随）微分需要存储或重新计算前向轨迹的内存开销。
- 稠密输出： 求解器提供稠密输出，允许在任意观测时间进行插值，而不必强制求解器在这些特定时刻终止步骤。
- 比较协议： 为了确保基准测试的公平性，作者在严格的预定义精度要求下，将
gradsolve 与 diffrax 以及串行 CPU 代码（如 gala、scipy、precession）进行对比。只有当代码在整个系综中的 99 分位误差满足物理精度阈值（例如轨道位置或距离的相对误差）时，才会被接受。
核心贡献
- 库开发： 发布了
gradsolve，这是一个 JAX 原生的库，用于生成用于自适应 ODE 集成和前向灵敏度计算的每线程 GPU 核函数。
- 架构转变： 证明了对于非刚性、小系统的 ODE 系综，每线程方法在 GPU 上显著优于向量化批处理方法，因为它消除了由步长异质性导致的空闲时间。
- 全面的基准测试： 在三个截然不同的天体物理领域提供了严谨的性能比较：
- 恒星流： 在扁平银河系势场中积分轨道及其对势参数的灵敏度。
- 暗能量： 计算多种宇宙学模型（包括空间平坦性约束解）下的膨胀历史和距离。
- 双黑洞： 模拟双黑洞种群中的自旋进动。
结果
论文报告了对于 106 条轨迹的系综，将 gradsolve 与串行 CPU 代码（运行在拥有 128 个核心的 AMD EPYC 9745 上）以及在同一 NVIDIA H200 GPU 上运行的最快被接受的 diffrax 配置进行对比后的加速倍数。
- 对比串行 CPU：
gradsolve 比运行在 128 个核心上的串行 CPU 代码快 8.5 到 1,500 倍。相对于单个 CPU 核心，其加速比高出数个数量级。
- 恒星轨道： 比 128 个核心快 8.5 倍。
- 暗能量： 比 128 个核心快 1,500 倍。
- 自旋进动： 比 128 个核心快 760 倍。
- 对比
diffrax（同一 GPU）： 对于大型系综（106 个成员），gradsolve 在保持相同精度的同时，比 diffrax 快 11 到 15 倍。
- 恒星轨道： 轨道 + 导数，快 15 倍。
- 暗能量： 包含空间平坦性求解，快 11 倍。
- 自旋进动： 快 11 倍。
- 精度： 所有测试的代码均达到了预定义的精度要求。
gradsolve 的 99 分位误差仍保持在各应用的物理极限内（例如，轨道位置为 10−4，进动角为 10−7 rad）。
- 加速来源： 当使用相同的数值方法（如 Verner 7 或 Tsit5）和容差比较
gradsolve 和 diffrax 时，步尝试次数几乎相同。性能差距归因于实现方式：gradsolve 避免了每一步都将整个批次状态读写到全局内存的带宽开销，并消除了向量化循环中“等待”的惩罚。
意义与范围
论文声称，gradsolve 通过将单次 GPU 上的计算时间从数小时缩短至数秒，使需要数百万次 ODE 集成的梯度分析变得“常规化”。
- 适用性： 该方法针对非刚性系统、小状态向量（3–9 个分量）以及少量微分参数（前向模式效率）进行了优化。
- 局限性： 论文指出，对于大状态向量（如网格流体力学或粒子网格模拟）或刚性系统（如爱因斯坦-玻尔兹曼方程），其他方法（如批处理数组求解器或隐式/Rosenbrock 方法）更为合适。虽然
gradsolve 包含一个用于刚性问题的 Rosenbrock 核函数，但主要的基准测试侧重于非刚性的天体物理应用。
- 对推断的影响： 在暗能量轮廓似然性示例中，加速将 ODE 求解的成本降低到了总迭代成本的一个极小部分，从而将计算瓶颈转移到了似然计算本身。这使得在可行的时间范围内探索更复杂的模型和更大的参数空间成为可能。
代码已公开，允许天文学界将这些高性能可微求解器集成到现有的 JAX 推断流水线中。
每周获取最佳 general relativity 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题，再试一次？
无垃圾邮件，随时退订。