GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

Alessio Spurio Mancini

Source-linked AI summary

GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

Alessio Spurio Mancini

arXiv:2609.02876v1Sep 2, 2026cs.MScs.DCcs.LGmath.NA

TL;DR

GPU ODE ensembles must balance fast solving against efficient reverse-mode gradients. GRADSOLVE records accepted adaptive steps and differentiates a fixed-step replay, achieving substantial speedups over Diffrax while retaining the same discrete-adjoint derivative. Its scope is low-dimensional ensembles, with performance depending on ensemble size, stiffness, accuracy, and fused-kernel limits.

Problem
GPU software has lacked a tool combining fused-kernel ensemble-solve speed with efficient reverse-mode gradients.
Method
GRADSOLVE records accepted adaptive steps and differentiates a fixed-step replay over the ensemble, holding recorded step sizes fixed.
Results
GRADSOLVE ran 2.8x faster than DiffEqGPU.jl forward-only and computed gradients 5.6–14.1x faster than Diffrax’s checkpointed adjoint once a record existed.
Takeaways & Limitations
GRADSOLVE is suited to low-dimensional GPU ensembles differentiated many times against one recorded mesh, while retaining Diffrax as a routing fallback.
Takeaways & Limitations
The advantage narrows with large ensembles and can reach parity on stiff systems at tight accuracy; fused kernels also have a state-dimension launch limit.

Abstract

from arXiv · show

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.

1 Introduction

GRADSOLVE addresses the GPU trade-off between fast ensemble solving and efficient reverse-mode differentiation by recording adaptive steps and differentiating a fixed-step replay. It matches the discrete adjoint used by Diffrax while targeting repeated gradients over one recorded mesh.

1 Introduction: GRADSOLVE is an open-source JAX library for reverse-mode differentiating low-dimensional ODE ensembles on NVIDIA GPUs.It records accepted adaptive steps and replays them as a fixed-step computation.
1 Introduction: 2.8x faster forward-only execution than DiffEqGPU.jl was measured for GRADSOLVE’s solver kernel.The forward-only path avoids differentiation overhead when no gradient is required.
1 Introduction: The returned gradient is the exact discrete adjoint of the recorded replay, matching the derivative Diffrax returns by default.The fixed-step chain avoids differentiating an adaptive loop with checkpointing and branch handling.
1 Introduction: GRADSOLVE targets ensemble calibrations, sensitivity studies, and small neural right-hand sides differentiated many times against one recorded set of steps.Diffrax remains available as a fallback.
1 Introduction: 5.6–14.1x faster gradients than Diffrax’s checkpointed adjoint were measured once a record existed, at matched forward-state accuracy.The comparison covered three GPU generations.

2 Differentiating an ODE solve

Differentiating ODE ensembles requires gradients through numerical solves whose adaptive trajectories, rejected trials, and GPU execution designs create distinct computational costs. Reverse mode uses one backward sweep for the full gradient, but adaptive branching and GPU scheduling complicate that sweep.

2 Differentiating an ODE solve: An ensemble contains n independent trajectories of the same ODE, solved together from potentially different initial conditions and with shared or varying parameters.The target is the gradient of a scalar loss formed from the ensemble solutions.
2 Differentiating an ODE solve: Adaptive Runge–Kutta solvers accept or reject trial steps using an error estimate and adjust the next step size accordingly.Embedded pairs obtain the error estimate from two weighted combinations of the same stage evaluations.
2 Differentiating an ODE solve: Different trajectories can accept different numbers of steps, while rejected trials, error estimates, and controller updates also contribute to the computation.This runtime variability complicates ensemble execution and differentiation.
2 Differentiating an ODE solve: Reverse mode yields the whole gradient in one backward sweep, whereas forward mode generally needs one pass per parameter direction.Reverse mode requires intermediate forward values to be stored or recomputed.
2 Differentiating an ODE solve: Adaptive accept/reject branching makes the solve discontinuous at threshold crossings, so differentiable solvers hold branch decisions and accepted step sizes fixed.Different fixed quantities can produce measurably different derivatives.
2 Differentiating an ODE solve: Fused kernels maximize forward speed but lack reverse-mode differentiation, while lockstep arrays and per-operation dispatch incur waiting or launch overheads.These designs expose the trade-off between forward throughput and reverse-mode support.

3 The gradsolve method: record and replay

