⚛️ general relativity

Differentiable astrophysics at scale: solving and differentiating ODE ensembles on the GPU

The paper introduces GRADSOLVE, a JAX library that leverages GPU parallelism to solve and differentiate ensembles of ordinary differential equations for astronomical models, achieving speedups of up to 1,500 times over serial CPU implementations and significantly outperforming existing JAX ODE solvers to enable routine gradient-based analyses of millions of trajectories.

Original authors: A. Spurio Mancini

Published 2026-09-24

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: A. Spurio Mancini

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern astronomy has entered an era where scientists do not just look at the sky; they simulate it with incredible precision to understand how the universe works. To test their ideas, researchers build complex mathematical models that describe everything from the swirling paths of stars in our galaxy to the expansion of space itself. These models often rely on solving a specific type of mathematical puzzle called an ordinary differential equation. Think of these equations as a set of instructions that tell a computer how a system changes over time, such as how a star moves under the pull of gravity or how the universe stretches as it ages. To find the best possible model, astronomers must run these instructions millions of times, tweaking the numbers slightly each time to see which version matches the real world most closely. This process is the backbone of modern discovery, but it has hit a wall: the sheer number of calculations required is so vast that even the fastest supercomputers struggle to finish the job in a reasonable amount of time.

A researcher has now broken through this bottleneck by moving these calculations from standard computer processors to graphics processing units, the powerful chips originally designed to render video games. They developed a new software tool called gradsolve that treats each of the millions of simulations as a separate worker, assigning each one its own dedicated lane on the computer chip. In traditional methods, a computer tries to solve these problems one after another or in rigid groups, often forcing fast simulations to wait idly for the slowest ones to catch up. This new approach allows every single simulation to run at its own speed, finishing exactly when it is ready without holding anyone else back. The result is a dramatic leap in speed that transforms what used to be a task taking days into one that can be completed in seconds.

The researcher tested this method on three very different problems from across the field of astronomy to prove its versatility. First, they looked at the paths of stars being stripped away from a globular cluster to form a long, thin stream across the Milky Way. By simulating the motion of one million stars at once, their new tool ran 1500 times faster than the best existing software running on a standard computer with 128 processors. Second, they applied it to the history of the universe's expansion, specifically testing theories about dark energy, the mysterious force driving the cosmos apart. Here, the tool calculated distances for a million different versions of the universe, again outperforming the old methods by a factor of 1500 compared to a single computer core. Finally, they examined the twisting motion, or precession, of two black holes spiraling toward each other before colliding. For a population of one million binary black holes, the new software was 760 times faster than the standard code used by astronomers today.

What makes this achievement particularly significant is that the new tool does not just run faster; it also calculates the necessary derivatives, or rates of change, in the same pass. In the past, finding these rates required a separate, equally expensive calculation that would have doubled the time needed. By handling the simulation and its sensitivity to changes simultaneously, the researcher has removed a major barrier to entry for complex statistical analysis. The speed gains are so substantial that analyses which were previously too computationally expensive to attempt are now becoming routine. For instance, a task that would take a single computer core nearly a year to complete can now be finished on a single graphics card in a matter of seconds.

The researcher was careful to ensure that this massive increase in speed did not come at the cost of accuracy. They set strict standards for how close the results needed to be to a known, highly precise reference solution. In every test, the new software met these rigorous accuracy requirements, proving that the speed came from a smarter way of organizing the work rather than from cutting corners. The study explicitly shows that for problems involving millions of independent calculations that do not require complex, stiff mathematical steps, this new approach is superior to the current state-of-the-art methods. It demonstrates that by rethinking how computers handle massive groups of simulations, astronomers can now explore the universe with a level of detail and speed that was previously out of reach, opening the door to more precise models of everything from the structure of our galaxy to the fundamental nature of dark energy.

Technical Summary: Differentiable Astrophysics at Scale

Problem Statement
Modern astrophysical inference increasingly relies on gradient-based methods, such as Hamiltonian Monte Carlo (HMC) and Fisher forecasts, which require the derivatives of model predictions with respect to their parameters. A common computational bottleneck in these analyses is the need to solve ensembles of ordinary differential equations (ODEs)—ranging from thousands to millions of trajectories—where each trajectory corresponds to a different set of parameters or initial conditions.

Current approaches face two primary limitations:

Serial CPU execution: Solving these ensembles on Central Processing Units (CPUs) is computationally expensive, often limiting the scale of analyses.
Vectorized GPU limitations: State-of-the-art differentiable ODE libraries in JAX, such as diffrax, typically vectorize the ensemble as a single array. In adaptive solvers, where step sizes vary based on local error estimates, this "lock-step" execution forces all trajectories to wait for the slowest member of the batch to finish a step. Consequently, the computational cost is dictated by the trajectory requiring the most steps rather than the mean, leading to significant inefficiencies on Graphics Processing Units (GPUs).

