Accelerating massive ensembles of ordinary differential equations
Abstract
Solving large ensembles of small, independent ordinary differential equations (ODEs) is a common task in computational science and engineering, arising for example in Bayesian parameter estimation, Monte Carlo uncertainty quantification, and the integration of uncoupled physical systems such as the linearised Einstein-Boltzmann equations of cosmology. Modern graphics processing units (GPUs) with thousands of cores are well suited to such parallel workloads, but realising this potential requires careful attention to GPU-specific architectural constraints. We present modax, a Python library of GPU-accelerated ODE solvers specifically optimised for solving large ensembles of low-dimensional, independent problems. Our solvers are compatible with the JAX ecosystem but are written in a low-level, thread-based programming paradigm better suited to ensembles of highly divergent trajectories than their pure-JAX counterparts. Benchmarks on the Lorenz, Robertson and van der Pol lattice systems compare modax against diffrax, DiffEqGPU$.$jl and torchdiffeq across ensemble size, dimensionality, and trajectory divergence. Our explicit Tsit5 solver achieves 30x speed-up over diffrax on Lorenz systems, and our implicit Rodas5P solver achieves over 700x speed-up on Robertson systems, demonstrating the superiority of linearly implicit Rosenbrock-Wanner methods on GPUs. Our implicit solver scales well to higher dimensions thanks to the use of sparsity-exploiting techniques and data structures. We conclude by applying modax to three examples from the field of cosmology - Bayesian parameter estimation from primordial abundances, Monte Carlo uncertainty quantification of the global 21cm signal, and the computation of uncoupled Fourier modes in the primordial power spectrum - and demonstrate significant efficiency improvements in our ability to model physical phenomena described by low-dimensional (<200D) ODEs.
Related
- A semi-analytical surrogate model for giant planet evolution: bypassing ordinary differential equation solvers with localised thermodynamics, softplus asymptotes, and B-spline photometry
- A Method Using Differential Equations for Characterizing Warped Product Lagrangian Submanifolds in Complex Space Forms
- CausticFlow: An Efficient Machine Learning Framework Combining Neural Differential Equations and Normalizing Flows for Binary Microlensing Parameter Inference
- CausticFlow: An Efficient Machine Learning Framework Combining Neural Differential Equations and Normalizing Flows for Binary Microlensing Parameter Inference
- Parameter estimation in differential equations: Mathematical foundations for satellite gravimetry, review and perspectives
BibTeX
@misc{berry2026accelerating,
title = {{Accelerating massive ensembles of ordinary differential equations}},
author = {Lawrence Berry and Will Handley and Oliver Hahn and Nils Schöneberg},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.35661},
eprint = {2609.35661},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.35661},
}
RIS
TY - UNPB
TI - Accelerating massive ensembles of ordinary differential equations
AU - Lawrence Berry
AU - Will Handley
AU - Oliver Hahn
AU - Nils Schöneberg
PY - 2026
DA - 2026/09/28
PB - arXiv
DO - 10.48550/arxiv.2609.35661
UR - https://arxiv.org/abs/2609.35661
AB - Solving large ensembles of small, independent ordinary differential equations (ODEs) is a common task in computational science and engineering, arising for example in Bayesian parameter estimation, Monte Carlo uncertainty quantification, and the integration of uncoupled physical systems such as the linearised Einstein-Boltzmann equations of cosmology. Modern graphics processing units (GPUs) with thousands of cores are well suited to such parallel workloads, but realising this potential requires careful attention to GPU-specific architectural constraints. We present modax, a Python library of GPU-accelerated ODE solvers specifically optimised for solving large ensembles of low-dimensional, independent problems. Our solvers are compatible with the JAX ecosystem but are written in a low-level, thread-based programming paradigm better suited to ensembles of highly divergent trajectories than their pure-JAX counterparts. Benchmarks on the Lorenz, Robertson and van der Pol lattice systems compare modax against diffrax, DiffEqGPU$.$jl and torchdiffeq across ensemble size, dimensionality, and trajectory divergence. Our explicit Tsit5 solver achieves 30x speed-up over diffrax on Lorenz systems, and our implicit Rodas5P solver achieves over 700x speed-up on Robertson systems, demonstrating the superiority of linearly implicit Rosenbrock-Wanner methods on GPUs. Our implicit solver scales well to higher dimensions thanks to the use of sparsity-exploiting techniques and data structures. We conclude by applying modax to three examples from the field of cosmology - Bayesian parameter estimation from primordial abundances, Monte Carlo uncertainty quantification of the global 21cm signal, and the computation of uncoupled Fourier modes in the primordial power spectrum - and demonstrate significant efficiency improvements in our ability to model physical phenomena described by low-dimensional (<200D) ODEs.
ER -