发表机构
Royal Holloway, University of London（伦敦大学皇家霍洛威学院）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
GRADSOLVE是开源JAX库，可在NVIDIA GPU上快速精确求导低维ODE集合，比现有工具在梯度计算上快5.6-14.1倍，支持多种积分器，已发布。
AI 中文摘要
常微分方程（ODE）是科学与工程领域模型的基础，许多应用需要其解关于参数的导数。独立轨迹的集合适合图形处理单元（GPU）运算，但当前GPU软件存在权衡：最快的集合求解器无法以其求解速度进行反向模式求导，而为求导构建的求解器求解速度更慢。目前尚无单一工具能以融合内核求解的速度提供反向模式梯度。我们提出GRADSOLVE，一个开源JAX库，用于在NVIDIA GPU上求解和反向模式求导低维ODE集合。它记录自适应求解器接受的步骤，并对这些步骤进行固定步长重放的求导；返回的梯度是这些步骤的精确离散伴随，与Diffrax默认返回的导数相同，且通过固定长度链而非自适应循环获得，成本更低。该库针对在一个记录网格上多次求导的集合，保留Diffrax作为备用，支持显式和Rosenbrock积分器。作为求解器使用时，GRADSOLVE的仅前向内核比this http URL快2.8倍；用于梯度计算时，在存在记录后，跨三代GPU，其计算梯度的速度比Diffrax的检查点伴随快5.6至14.1倍，优势在大集合上缩小，在刚性系统上，精度要求严格时则持平。GRADSOLVE已发布于this https URL。
英文摘要
Ordinary differential equations (ODEs) underlie models in science and engineering, and many applications need derivatives of their solutions with respect to parameters. Ensembles of independent trajectories suit graphics processing units (GPUs), but current GPU software forces a trade-off: the fastest ensemble solvers cannot be differentiated in reverse mode at the speed they solve, and the solvers built for differentiation solve more slowly. No single tool has yet offered a reverse-mode gradient at the speed of a fused-kernel solve. We present GRADSOLVE, an open-source JAX library for solving and reverse-mode differentiating low-dimensional ODE ensembles on NVIDIA GPUs. It records the steps an adaptive solver accepts and differentiates a fixed-step replay of them; the returned gradient is the exact discrete adjoint of those steps, the same derivative Diffrax returns by default, obtained more cheaply from a fixed-length chain than from an adaptive loop. It targets ensembles differentiated many times against one recorded mesh, keeps Diffrax as a fallback, and supports explicit and Rosenbrock integrators. Used as a solver, GRADSOLVE's forward-only kernel ran 2.8x faster than DiffEqGPU.jl; used for gradients, once a record exists, it computed them 5.6-14.1x faster than Diffrax's checkpointed adjoint at matched forward-state accuracy across three GPU generations, the advantage narrowing on large ensembles and, on stiff systems, down to parity at tight accuracy. GRADSOLVE is released at https://github.com/ECLIPSE-AI4Science/gradsolve.
Comments38 pages, 12 figures. GRADSOLVE available at https://github.com/ECLIPSE-AI4Science/gradsolve
