Parallel-in-Time Probabilistic Numerical ODE Solvers

TL;DR

This paper builds on the time-parallel formulation of iterated extended Kalman smoothers to formulate a parallel-in-time probabilistic numerical ODE solver and reduces the span cost from linear to logarithmic in the number of time steps.

Abstract

Probabilistic numerical solvers for ordinary differential equations (ODEs) treat the numerical simulation of dynamical systems as problems of Bayesian state estimation. Aside from producing posterior distributions over ODE solutions and thereby quantifying the numerical approximation error of the method itself, one less-often noted advantage of this formalism is the algorithmic flexibility gained by formulating numerical simulation in the framework of Bayesian filtering and smoothing. In this paper, we leverage this flexibility and build on the time-parallel formulation of iterated extended Kalman smoothers to formulate a parallel-in-time probabilistic numerical ODE solver. Instead of simulating the dynamical system sequentially in time, as done by current probabilistic solvers, the proposed method processes all time steps in parallel and thereby reduces the span cost from linear to logarithmic in the number of time steps. We demonstrate the effectiveness of our approach on a variety of ODEs and compare it to a range of both classic and probabilistic numerical ODE solvers.

Figures (6)

• Figure 1: Trajectories, errors, and error estimates computed by the parallel-in-time solver. Top row: ODE solution trajectories. Visually, all three test problems seem to be solved accurately. Bottom row: Numerical errors (lines) and error estimates (shaded area). Ideally, for good calibration, the error should be of similar magnitude than the error estimate. The posterior appears underconfident on the logistic equation, and well-calibrated on the rigid body and Van der Pol problems.
• Figure 2: The sequential and parallel IEKS compute numerically identical solutions. For all three problems and all considered grid sizes, the sequential IEKS and the parallel ParaIEKS achieve (numerically) identical errors. This is expected, as both versions compute the same quantities and only differ in their implementation.
• Figure 3: Work-precision diagrams for the sequential and parallel IEKS-based ODE solver. Top row: Runtime in seconds per error (lower-left is better). Bottom row: Speed-up of the parallel ParaIEKS over the sequential IEKS (higher is better). Across all problems, grid sizes, and priors, ParaIEKS outperforms the sequential IEKS.
• Figure 4: The time-parallel ParaIEKS shows logarithmic scaling in the grid size and benefits from GPU improvements. In comparison, the sequential IEKS and the classic Dopri5 and Kvaerno5 solvers show the expected linear runtime complexity (left). The sequential methods also do not show relevant changes in runtime for GPUs with more CUDA cores, whereas ParaIEKS's runtime improves (right).
• Figure 5: Benchmarking ParaIEKS against other commonly used numerical ODE solvers. Top row: Work-precision diagrams showing runtimes per error for a range of different ODE solvers (lower-left is better). Bottom row: Errors per specified grid size (lower-left is better). Per grid size, the closely related EKS, IEKS, and ParaIEKS solvers often coincide; Kvaerno5 achieves the lowest error per step as it has the highest order. In terms of runtime, ParaIEKS outperforms all other methods on medium-to-high accuracy settings due to its logarithmic time complexity.

Theorems & Definitions (7)

• Remark 1: Alternative Gauss--Markov priors
• Remark 2: Information operators for other differential equation problems
• Remark 3: On Prefix-Sums
• Remark 4: Linearizing with approximate Jacobians (EK0 & DiagonalEK1)
• Remark 5: Statistical linear regression
• Remark 6: Other choices for the initial state trajectory
• Remark 7: Convergence of the time-parallel IEKS