Methodology
The author introduces gradsolve, a JAX library designed to move the integration of ODE ensembles to GPUs with a "one trajectory per thread" architecture.

Per-Thread Integration: Unlike vectorized approaches, gradsolve generates a GPU kernel where each thread integrates a single trajectory independently. Each thread maintains its own state, intermediate evaluations, and sensitivities in private memory. This allows trajectories to proceed with their own adaptive step sizes without waiting for others, confining synchronization only to groups of 32 threads (warps) on NVIDIA GPUs.
Forward Sensitivities: The library integrates the forward sensitivity equations (derivatives of the state with respect to parameters) alongside the ODEs within the same kernel. This forward-mode differentiation is efficient for problems with few parameters and many outputs, avoiding the memory overhead of reverse-mode (adjoint) differentiation which requires storing or recomputing the forward trajectory.
Dense Output: The solver provides dense output, allowing the solution to be interpolated at arbitrary observation times without forcing the solver to terminate steps at those specific times.
Comparison Protocol: To ensure fair benchmarking, the author compares gradsolve against diffrax and serial CPU codes (e.g., gala, scipy, precession) under strict, pre-defined accuracy requirements. Codes are accepted only if the 99th percentile of their error across the ensemble meets a physical accuracy threshold (e.g., relative error in orbital position or distance).

Key Contributions

Library Development: The release of gradsolve, a JAX-native library that generates per-thread GPU kernels for adaptive ODE integration and forward sensitivity calculation.
Architectural Shift: Demonstrating that for non-stiff, small-system ODE ensembles, a per-thread approach significantly outperforms vectorized batch approaches on GPUs by eliminating idle time caused by step-size heterogeneity.
Comprehensive Benchmarking: Providing rigorous performance comparisons across three distinct astrophysical domains:
- Stellar Streams: Integrating orbits in a flattened Galactic potential and their sensitivities to potential parameters.
- Dark Energy: Calculating expansion histories and distances for thawing quintessence models across many cosmologies, including the solution for spatial flatness constraints.
- Binary Black Holes: Simulating spin precession across a population of binary black holes.

Results
The paper reports speed-up factors for ensembles of $10^6$ trajectories, comparing gradsolve against both serial CPU codes (running on all 128 cores of an AMD EPYC 9745) and the fastest accepted diffrax configuration on the same NVIDIA H200 GPU.

vs. Serial CPU: gradsolve is 8.5 to 1,500 times faster than the serial CPU code running on all 128 cores. Against a single CPU core, the speed-up is several orders of magnitude.
- Stellar Orbits: 8.5 $\times$ faster than 128 cores.
- Dark Energy: 1,500 $\times$ faster than 128 cores.
- Spin Precession: 760 $\times$ faster than 128 cores.
vs. diffrax (Same GPU): gradsolve is 11 to 15 times faster than diffrax for large ensembles ( $10^6$ $1 0^{6}$ members) while maintaining identical accuracy.
- Stellar Orbits: 15 $\times$ faster (orbits + derivatives).
- Dark Energy: 11 $\times$ faster (including the solve for spatial flatness).
- Spin Precession: 11 $\times$ faster.
Accuracy: All tested codes met the pre-defined accuracy requirements. The 99th percentile of errors for gradsolve remained within the physical limits set for each application (e.g., $10^{-4}$ for orbital positions, $10^{-7}$ rad for precession angles).
Source of Speed-up: When comparing gradsolve and diffrax using the same numerical method (e.g., Verner 7 or Tsit5) and tolerance, the number of step attempts is nearly identical. The performance gap is attributed to the implementation: gradsolve avoids the memory bandwidth overhead of reading/writing the entire batch state to global memory at every step and eliminates the "waiting" penalty of the lock-step vectorized loop.

Significance and Scope
The paper claims that gradsolve makes gradient-based analyses requiring ensembles of millions of ODE integrations "routine" by reducing computation times from hours to seconds on a single GPU.

Applicability: The approach is optimized for non-stiff systems with small state vectors (3–9 components) and few differentiated parameters (forward-mode efficiency).
Limitations: The paper notes that for problems with large state vectors (e.g., grid hydrodynamics or particle-mesh simulations) or stiff systems (e.g., Einstein-Boltzmann equations), other approaches (batched array solvers or implicit/Rosenbrock methods) are more appropriate. While gradsolve includes a Rosenbrock kernel for stiff problems, the primary benchmarks focus on non-stiff astrophysical applications.
Impact on Inference: In the dark energy profile likelihood example, the speed-up reduces the cost of the ODE solve to a fraction of the total iteration cost, shifting the computational bottleneck to the likelihood calculation itself. This enables more complex models and larger parameter spaces to be explored within feasible timeframes.

The code is publicly available, allowing the astronomical community to integrate these high-performance differentiable solvers into existing JAX-based inference pipelines.

Technical Summary: Differentiable Astrophysics at Scale

More like this