Stable and Scalable Probabilistic Numerical Solvers for Stiff and High-Dimensional ODEs

Abstract: Filtering-based probabilistic numerical solvers for ordinary differential equations (ODEs) have been established as a flexible and efficient simulation framework with built-in numerical uncertainty quantification. However, problems that are both stiff and high-dimensional remain a challenge, as current methods are either stable and have cubic cost in the ODE dimension, or scale linearly at the expense of stability. In this paper, we close this gap and develop probabilistic ODE solvers that are both stable and scalable. We propose two complementary strategies. First, we develop a matrix-free update step that uses Jacobian-vector products, iterative linear solvers, and stochastic covariance estimation to enable linear scaling, all while retaining stability. Second, we propose iterative re-linearization to further improve stability without sacrificing scalability, turning probabilistic ODE solvers into fully implicit methods. We evaluate the proposed approaches on a range of stiff and high-dimensional problems and demonstrate improved stability and scalability over established probabilistic solvers.

• The paper introduces MatfreeEK1, which leverages Jacobian-vector products and iterative solvers to achieve O(d) scalability without sacrificing stability.
• It employs iterative re-linearization (IEKF) to enhance robustness, yielding implicit updates that maintain solver performance in stiff regimes.
• Empirical results on systems like Lorenz96 and large-scale PDEs validate its efficiency, demonstrating reduced runtimes and improved stability over traditional methods.

Stable and Scalable Probabilistic Numerical Solvers for Stiff and High-Dimensional ODEs

Overview

The paper "Stable and Scalable Probabilistic Numerical Solvers for Stiff and High-Dimensional ODEs" (2606.08203) addresses a fundamental challenge in numerical simulation: the implementation of probabilistic ODE solvers that provide both stability for stiff systems and scalability in high-dimensional regimes. Probabilistic numerics, particularly Bayesian filtering-based ODE solvers (ODE filters), have demonstrated flexibility, built-in uncertainty quantification, and theoretical convergence guarantees, but have historically suffered from a stability-cost trade-off. The authors bridge this gap with two key contributions: (1) the matrix-free MatfreeEK1 solver, which exploits Jacobian-vector products, iterative linear solvers, and stochastic covariance estimation to achieve $O(d)$ scaling without sacrificing stability; and (2) iterative re-linearization (IEKF) to enhance stability, turning the solver into a fully implicit method while maintaining scalability.

Technical Contributions

Matrix-Free Filtering and Covariance Estimation

The MatfreeEK1 approach circumvents the cubic computational bottleneck of traditional A-stable ODE filters by:

• Utilizing exact Jacobian evaluations via automatic differentiation (JVPs/VJPs), avoiding explicit instantiation of large matrices.
• Applying iterative linear solvers (CG, block-GMRES) for mean and covariance updates.
• Introducing stochastic block-diagonal covariance estimation, which maintains the computational benefits of structured priors (e.g., IWP) and ensures linearly-scaling memory and runtime costs.

This enables full-Jacobian implicit filtering with linear complexity in the ODE dimension, a claim substantiated by empirical results on the Lorenz96 system, where MatfreeEK1 achieves wall-clock runtime parity with previous $O(d)$ methods and vastly outperforms EK1 and ExpEK in high-dimensional settings.

Iterative Re-Linearization with IEKF

Iterative Extended Kalman Filtering (IEKF) refines the linearization point at each update, akin to Gauss-Newton iterations for MAP estimation. When combined with MatfreeEK1, stability is further enhanced, especially for explicit or diagonal-Jacobian approximations in DiagonalEK1. The IEKF iterations do not affect per-step computational scaling if iteration counts remain dimension-independent.

Empirical analysis demonstrates that while MatfreeEK1's base approach suffices for stability under adaptive step-size control, IEKF iterations significantly improve robustness at fixed step sizes for certain solvers, notably DiagonalEK1.

Experimental Results

Computational Scalability

MatfreeEK1 consistently demonstrates $O(d)$ scaling on the Lorenz96 system—wall-clock time per solver step increases linearly with dimension, matching explicit methods while providing implicit stability.

Stability Across Stiffness Regimes

On the 1D Brusselator and several stiff PDEs (Burgers, Fisher-KPP), MatfreeEK1 maintains a constant step count as stiffness increases, mirroring the performance of A- or L-stable solvers (EK1, ExpEK), while explicit and diagonal methods require an increasing number of steps (stability-limited). Work-precision diagrams reconfirm MatfreeEK1’s stability, showing a consistent accuracy-cost trade-off across tolerance ranges.

Practical Performance on Large-Scale PDEs

For large-scale discretized PDEs (e.g., FitzHugh-Nagumo with $d=125{,}000$), MatfreeEK1 not only scales but retains stability, solving problems where classic implicit methods (EK1) are computationally infeasible ($O(d^3)$ runtime). In adaptive step-size scenarios, MatfreeEK1 achieves convergent accuracy with significantly reduced runtime (e.g., 27 minutes versus hours for large 2D Fisher-KPP instances).

Iterative Re-Linearization Benefits

IEKF substantially improves DiagonalEK1's stability and runtime, making it competitive in both fixed and adaptive step-size regimes. For MatfreeEK1, IEKF primarily benefits cases where step-size adaptation is absent or restricted, as exact linearization supersedes further iterative gains in typical scenarios.

Limitations and Theoretical Considerations

• Stability Guarantees: While empirical evidence shows MatfreeEK1 is stable, formal A-stability lacks an analogous theoretical proof due to the posterior covariance approximation. The effect of stochastic sample counts in covariance estimation remains uncharacterized and warrants further analysis.
• Posterior Calibration: Uncertainty quantification is not directly evaluated; calibration procedures are adapted from prior work, but the statistical reliability of the probabilistic output is left for future investigation.
• IEKF Utility: The scope of IEKF’s benefit is still partially delineated, particularly its interplay with step-size adaptation and covariance estimation.

Implications and Future Directions

The introduction of MatfreeEK1 and scalable iterative implicit filtering extends the applicability of probabilistic numerical ODE solvers to domains previously restricted by computational cost or instability. This is vital for simulation and inference in high-dimensional PDE contexts (e.g., spatio-temporal biological systems, climate modeling, large-scale Bayesian inverse problems). The convergence of probabilistic numerics and classical iterative methods (Krylov solvers, structured priors) exemplifies a trend toward blending uncertainty quantification and scalable computation.

Further theoretical development in stability analysis, calibration of uncertainty, and automatic adaptation of IEKF iterations is anticipated. Integration with physics-informed inference, real-time simulation, and large-scale Bayesian workflows is expected to expand as these methods mature.

Conclusion

This work presents two complementary strategies that enable stable and scalable probabilistic ODE solvers for stiff and high-dimensional systems. The matrix-free MatfreeEK1 achieves implicit full-Jacobian updates at linear scaling, and iterative re-linearization further enhances stability. The empirical results indicate substantial computational and stability improvements over established solvers. The contributions set a new standard for probabilistic numerical methods, expanding their utility in contemporary computational science and opening pathways for rigorous uncertainty quantification in large-scale simulation (2606.08203).