Adaptive Probabilistic ODE Solvers Without Adaptive Memory Requirements
This solves a critical memory bottleneck for researchers using probabilistic ODE solvers in scientific computing, enabling reliable long time series simulations.
The paper tackles the problem of probabilistic ODE solvers with adaptive step sizes having unpredictable memory demands that cause simulations to fail, and develops an adaptive probabilistic solver with fixed memory requirements that eliminates memory issues, accelerates simulations by orders of magnitude, and enables compatibility with JAX.
Despite substantial progress in recent years, probabilistic solvers with adaptive step sizes can still not solve memory-demanding differential equations -- unless we care only about a single point in time (which is far too restrictive; we want the whole time series). Counterintuitively, the culprit is the adaptivity itself: Its unpredictable memory demands easily exceed our machine's capabilities, making our simulations fail unexpectedly and without warning. Still, dropping adaptivity would abandon years of progress, which can't be the answer. In this work, we solve this conundrum. We develop an adaptive probabilistic solver with fixed memory demands building on recent developments in robust state estimation. Switching to our method (i) eliminates memory issues for long time series, (ii) accelerates simulations by orders of magnitude through unlocking just-in-time compilation, and (iii) makes adaptive probabilistic solvers compatible with scientific computing in JAX.