Prob-nnGParareal: Local nnGP in Parareal Solvers

• The paper introduces Prob-nnGParareal, a variant of Prob-GParareal that replaces the full GP with local nnGP fits, significantly reducing computational complexity from O((Nk)^3) to O(m^3).
• It employs a localized neighbor selection strategy via kd-tree searches to train on only the m nearest points, ensuring uncertainty quantification and scalability for parallel-in-time differential equation solvers.
• Benchmark tests on viscous Burgers’ equation demonstrate marked improvements in convergence speed and runtime, while achieving low bias and mean squared error compared to conventional Parareal methods.

Prob-nnGParareal is a variant of Prob-GParareal in which the Gaussian-process correction model used within Parareal is replaced by a nearest-neighbors Gaussian process, or nnGP. In the formulation introduced in "Prob-GParareal: A Probabilistic Numerical Parallel-in-Time Solver for Differential Equations" (Gattiglio et al., 4 Sep 2025), the method is a probabilistic Parallel-in-Time solver for ODEs and PDEs that retains uncertainty quantification for the Parareal correction while reducing the prohibitive scaling of standard GP regression. Its defining feature is the use of local GP fits based on the $m \ll Nk$ nearest training points to a query state, rather than a GP trained on all accumulated correction data.

1. Position within the Parareal, GParareal, and Prob-GParareal lineage

Parareal solves an initial value problem over a time domain partitioned into $N$ subintervals by alternating a cheap coarse solver and a fine solver in a predictor–corrector iteration. In the classical formulation, the update is

$$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$

with convergence in $K_{\mathrm{Para}}<N$ iterations giving a speed-up approximately $N/K_{\mathrm{Para}}$ (Gattiglio et al., 2024).

GParareal replaces the raw discrepancy $\delta(u)=\mathcal F(u)-\mathcal G(u)$ by a GP surrogate trained on all previous discrepancy evaluations. This reduces the number of Parareal iterations in many settings, but its covariance inversion grows as $O((kN)^3)$ per iteration because the training set accumulates across iterations (Gattiglio et al., 2024). Prob-GParareal extends that construction by sampling from the full GP posterior rather than using only a posterior mean, thereby quantifying uncertainty in the Parareal correction and propagating numerical uncertainty across time (Gattiglio et al., 4 Sep 2025).

Prob-nnGParareal is the nearest-neighbors counterpart of Prob-GParareal. It inherits the probabilistic formulation of Prob-GParareal while replacing the full GP with local nnGP fits, reducing the cubic dependence on $Nk$ to a cubic dependence on the neighbor count $m$ (Gattiglio et al., 4 Sep 2025). This places it simultaneously in two lineages: the probabilistic extension of GParareal and the scalability-oriented nnGParareal construction of "Nearest Neighbors GParareal: Improving Scalability of Gaussian Processes for Parallel-in-Time Solvers" (Gattiglio et al., 2024).

2. Local probabilistic correction model

The correction function is defined coordinatewise by

$f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$

with GP prior

$N$0

The kernel may be, for example, the squared-exponential kernel

$N$1

or a Matérn kernel $N$2 (Gattiglio et al., 4 Sep 2025).

The probabilistic Parareal update is expressed in terms of random states $N$3: $N$4 with

$N$5

where $N$6 and $N$7 are the posterior mean and covariance of the local nnGP approximation to $N$8 (Gattiglio et al., 4 Sep 2025).

The essential algorithmic modification is local data selection. Rather than train on all

$N$9

the method trains at a test location $$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$0 only on the $$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$1 nearest neighbor inputs $$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$2, selected by Euclidean distance, together with outputs

$$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$3

The resulting nnGP posterior for each coordinate is

$$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$4

$$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$5

This replaces a global GP inversion by repeated local inversions of size $$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$6 (Gattiglio et al., 4 Sep 2025).

A plausible implication is that the method trades a globally coherent covariance model for locality-adapted correction estimates. The construction is explicitly local in state space rather than global in iteration history.

3. Iterative workflow and distributional convergence