GRADSOLVE separates adaptive integration from differentiation: it records accepted step sizes, then differentiates a fixed-step replay of the same trajectory. This regular chain supports GPU execution while preserving the discrete-adjoint convention and routing among general, fused, and stiff implementations.

3.1 The construction: GRADSOLVE records each trajectory’s accepted adaptive step sizes, which define a mesh for a second, fixed-step replay.Rejected trials, error estimates, and controller decisions are omitted from the differentiated replay.
3.1 The construction: Across ensembles, variable-length records are zero-padded into one rectangular array so the trajectories replay as a single regular GPU computation.Zero-size padded steps leave states and derivatives unchanged.
3.1 The construction: The replay is a chain of S one-step maps using recorded h_i values, reproducing the accepted trajectory without branching.For non-stiff systems it uses the explicit method map; stiff systems use a different one-step map.
3.1 The construction: GRADSOLVE returns the exact reverse-mode derivative of the replay with the recorded mesh held fixed, matching the discrete-adjoint convention described for Diffrax.Stopping the mesh dependence omits the term describing how adaptive step choices would change with parameters.
3.2 The integrators: Non-stiff problems use Tsit5 by default or Verner’s seventh-order pair, while stiff problems use Rodas5P or a fused-kernel variant.Stiff replay differentiates through linear solves automatically; the backward pass requires a solve with the transposed system matrix.
3.3 The engines and the router: The router selects engines using stiffness, state dimension, and gradient demand; fused engines are limited to d ≤ 64, while the general path accepts every supported right-hand side.The general recorder is a batched JAX loop, whereas the fused recorder uses a Warp kernel for registered fields.
3.3 The engines and the router: With either recorder, recording costs about one forward ensemble solve, after which the fixed replay can be differentiated repeatedly against the recorded mesh.Forward-only requests can use fused kernels without paying for differentiability.

4 Benchmarks

Across matched-accuracy benchmarks, GRADSOLVE improves forward-only throughput and reverse-mode gradient cost, with the largest gains on non-stiff systems and smaller or conditional gains on stiff systems.

4.1 Forward-only throughput: 2.8× faster in double precision than DiffEqGPU.jl, GRADSOLVE’s forward-only kernel provides the fastest measured integrator on the Lorenz ensemble.At matched tolerances and n = 1.05 × 10^6 trajectories, the single-precision speedup was 1.95×.
4.1 Forward-only throughput: 1.4–2.0× faster throughout, GRADSOLVE’s fused kernel maintains its advantage when Lorenz trajectories have widely varying difficulty.Both fused kernels preserve per-trajectory speed rather than slowing to their hardest trajectory.
4.2 Forward-only stiff kernels: 1.5–4.5× slower than DiffEqGPU.jl’s hand-tuned stiff kernels, the automatically generated stiff forward kernel pays a generality cost.A hand-written kernel using the same second-order method was 1.5–3.1× faster than DiffEqGPU.jl on HIRES, but up to 1.1–5.3× slower on Robertson at loose accuracy.
4.3 Reverse-mode gradient cost: 3.2–3.5× faster than Diffrax’s forward mode on Lorenz, GRADSOLVE’s reverse-mode replay remains faster, while its own forward mode reaches 7.3–7.9×.Diffrax’s forward mode was itself 2.6–2.7× faster than its checkpointed adjoint for this problem.

5 Where the speedup comes from and when to expect it

The measured gradient advantage comes primarily from replacing an adaptive, checkpointed differentiated loop with a fixed-length replay, while uneven step counts and recording contribute little. The benefit is strongest for small systems and repeated gradients, but narrows for large ensembles and stiff systems at tight accuracy.

