Adaptive Probabilistic ODE Solvers Without Adaptive Memory Requirements

TL;DR

This paper introduces an adaptive probabilistic ODE solver that maintains fixed memory usage, enabling efficient, long-term simulations without memory overload, and integrates seamlessly with modern scientific computing frameworks.

Contribution

We develop a novel adaptive probabilistic ODE solver with fixed memory requirements, overcoming limitations of traditional adaptive methods and enhancing computational efficiency.

Findings

Eliminates memory issues for long time series

Accelerates simulations via just-in-time compilation

Ensures compatibility with JAX-based scientific computing

Abstract

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.