Search papers, labs, and topics across Lattice.
GRADSOLVE is an open-source JAX library designed to efficiently solve and reverse-mode differentiate low-dimensional ODE ensembles on NVIDIA GPUs. By recording the steps of an adaptive solver and differentiating a fixed-step replay, it achieves significant speed improvements, outperforming existing methods like Diffrax by up to 14.1 times for gradient computations while maintaining accuracy. This advancement addresses the long-standing trade-off between speed and differentiability in GPU-based ODE solvers, making it a valuable tool for applications requiring rapid derivative calculations.
GRADSOLVE achieves up to 14.1x faster gradient computations for ODE ensembles, revolutionizing the speed-accuracy trade-off in GPU-based solvers.
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.