Source of the speedup: 1.33–1.34× step-count imbalance could explain only a small fraction of the measured 7.8–9.3× gradient advantage.The slowest Lorenz trajectory accepted only 1.33–1.34 times the average number of steps.
Source of the speedup: 1.9–2.6× and 2.2–2.4× CPU replay speedups with general and fused recording showed that the recorder was not the main source of the gap.Both approaches differentiated the same fixed-step replay, and their curves tracked closely.
Source of the speedup: A fixed mesh alone did not explain the advantage: it was slower than matched adaptive Diffrax on Lorenz and Robertson, while HIRES gained speed at much worse error.The HIRES fixed-mesh run reached 5.9 × 10−4 error versus 1.3 × 10−7 for the replay.
Mechanism: The GPU executes a fixed-length scan as one regular program, whereas adaptive differentiation retains bounded-loop and checkpoint bookkeeping.The CPU showed no corresponding scan advantage, supporting a GPU-execution explanation.
Mechanism: 11–15% of one record-plus-gradient cost was recording, while about 80% was the backward sweep that replay simplifies.The in-kernel adjoint adds memory savings rather than speed in these ranges.
Source of the speedup: 7.9–11.1× faster gradients came from the fixed-step scan versus Diffrax’s StepTo loop on the A100, accounting for essentially the whole advantage.The control held the method and mesh fixed, isolating the differentiated computation.
When to expect it: The advantage is largest for state dimensions up to eight, smaller ensembles, and many gradients against one recorded mesh; it declines with large ensembles and can reach parity on stiff systems at tight accuracy.For ensembles from n = 4096 to n = 1 048 576, the measured range was 6.2–8.5×; stiff-system margins ranged from 5.6× to parity or slightly below.

6 Discussion and conclusion

GRADSOLVE combines adaptive recording with fixed-step replay to provide reverse-mode discrete adjoints for low-dimensional ODE ensembles on GPUs. Its specialized kernels target repeated differentiation against one mesh, while Diffrax and a general path cover cases beyond fused-kernel boundaries.

Contribution: GRADSOLVE records accepted adaptive steps, then differentiates their fixed-step replay to return a discrete adjoint of the same kind Diffrax produces.The replay holds the recorded steps fixed as data, making the gradient the exact derivative of that replay.
Scope: GRADSOLVE specializes in low-dimensional GPU ensembles differentiated many times against a recorded mesh, while retaining Diffrax as a routing fallback.The general path supports user-defined right-hand sides, and specialized engines apply only within stated implementation boundaries.
Performance: 9.5× and 7.2× per gradient against Diffrax were reported for the Verner and Rodas5P integrators, respectively.Both methods use the same replay-based path and show comparable advantage in the cited comparison.
Performance: The advantage is largest on non-stiff systems, narrows as ensembles grow, and reaches parity with Diffrax’s forward mode on a stiff system at tightest accuracy.These boundaries qualify the performance benefit rather than the correctness of the general path.
Release: GRADSOLVE is released as open-source Python software built on JAX, with fused Warp kernels and Diffrax available as baseline and fallback.The repository is hosted at https://github.com/ECLIPSE-AI4Science/gradsolve.

Appendix A Benchmark problem definitions

The benchmark suite covers explicit, mildly stiff, stiff, chaotic, orbital, neural, and controlled ODE ensembles. Problems vary state dimension, initial conditions, parameters, stiffness, and trajectory step counts to test the solver across distinct regimes.

Ensemble construction: Parameter and initial-condition variation is used to create ensembles with shared or differing trajectories, including Lorenz ρ ∈ [0, 21] and galactic-orbit initial-state ranges.The Lorenz timing range lies below the chaotic-attractor threshold, while chaotic behavior enters through other studies and a ρ = 28 timing point.
Problem suite: The benchmark systems include Lorenz, Lorenz-96, Van der Pol, Robertson, HIRES, galactic orbits, a neural Robertson variant, and a linear control.Lorenz-96 is used only for a high-dimensional gradient check, while the linear family provides an equal-step control and dimension-limit probe.
Stiff systems: Robertson uses rates k1 = 0.04, k2 = 3 × 10^7, and k3 = 10^4 over t ∈ [0, 10^4], spanning nearly nine orders of magnitude.The Rodas5P and fitting studies vary rates or draw true parameters around these values.
Learned dynamics: The neural Robertson system replaces the autocatalytic term with a non-negative neural network whose weights θ receive gradients through the stiff solve.The network has two width-16 hidden layers, tanh activations, and a softplus output.
Dimensions: The ensembles span state dimensions from low-dimensional chemical and dynamical systems to d = 96 for the Lorenz-96 gradient check.The linear family separately probes the fused kernels’ state-dimension limit.

Appendix B Benchmark protocol and reproducibility

The appendix specifies the measurement, baseline, timing, fitting, hardware, software, and gradient-check procedures needed to reproduce the reported evaluations.

Reproducibility: The protocol defines accuracy matching, baseline configurations, timing rules and uncertainties, fitting studies, hardware and software, and gradient correctness checks.These procedures are organized across Appendices B.1–B.6.

Appendix B.1

