Time-vectorized numerical integration for systems of ODEs

TL;DR

This paper introduces a novel vectorized implicit method for efficiently integrating stiff ODE systems and computing gradients, significantly accelerating computations on modern GPUs, especially with sparse data.

Contribution

The main contribution is the development of a fully vectorized implicit integration approach for stiff ODEs, leveraging parallel cyclic reduction and enabling over 100x speedups.

Findings

Achieved over 100x speedup compared to sequential methods.

Demonstrated effectiveness on stiff and non-stiff ODE models.

Provided open-source implementation for broader use.

Abstract

Stiff systems of ordinary differential equations (ODEs) and sparse training data are common in scientific problems. This paper describes efficient, implicit, vectorized methods for integrating stiff systems of ordinary differential equations through time and calculating parameter gradients with the adjoint method. The main innovation is to vectorize the problem both over the number of independent times series and over a batch or "chunk" of sequential time steps, effectively vectorizing the assembly of the implicit system of ODEs. The block-bidiagonal structure of the linearized implicit system for the backward Euler method allows for further vectorization using parallel cyclic reduction (PCR). Vectorizing over both axes of the input data provides a higher bandwidth of calculations to the computing device, allowing even problems with comparatively sparse data to fully utilize modern GPUs and achieving speed ups of greater than 100x, compared to standard, sequential time integration. We demonstrate the advantages of implicit, vectorized time integration with several example problems, drawn from both analytical stiff and non-stiff ODE models as well as neural ODE models. We also describe and provide a freely available open-source implementation of the methods developed here.