Time-vectorized numerical integration for systems of ODEs
This addresses computational bottlenecks in scientific problems involving stiff ODEs, such as neural ODE models, by enabling better GPU utilization, though it is incremental in optimizing existing methods.
The paper tackles the problem of efficiently integrating stiff systems of ordinary differential equations (ODEs) with sparse training data by developing implicit, vectorized methods that vectorize over independent time series and sequential time steps, achieving speedups of over 100x compared to standard sequential integration.
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.