The high-level pseudocode in (Gattiglio et al., 4 Sep 2025) treats the state on each interval as an empirical distribution rather than a single iterate. The inputs are an initial measure $$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$7, a sample count $$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$8, coarse and fine solvers $$U_i^{(k+1)}=\mathcal G(U_{i-1}^{(k+1)})+\bigl[\mathcal F(U_{i-1}^{(k)})-\mathcal G(U_{i-1}^{(k)})\bigr],$$9 and $K_{\mathrm{Para}}<N$0, a tolerance $K_{\mathrm{Para}}<N$1, and neighbor count $K_{\mathrm{Para}}<N$2.

Initialization proceeds by drawing $K_{\mathrm{Para}}<N$3 and propagating them through the coarse solver: $K_{\mathrm{Para}}<N$4 At each subsequent iteration, the method first gathers summaries by computing

$K_{\mathrm{Para}}<N$5

for each interval, and then evaluates $K_{\mathrm{Para}}<N$6 and $K_{\mathrm{Para}}<N$7 (Gattiglio et al., 4 Sep 2025).

The training data are augmented via

$K_{\mathrm{Para}}<N$8

A kd-tree is then built over the inputs of $K_{\mathrm{Para}}<N$9, and no global GP inversion is performed. For each sample trajectory and interval, the algorithm finds the $N/K_{\mathrm{Para}}$0 nearest neighbors of $N/K_{\mathrm{Para}}$1, computes $N/K_{\mathrm{Para}}$2 and $N/K_{\mathrm{Para}}$3, draws

$N/K_{\mathrm{Para}}$4

and sets

$N/K_{\mathrm{Para}}$5

Convergence is checked in distribution, for example by a Wasserstein distance on empirical measures, with termination when

$N/K_{\mathrm{Para}}$6

(Gattiglio et al., 4 Sep 2025).

This workflow distinguishes Prob-nnGParareal from deterministic nnGParareal. In nnGParareal, the local GP posterior mean supplies a deterministic correction. In Prob-nnGParareal, the posterior also drives sampling, so the output is a time-indexed family of empirical distributions rather than a single corrected state sequence.

4. Complexity reduction and error-control structure

For the full-GP version Prob-GParareal, the correction cost at iteration $N/K_{\mathrm{Para}}$7 is

$N/K_{\mathrm{Para}}$8

plus

$N/K_{\mathrm{Para}}$9

For Prob-nnGParareal, the corresponding correction cost is

$\delta(u)=\mathcal F(u)-\mathcal G(u)$0

plus

$\delta(u)=\mathcal F(u)-\mathcal G(u)$1

Because $\delta(u)=\mathcal F(u)-\mathcal G(u)$2, this is described as a drastic reduction, typically $\delta(u)=\mathcal F(u)-\mathcal G(u)$3 (Gattiglio et al., 4 Sep 2025).

The deterministic precursor nnGParareal is analyzed similarly in terms of local GP inversions and kd-tree neighbor search. There, the total model cost per iteration is

$\delta(u)=\mathcal F(u)-\mathcal G(u)$4

and the complexity reduction is summarized as a shift from $\delta(u)=\mathcal F(u)-\mathcal G(u)$5 to $\delta(u)=\mathcal F(u)-\mathcal G(u)$6 (Gattiglio et al., 2024). This suggests that Prob-nnGParareal inherits the same geometric advantage of locality, but with an added sampling layer.

The error-control statement in (Gattiglio et al., 4 Sep 2025) is phrased in terms of variance propagation and Wasserstein distance. Under standard Lipschitz and GP-RKHS assumptions, if

$\delta(u)=\mathcal F(u)-\mathcal G(u)$7

then for interval $\delta(u)=\mathcal F(u)-\mathcal G(u)$8 and iteration $\delta(u)=\mathcal F(u)-\mathcal G(u)$9,

$O((kN)^3)$0

and

$O((kN)^3)$1

where

$O((kN)^3)$2

involves the local fill distance $O((kN)^3)$3 of the training set around $O((kN)^3)$4 (Gattiglio et al., 4 Sep 2025).

For the nnGP variant, the same type of bound holds with the fill distance computed over the $O((kN)^3)$5-neighbor subset, assuming the neighbor set adequately covers the local region. This makes the dependence on local data density explicit: the uncertainty and transport-error bounds are tied to how well the selected neighbors cover the relevant part of the state space.

