Time-vectorized numerical integration for systems of ODEs
Mark C. Messner Tianchen Hu Tianju Chen
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.