Prob-GParareal: A Probabilistic Numerical Parallel-in-Time Solver for Differential Equations

Technical Summary: Prob-GParareal

Problem Statement
Parallel-in-Time (PinT) methods, particularly the Parareal algorithm, have emerged as a powerful approach to accelerate the solution of large-scale ordinary and partial differential equations (ODEs and PDEs) by enabling concurrent computations over the time domain. While extensions like GParareal have improved efficiency by using Gaussian Processes (GPs) to learn the discrepancy between coarse and fine solvers, existing PinT methods remain deterministic. They lack a principled mechanism to quantify the epistemic uncertainty arising from numerical approximations, finite time steps, and incomplete convergence. Furthermore, current probabilistic numerics approaches for ODEs often struggle with chaotic dynamics or require restrictive assumptions (e.g., linearization) that limit their applicability to complex, nonlinear systems. There is a critical gap in the literature regarding probabilistic counterparts to established PinT methods that can provide uncertainty quantification (UQ) without sacrificing scalability or compatibility with existing solvers.

Methodology
The paper introduces Prob-GParareal, a probabilistic extension of the GParareal algorithm. The core methodology involves modeling the Parareal correction function (the discrepancy between the expensive fine solver $F$ and the cheap coarse solver $G$) not just as a point estimate, but as a distribution using Gaussian Processes.

1. Probabilistic Formulation: Unlike GParareal, which utilizes only the posterior mean of the GP to correct the coarse solution, Prob-GParareal incorporates the full posterior distribution. The update rule is formulated as a Markov process where the solution at time step $i$ and iteration $k$, denoted $U_{i,k}$, is a random variable:
   $$U_{i,k} = G(U_{i-1,k}) + Z_{i,k}$$
   Here, $Z_{i,k}$ is a random variable drawn from a multivariate normal distribution conditioned on the previous step's realization, with a mean equal to the GP posterior mean and a covariance derived from the GP posterior variance.

2. Sampling Scheme: To propagate this uncertainty through the generally nonlinear coarse solver $G$ without linearizing the vector field (a common requirement in other probabilistic ODE filters), the authors employ an ancestral sampling procedure. The algorithm maintains a set of $n$ samples (realizations) at each time interval. In each iteration, these samples are propagated through the coarse solver, and new correction terms are sampled from the GP posterior conditioned on the specific realization of the previous step.

3. Scalability (Prob-nnGParareal): To address the cubic computational complexity of standard GP training ($O(N^3)$), the framework is extended to Prob-nnGParareal. This variant uses Nearest-Neighbor Gaussian Processes (nnGPs), training on a reduced subset of the data (nearest neighbors to the prediction point) rather than the entire dataset. This significantly reduces the training cost while maintaining predictive accuracy.

4. Stopping Criteria: The algorithm introduces flexible termination rules. Beyond traditional convergence based on solution accuracy, it supports early termination based on:

   • Predefined solution variance thresholds.
   • Computational budget constraints.
   • The observation that the predictive variance accurately reflects the uncertainty induced by incomplete convergence.

Key Contributions

• First Probabilistic PinT Solver: To the authors' knowledge, this is the first probabilistic extension of the (G)Parareal algorithm, bridging the fields of PinT computation and probabilistic numerics.
• Uncertainty Quantification: The method produces probabilistic forecasts of the system's evolution, explicitly quantifying the dynamics of error across both the temporal domain and algorithm iterations. It handles both non-chaotic and chaotic systems, addressing known challenges in representing chaotic behavior with probabilistic ODE filters.
• Support for Random Initial Conditions: The framework generalizes deterministic initial value problems to accommodate probabilistic initial conditions, enabling the solution of systems where precise initial data is unavailable.
• Solver Agnosticism: The method is compatible with any existing deterministic Parareal implementation and any choice of numerical solver ($F$ and $G$), requiring no major modifications to the underlying simulation specifics.
• Theoretical Guarantees: The paper derives theoretical error bounds for the Prob-GParareal solution relative to the true solution (in terms of Wasserstein-2 distance) and relative to the deterministic GParareal solution. These bounds depend on the posterior variance of the GP and the fill distance of the training data.
• Scalability: By integrating nnGPs, the method achieves improved scalability for higher-dimensional systems and PDEs.

Empirical Results
The authors validated Prob-GParareal on five benchmark ODE systems: FitzHugh-Nagumo, Rössler (chaotic), Hopf bifurcation, Double Pendulum, and Lorenz (chaotic).

• Accuracy: The method demonstrated accuracy comparable to deterministic GParareal and Parareal across all systems.
• Calibration: The probabilistic forecasts were well-calibrated, with the predictive variance accurately reflecting the uncertainty. For chaotic systems (Lorenz, Rössler), the uncertainty naturally increased over time, consistent with the Lyapunov time, and the 95% confidence intervals captured the divergence of trajectories.
• Early Termination: Experiments showed that stopping the algorithm before full convergence (e.g., at 80% of the required iterations) yielded results with similar performance metrics (MAD, Energy Score, Variogram Score) to the fully converged solution, provided the variance-based stopping criteria were met. This offers significant runtime savings.
• PDE Application: The Prob-nnGParareal variant was successfully applied to the Viscous Burgers equation (discretized as a 128-dimensional ODE system). It achieved comparable accuracy to deterministic counterparts with reduced computational cost and provided valid uncertainty estimates.

Significance and Claims
The paper claims to bridge a critical gap in the development of probabilistic counterparts to established PinT methods. Its primary significance lies in providing a tool for Uncertainty Quantification in parallel-in-time simulations without requiring linearization of the system or restricting the choice of numerical solvers.

The authors emphasize that Prob-GParareal allows for "flexible and controlled resource allocation," enabling users to trade off computational cost against solution precision based on variance thresholds. While the method relies on the assumption that the fine solver represents the true solution and that coordinate-wise independence (intrinsic coregionalization) is a sufficient approximation for the GP covariance, the authors argue that these assumptions enable computational feasibility for nonlinear and chaotic systems where other probabilistic methods fail. The work positions itself as a foundational step toward integrating UQ into high-performance PinT workflows, particularly for applications involving chaotic dynamics and large-scale PDEs.