5. Benchmark behavior on viscous Burgers’ equation

The principal PDE example in (Gattiglio et al., 4 Sep 2025) is viscous Burgers’ equation,

$O((kN)^3)$6

on $O((kN)^3)$7 with $O((kN)^3)$8 and periodic boundary conditions, spatially discretized to $O((kN)^3)$9 ODEs. The time domain is partitioned into $Nk$0 intervals over $Nk$1. The coarse solver is $Nk$2 with $Nk$3 steps, and the fine solver is $Nk$4 with $Nk$5 steps (Gattiglio et al., 4 Sep 2025).

The reported convergence and timing figures are as follows.

Method	Convergence / model build	Total runtime
Parareal	$Nk$6	$Nk$7
GParareal	$Nk$8, model build $Nk$9	$m$0
nnGParareal	$m$1, model $m$2 min	$m$3
Prob-nnGParareal $m$4	$m$5, model $m$6 min	$m$7 min

For Prob-nnGParareal, the additional reported metrics relative to the fine solver are bias $m$8 and MSE $m$9. Proper scoring rules for the probabilistic forecast are MAD $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$0, VS $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$1, and ES $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$2 (Gattiglio et al., 4 Sep 2025).

The same experiment is used to illustrate scalability relative to deterministic nearest-neighbor GP Parareal. In that setting, the method is reported as giving approximately a $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$3 speed-up over nnGParareal on this PDE example (Gattiglio et al., 4 Sep 2025). This suggests that, at least for the reported Burgers configuration, the probabilistic local-GP formulation is not merely a UQ layer but also a computationally competitive solver configuration.

6. Integration, limitations, hyperparameters, and nomenclature

Prob-nnGParareal is designed as a minimally invasive modification of an existing Parareal implementation: it replaces only the GP training and prediction module, while leaving the coarse/fine solver calls and the domain decomposition untouched (Gattiglio et al., 4 Sep 2025). This is one reason it sits naturally beside earlier GP-based Parareal variants rather than requiring a new PinT framework.

The limitations recorded in (Gattiglio et al., 4 Sep 2025) are local-data limitations rather than conceptual ones. If $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$4 is too small, or if the training set is not locally dense, the nearest-neighbor approximation degrades; dynamic adaptation of $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$5 may therefore be required. kd-tree search incurs $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$6 per query, which can remain challenging in extreme high dimension. The uncertainty-quantification bounds rely on fill-distance assumptions, so if the sample states $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$7 drift away from the training-data manifold, the error may grow.

The hyperparameter guidelines are likewise local. The number of neighbors should start with $f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$8–$f_c^{(s)}(u)\equiv(F-G)^{(s)}(u),$9 and be increased until local GP posterior variance decays appropriately. Kernel lengthscales $N$00 and regularization $N$01 may be tuned by marginal likelihood locally on each neighbor set, or fixed by “rulers” based on typical distances in $N$02. The number of samples $N$03 should be chosen so that empirical measures converge; $N$04–$N$05 is stated as often sufficient. The convergence tolerance $N$06 is defined on the Wasserstein distance between successive empirical distributions, with values $N$07–$N$08 described as typically yielding good accuracy without excessive iterations. Early stopping can be based on forecast variance when it stops decreasing significantly (Gattiglio et al., 4 Sep 2025).

A recurrent source of confusion is nomenclature. In the usage tied to (Gattiglio et al., 4 Sep 2025), Prob-nnGParareal denotes the nearest-neighbors-GP variant of Prob-GParareal. By contrast, a separate summary associated with "Parallel-in-Time Solutions with Random Projection Neural Networks" (Betcke et al., 2024) uses the same label for a Parareal method whose coarse propagator is a Random Projection Neural Network. The two constructions are technically distinct: one localizes GP corrections within a probabilistic Parareal framework, whereas the other replaces the coarse propagator by an RPNN. The nearest-neighbors-GP meaning is the one explicitly introduced as a variant of Prob-GParareal (Gattiglio et al., 4 Sep 2025).