Solver comparisons are made at matched achieved forward accuracy rather than nominal tolerance, using componentwise relative error against a high-accuracy reference.

Accuracy matching: Methods are compared at matched achieved accuracy because different-order solvers reach different errors at the same nominal tolerance.GRADSOLVE’s median componentwise relative error is measured against an analytic or high-accuracy numerical reference.

Baseline configuration

The baseline comparisons use synchronized GPU measurements, specified Diffrax configurations, reverse-mode timing, and several forward-mode and replay controls.

Every device is synchronized before reading timestamps in both JAX and PyTorch.
Diffrax uses Tsit5 for non-stiff systems, Kvaerno5 for stiff systems, adaptive PID control, and its default RecursiveCheckpointAdjoint.
Checkpoint counts are selected as the fastest of three measured settings based on attempted-step counts and the library default.
All baselines are timed in reverse mode, while ForwardMode is additionally measured for systems with one to three inputs per trajectory.
Forward-mode gradient assembly uses one pass per input, with either whole-ensemble passes or per-trajectory jacfwd inside vmap; the faster placement is quoted.
The replay is also differentiated in forward mode, while Diffrax’s approximate continuous adjoint is excluded from timing.
PyTorch comparisons use a separate PCIe A100 and therefore compare across machines; torchode uses torch.compile, whereas torchdiffeq uses its eager continuous adjoint.

Appendix B.3

Timing procedures are designed to reduce transient effects while separating warm-up and compilation costs from repeated computation.

Per-gradient times use the best of three timed calls, or five for higher-order engines, after warm-up and one-time replay recording.Fitting-loop times are averaged over repetitions.

Fitting protocol

The fitting studies compare concurrent parameter-recovery runs under identical optimizers while reusing recorded meshes and excluding failed engine–problem combinations.

Fitting studies run B concurrent fits, with B ∈ {1, 8, 64, 256}, and each fit recovers its own parameter under the identical optimizer.The neural-network study uses B as minibatch size rather than concurrent-fit count.
Recorded meshes are treated as data throughout each fit, and reported fits do not re-record midway.When fixed-cadence re-recording is used, its time is included inside the loop.
Failed engine–problem combinations are dropped, while successful fits use Adam with problem-specific learning rates and initial parameter offsets.

Hardware and software

Measurements span NVIDIA A100, H100, and RTX 4090 GPUs, with A100 results collected on two machine configurations and software environments documented separately.

Most experiments use double precision on one NVIDIA A100, H100, or RTX 4090 GPU.The RTX 4090 has 24 GB and a double-precision arithmetic rate one sixty-fourth of its single-precision rate.
A100 measurements use both PCIe and SXM4 packagings across two machines, whose differences are named where relevant.
The CPU experiment runs on an Apple M3 Max, while Figure 6(b)’s CPU curve uses the CPU hosting the A100.
Table 6 documents the measurement software environment and distinguishes the Julia forward comparison from its reverse-mode attempt.

Gradient checks

The gradient checks distinguish replay differentiation correctness from sensitivity accuracy against the underlying ODE and verify agreement with an independently stepped implementation. Across these tests, GRADSOLVE’s gradients closely match finite-difference, closed-form, and shared-mesh references, while forward-state error remains representative across Robertson species.

Replay differentiation checks: Finite-difference checks agree with GRADSOLVE to about 9 × 10−11 on a laptop CPU and 7 × 10−9 for stiff Robertson replay on an A100.For a linear d = 32 case where finite differences reach the roundoff floor, the replay agrees with closed-form sensitivity to 1.6 × 10−7 relative.
Sensitivity accuracy: Independent high-accuracy references show relative gradient differences of 4.6 × 10−6 on Lorenz and 6.8 × 10−6 on Robertson.These differences are at the systems’ forward discretization-error levels rather than machine precision, supporting sensitivity accuracy beyond self-consistency.
Cross-engine agreement: Shared-mesh gradients from Diffrax agree with GRADSOLVE to 2.9 × 10−15 on Lorenz and 1.8 × 10−13 on Van der Pol.Both engines use the same recorded times, so the comparison isolates agreement between their differentiated fixed-step computations.
Implementation cross-checks: At Lorenz-96 state dimension 96, independent JAX replay and Warp backward sweeps agree to about 3.5 × 10−11.The stiff counterpart was tested only through d ≈12, and parts of the in-kernel adjoint are handwritten because in-place linear solves impede Warp autodiff.