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.

翻译：刚性常微分方程组和稀疏训练数据在科学问题中普遍存在。本文描述了高效、隐式、矢量化的方法，用于对刚性常微分方程组进行时间积分，并通过伴随方法计算参数梯度。主要创新在于将问题同时在独立时间序列数量和批量（或“块”）连续时间步上矢量化，从而有效实现了隐式常微分方程组组装的矢量化。针对后向欧拉法的线性化隐式系统的块双对角结构，允许利用并行循环归约（PCR）进一步矢量化。在输入数据的两个轴向上同时矢量化，为计算设备提供了更高的计算带宽，使得即使是数据相对稀疏的问题也能充分利用现代GPU，相较于标准的序贯时间积分实现超过100倍的加速。我们通过多个示例问题（包括解析刚性/非刚性ODE模型及神经ODE模型）展示了隐式矢量化时间积分的优势。此外，本文还描述并提供了所开发方法的开源免费实现。