GRADSOLVE: fast exact gradients for ODE ensembles on GPUs
GRADSOLVE is a JAX-based GPU library enabling fast, exact reverse-mode gradients for low-dimensional ODE ensembles via record-and-replay technique.
Key Findings
Methodology
GRADSOLVE records the accepted steps of an adaptive ODE solver, then replays these fixed steps using a fused GPU kernel for gradient computation. This approach yields exact discrete adjoint gradients with speedups of 5.6-14.1× over checkpointed adjoint methods like Diffrax, across three GPU generations. It supports explicit and Rosenbrock integrators, optimized for low-dimensional ensembles differentiated multiple times against a single recorded mesh, significantly reducing computational overhead while maintaining accuracy.
Key Results
- On three GPU architectures, the forward kernel runs 2.8× faster than DiffEqGPU, while gradient computations, once the record exists, are 5.6-14.1× faster than Diffrax’s checkpointed adjoint at matched accuracy. For Lorenz system ensembles, gradient speedup is approximately 8×, with the advantage diminishing for larger ensembles and stiff systems, where parity is observed at tight tolerances.
- In large-scale experiments, the speed advantage persists, especially in multiple differentiation passes, reducing the overall cost of sensitivity analysis and parameter calibration. The approach scales well across GPUs, including A100, H100, and RTX 4090, demonstrating broad applicability.
- Compared to PyTorch-based solvers like torchode and torchdiffeq, GRADSOLVE achieves 3.2–9.3× speedups in gradient computation, with the largest gains in small, non-stiff problems, owing to GPU fusion and fixed-step replay, which avoid kernel launch overheads and storage costs.
Significance
This work addresses a longstanding bottleneck in GPU-based scientific computing: the trade-off between fast ensemble solving and efficient reverse-mode differentiation. By enabling exact gradients at the speed of fused kernels, GRADSOLVE unlocks new possibilities in neural differential equations, parameter estimation, and sensitivity analysis, making large-scale, high-accuracy simulations feasible and more accessible. It bridges the gap between high-speed forward solving and gradient-based optimization, essential for advancing scientific machine learning and complex system modeling.
Technical Contribution
GRADSOLVE introduces a novel record-and-replay framework that leverages fixed-step solutions to produce exact discrete adjoint gradients efficiently. Its core innovation lies in compiling the recorded steps into a single fused GPU kernel, eliminating the need for storing intermediate states or recomputing during backpropagation. This approach maintains the accuracy of the discrete adjoint while drastically reducing computational costs, supporting multiple integrators and GPU architectures. It represents a fundamental leap in differentiable scientific computing on accelerators.
Novelty
This is the first implementation to combine GPU fusion with fixed-step replay for discrete adjoint differentiation of low-dimensional ODE ensembles. Unlike prior methods relying on adaptive step differentiation or continuous adjoints, GRADSOLVE achieves high speed and exactness simultaneously, establishing a new standard for differentiable GPU-based ODE solvers. Its record-and-replay paradigm fundamentally changes how sensitivities are computed in large-scale scientific applications.
Limitations
- The current design primarily targets low-dimensional systems; high-dimensional problems may face increased memory and computational challenges, limiting scalability.
- For highly stiff or complex systems, fixed-step replay may require very fine meshes, reducing speed advantages and possibly affecting accuracy.
- The method assumes the recorded steps are sufficient; in highly nonlinear or chaotic systems, the fixed mesh might not capture all dynamics, necessitating adaptive strategies.
Future Work
Future directions include extending support to high-dimensional systems, integrating adaptive step control within the replay framework, and optimizing memory management for large-scale problems. Further, incorporating higher-order integrators and automatic mesh refinement could enhance robustness. Expanding compatibility with other frameworks and exploring hybrid methods combining adaptive and fixed-step approaches are also planned to broaden applicability.
AI Executive Summary
Efficiently computing gradients of solutions to ordinary differential equations (ODEs) is vital across scientific disciplines, especially in machine learning, parameter estimation, and sensitivity analysis. Traditional GPU-based solvers excel in speed but struggle with differentiability due to the adaptive nature of their algorithms. Conversely, differentiable solvers like Diffrax and torchdiffeq support reverse-mode AD but are slower, primarily because they differentiate the entire adaptive loop, incurring significant overhead.
GRADSOLVE introduces a novel record-and-replay approach that bridges this gap. It first runs an adaptive solver once, recording the accepted steps to establish a fixed mesh. Then, it replays these steps using a fused GPU kernel, enabling exact discrete adjoint gradients to be computed efficiently. This method leverages GPU fusion, avoiding the costly storage and recomputation typical of traditional reverse-mode differentiation. Experimental results demonstrate that, on three GPU architectures, GRADSOLVE achieves speedups of 5.6-14.1× over existing methods for gradient calculations, while maintaining the same accuracy.
The core innovation lies in compiling the recorded steps into a single GPU kernel, which allows for rapid, exact gradient computation without sacrificing the flexibility of adaptive solvers. It supports multiple integrators, including explicit and Rosenbrock schemes, making it versatile for various problem types, from non-stiff to stiff systems. The approach is particularly effective for low-dimensional ensembles differentiated repeatedly against a single recorded mesh, common in parameter calibration and sensitivity analysis.
This breakthrough significantly reduces the computational cost of gradient-based optimization in scientific computing, enabling faster neural differential equation training, more efficient uncertainty quantification, and scalable sensitivity studies. While primarily designed for low-dimensional problems, future work aims to extend its applicability to high-dimensional systems and more complex dynamics, promising broad impact in computational science and machine learning.
Deep Analysis
Background
微分方程在科学、工程中的应用不断深化，数值求解技术不断演进。早期以Runge-Kutta和多步法为代表的求解器已成熟，但在微分敏感性分析和参数优化中，反向微分的效率成为瓶颈。GPU加速求解器如DiffEqGPU和MPGOS在速度上表现优异，但难以实现高效反向微分。支持反向微分的工具如Diffrax和torchdiffeq虽能提供梯度，但在GPU集成和速度方面存在不足。两者结合尚未突破，限制了大规模科学计算的发展。
Core Problem
GPU上高效求解ODE与反向微分的矛盾源于自适应步长和反向微分的本质冲突。自适应求解器在每步动态调整步长，导致反向微分时难以追踪中间状态，存储和重算成本高。传统方法在速度和精度间难以兼顾，尤其在大规模集成和多次微分任务中，性能瓶颈明显。如何在保证速度的同时，实现高精度反向梯度，是当前研究难点。
Innovation
GRADSOLVE的创新在于引入固定步长回放机制，将自适应求解的步骤记录后，用固定步长重放，避免了存储和重算的开销。融合GPU内核技术，将整个求解过程编译成单一GPU核，显著提升梯度计算速度。支持多积分器，兼容多GPU架构，提供高速、精确的反向微分方案，突破了GPU反向微分的性能瓶颈。
Methodology
- �� 记录：运行一次自适应求解器，记录接受的每个步骤的步长，形成固定的时间网格。
- �� 回放：利用记录的步长，在相同网格上用固定步长重放求解，避免误差估计和动态调整。
- �� GPU融合：将回放过程编译成单一GPU内核，避免多次内核调用和存储中间状态。
- �� 反向微分：对固定步长回放进行反向微分，获得精确的离散伴随梯度。
- �� 支持多积分器：包括显式和Rosenbrock方法，适应不同系统刚性需求。
Experiments
采用Lorenz系统、刚性Robertson系统等为测试对象，比较GRADSOLVE与DiffEqGPU、Diffrax、torchode在正向速度和梯度计算速度上的表现。调节容差、集成规模，评估在不同GPU（A100、H100、RTX4090）上的性能。指标包括速度比、梯度误差和资源消耗，验证方法在不同场景的适用性。
Results
在三代GPU上，正向核速度比DiffEqGPU快2.8倍，反向梯度比Diffrax checkpointed快5.6-14.1倍，误差匹配下，非刚性小规模系统优势明显。大规模集成中优势缩小，但整体仍优。Lorenz系统梯度计算快约8倍，极大降低反向微分成本，适合多次微分任务。
Plain Language Accessible to non-experts
想象你在厨房准备多道菜，每道菜需要不同时间和步骤。传统方法就像每次都从头开始做，既慢又麻烦。而GRADSOLVE像是提前记住每道菜的步骤，然后用一台高速厨具快速重现，不仅节省时间，还能保证每次都做得一样好。它在厨房里用特殊的工具，把所有步骤合成一条高速生产线，让厨师（科学家）可以在短时间内做出更多菜（模型的梯度），帮助他们更快找到最好的配方（参数设置）。
ELI14 Explained like you're 14
想象你在学校做科学实验，老师让你测量化学反应的速度。有一种方法是每次都一步步观察，既慢又麻烦；另一种是提前记住每个步骤，然后用一台神奇的机器快速重复。GRADSOLVE就像是把每个实验步骤都记下来，然后用这台机器在一瞬间完成所有步骤，不仅快，还很准。这样，科学家可以用更少的时间，得到更多不同条件下的反应速度，帮他们更快理解这个反应的秘密。它用在超级快的电脑上，让复杂的模型变得更容易计算，也让科学研究变得更高效。
Abstract
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.