GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

• URL: http://arxiv.org/abs/2609.02876v1
• Date: Wed, 02 Sep 2026 17:56:10 GMT
• Title: GRADSOLVE: fast exact gradients for ODE ensembles on GPUs
• Abstract summary: We present GRADSOLVE, an open-source library for solving and reverse-mode differentiating low-dimensional ODE ensembles on NVIDIA GPUs.<n>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.<n>It targets ensembles differentiated many times against one recorded mesh, keeps Diffrax as a fallback, and supports explicit and Rosenbrock.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• GPU-Enabled Large-Scale Optimization Using Randomized Linear Algebra [13.9407640377773]
  This paper introduces rlaopt, a PyTorch-based package for large-scale optimization and scientific computing using randomized numerical linear algebra (RandNLA)<n>Rlaopt addresses this gap by providing GPU-enabled solvers for positive-definite linear systems and convex empirical risk minimization with constraints and regularizers.<n>Rlaopt automatically checks compatibility with the selected solver and performs the required problem decomposition.
  arXiv  Detail & Related papers  (2026-09-08T02:14:05Z)
• Differentiate the Solver, Not the Equation: Reverse-Sweep Adjoints for Block Implicit Simulation [28.75201832533547]
  Differentiable simulation is a key component in learning, control, and inverse problems.<n>We propose solver-level differentiation, which differentiates the executed solver itself.<n>This work highlights solver structure as a practical organizing principle for efficient differentiable simulation.
  arXiv  Detail & Related papers  (2026-08-09T08:03:02Z)
• torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch [0.2960141730774496]
  We present torchsla, an open-source PyTorch library that enables GPU-accelerated, scalable, and differentiable sparse linear algebra.<n>torchsla supports multiple backends (SciPy, cuDSS, PyTorch-native) and seamlessly integrates with PyTorch autograd for end-to-end differentiable simulations.
  arXiv  Detail & Related papers  (2026-01-20T14:06:01Z)
• DInf-Grid: A Neural Differential Equation Solver with Differentiable Feature Grids [73.28614344779076]
  We present a differentiable grid-based representation for efficiently solving differential equations (DEs)<n>Our results demonstrate a 5-20x speed-up over coordinate-based methods, solving differential equations in seconds or minutes while maintaining comparable accuracy and compactness.
  arXiv  Detail & Related papers  (2026-01-15T18:59:57Z)
• INC: An Indirect Neural Corrector for Auto-Regressive Hybrid PDE Solvers [61.84396402100827]
  We propose the Indirect Neural Corrector ($mathrmINC$), which integrates learned corrections into the governing equations.<n>$mathrmINC$ reduces the error amplification on the order of $t-1 + L$, where $t$ is the timestep and $L$ the Lipschitz constant.<n>We test $mathrmINC$ in extensive benchmarks, covering numerous differentiable solvers, neural backbones, and test cases ranging from a 1D chaotic system to 3D turbulence.
  arXiv  Detail & Related papers  (2025-11-16T20:14:28Z)
• Enabling Automatic Differentiation with Mollified Graph Neural Operators [73.52999622724101]
  We propose the mollified graph neural operator ($m$GNO), the first method to leverage automatic differentiation and compute exact gradients on arbitrary geometries.<n>For a PDE example on regular grids, $m$GNO paired with autograd reduced the L2 relative data error by 20x compared to finite differences.<n>It can also solve PDEs on unstructured point clouds seamlessly, using physics losses only, at resolutions vastly lower than those needed for finite differences to be accurate enough.
  arXiv  Detail & Related papers  (2025-04-11T06:16:30Z)
• 3DGS$^2$: Near Second-order Converging 3D Gaussian Splatting [26.94968605302451]
  3D Gaussian Splatting (3DGS) has emerged as a mainstream solution for novel view synthesis and 3D reconstruction.<n>This paper introduces a (near) second-order convergent training algorithm for 3DGS, leveraging its unique properties.
  arXiv  Detail & Related papers  (2025-01-22T22:28:11Z)
• 3DGS-LM: Faster Gaussian-Splatting Optimization with Levenberg-Marquardt [65.25603275491544]
  We present 3DGS-LM, a new method that accelerates the reconstruction of 3D Gaussian Splatting (3DGS)<n>Our method is 20% faster than the original 3DGS while obtaining the same reconstruction quality optimization.
  arXiv  Detail & Related papers  (2024-09-19T16:31:44Z)
• Adaptive Elastic Training for Sparse Deep Learning on Heterogeneous Multi-GPU Servers [65.60007071024629]
  We show that Adaptive SGD outperforms four state-of-the-art solutions in time-to-accuracy. We show experimentally that Adaptive SGD outperforms four state-of-the-art solutions in time-to-accuracy.
  arXiv  Detail & Related papers  (2021-10-13T20:58:15Z)
• Nesterov Accelerated ADMM for Fast Diffeomorphic Image Registration [63.15453821022452]
  Recent developments in approaches based on deep learning have achieved sub-second runtimes for DiffIR. We propose a simple iterative scheme that functionally composes intermediate non-stationary velocity fields. We then propose a convex optimisation model that uses a regularisation term of arbitrary order to impose smoothness on these velocity fields.
  arXiv  Detail & Related papers  (2021-09-26T19:56:45Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.