Title: Probabilistic Exponential Integrators
Venue: NeurIPS 2023
Accepted: yes
Keywords: Probabilistic numerics

differential equations

exponential integrators

Kalman filters

Gaussian processes
Abstract: Probabilistic solvers provide a flexible and efficient framework for simulation, uncertainty quantification, and inference in dynamical systems. However, like standard solvers, they suffer performance penalties for certain stiff systems, where small steps are required not for reasons of numerical accuracy but for the sake of stability. This issue is greatly alleviated in semi-linear problems by the probabilistic exponential integrators developed in this paper. By including the fast, linear dynamics in the prior, we arrive at a class of probabilistic integrators with favorable properties. Namely, they are proven to be L-stable, and in a certain case reduce to a classic exponential integrator---with the added benefit of providing a probabilistic account of the numerical error. The method is also generalized to arbitrary non-linear systems by imposing piece-wise semi-linearity on the prior via Jacobians of the vector field at the previous estimates, resulting in probabilistic exponential Rosenbrock methods. We evaluate the proposed methods on multiple stiff differential equations and demonstrate their improved stability and efficiency over established probabilistic solvers. The present contribution thus expands the range of problems that can be effectively tackled within probabilistic numerics.

Probabilistic Exponential Integrators

Abstract

Probabilistic solvers provide a flexible and efficient framework for simulation, uncertainty quantification, and inference in dynamical systems. However, like standard solvers, they suffer performance penalties for certain stiff systems, where small steps are required not for reasons of numerical accuracy but for the sake of stability. This issue is greatly alleviated in semi-linear problems by the probabilistic exponential integrators developed in this paper. By including the fast, linear dynamics in the prior, we arrive at a class of probabilistic integrators with favorable properties. Namely, they are proven to be L-stable, and in a certain case reduce to a classic exponential integrator-with the added benefit of providing a probabilistic account of the numerical error. The method is also generalized to arbitrary non-linear systems by imposing piece-wise semi-linearity on the prior via Jacobians of the vector field at the previous estimates, resulting in probabilistic exponential Rosenbrock methods . We evaluate the proposed methods on multiple stiff differential equations and demonstrate their improved stability and efficiency over established probabilistic solvers. The present contribution thus expands the range of problems that can be effectively tackled within probabilistic numerics.

1 Introduction

Dynamical systems appear throughout science and engineering, and their accurate and efficient simulation is a key component in many scientific problems. There has also been a surge of interest in the intersection with machine learning, both regarding the usage of machine learning methods to model and solve differential equations [36, 18, 35], and in a dynamical systems perspective on machine learning methods themselves [8, 5]. This paper focuses on the numerical simulation of dynamical systems within the framework of probabilistic numerics , which treats the numerical solvers themselves as probabilistic inference methods [11, 12, 33]. In particular, we expand the range of problems that can be tackled within this framework and introduce a new class of stable probabilistic numerical methods for stiff ordinary differential equations (ODEs).

Stiff equations are problems for which certain implicit methods perform much better than explicit ones [10]. But implicit methods come with increased computational complexity per step, as they typically require solving a system of equations. Exponential integrators are an alternative class of methods for efficiently solving large stiff problems [48, 16, 7, 15]. They are based on the observation that, if the ODE has a semi-linear structure, the linear part can be solved exactly and only the non-linear part needs to be numerically approximated. The resulting methods are formulated in an explicit manner and do not require solving a system of equations, while achieving similar or better stability than implicit methods. However, such methods have not yet been formulated probabilistically.

Filip Tronarp

t

0

3

y

(

t

)

0

1

a.

Explicit method (not A-stable)

dt

=0.01

t

0

3

b.

Semi-implicit method (A-stable)

dt

=0.25

t

0

3

c.

Exponential integrator (L-stable)

dt

=0.25

In this paper we develop probabilistic exponential integrators , a new class of probabilistic numerical solvers for stiff semi-linear ODEs. We build on the ODE filters which have emerged as an efficient and flexible class of probabilistic numerical methods for general ODEs [40, 21, 45]. They have known convergence rates [21, 46], which have also been demonstrated empirically [2, 26, 24], they are applicable to a wide range of numerical differential equation problems [23, 25, 3], their probabilistic output can be integrated into larger inference problems [20, 39, 47], and they can be formulated parallel-in-time [4]. But while it has been shown that the choice of underlying Gauss-Markov prior does influence the resulting ODE solver [30, 45, 2], there has not yet been strong evidence for the utility of priors other than the well-established integrated Wiener process. Probabilistic exponential integrators provide this evidence: in the probabilistic numerics framework, 'solving the linear part of the ODE exactly' corresponds to an appropriate choice of prior.

Contributions Our main contribution is the development of probabilistic exponential integrators, a new class of stable probabilistic solvers for stiff semi-linear ODEs. We demonstrate the close link of these methods to classic exponential integrators in Proposition 1, provide an equivalence result to a classic exponential integrator in Proposition 2, and prove their L-stability in Proposition 3. To enable a numerically stable implementation, we present a quadrature-based approach to directly compute square-roots of the process noise covariance in Section 3.2. Finally, in Section 3.6 we also propose probabilistic exponential Rosenbrock methods for problems in which semi-linearity is not known a priori. We evaluate all proposed methods on multiple stiff problems and demonstrate the improved stability and efficiency of the probabilistic exponential integrators over existing probabilistic solvers.

2 Numerical ODE solutions as Bayesian state estimation

Let us first consider an initial value problem with some general non-linear ODE, of the form

with vector field f : R d × R → R d , initial value y 0 ∈ R d , and time span [0 , T ] . Probabilistic numerical ODE solvers aim to compute a posterior distribution over the ODE solution y ( t ) such that it satisfies the ODE on a discrete set of points T = { t n } N n =0 ⊂ [0 , T ] , that is

We call this quantity, and approximations thereof, a probabilistic numerical ODE solution . Probabilistic numerical ODE solvers thus compute not just a single point estimate of the ODE solution, but a posterior distribution which provides a structured estimate of the numerical approximation error.

In the following, we briefly recapitulate the probabilistic ODE filter framework of Schober et al. [40] and Tronarp et al. [45] and define the prior, data model, and approximate inference scheme. In Section 3 we build on these foundations to derive the proposed probabilistic exponential integrator.

2.1 Gauss-Markov prior

A priori , we model y ( t ) with a Gauss-Markov process, defined by a stochastic differential equation

with state Y ( t ) ∈ R d ( q +1) , model matrices A ∈ R d ( q +1) × d ( q +1) , B ∈ R d ( q +1) × d , diffusion scaling κ ∈ R , and smoothness q ∈ N . More precisely, A and B are chosen such that the state is structured as Y ( t ) = [ Y (0) ( t ) , . . . , Y ( q ) ( t )] , and then Y ( i ) ( t ) models the i -th derivative of y ( t ) . The initial value Y 0 ∈ R d ( q +1) must be chosen such that it enforces the initial condition, that is, Y (0) (0) = y 0 .

One concrete example of such a Gauss-Markov process that is commonly used in the context of probabilistic numerical ODE solvers is the q -times Integrated Wiener process, with model matrices

Alternatives include the class of Matérn processes and the integrated Ornstein-Uhlenbeck process [46]-the latter plays a central role in this paper and will be discussed in detail later.

Y ( t ) satisfies linear Gaussian transition densities of the form [44]

with transition matrices Φ( h ) and Q ( h ) given by

These quantities can be computed with a matrix fraction decomposition [44]. For q -times integrated Wiener process priors, closed-form expressions for the transition matrices are available [21].

2.2 Information operator

The likelihood, or data model, of a probabilistic ODE solver relates the uninformed prior to the actual ODE solution of interest with an information operator I [6], defined as

where E i ∈ R d × d ( q +1) are selection matrices such that E i Y ( t ) = Y ( i ) ( t ) . I [ Y ] then captures how well Y solves the given ODE problem. In particular, I maps the true ODE solution y to the zero function, i.e. I [ y ] ≡ 0 . Conversely, if I [ y ]( t ) = 0 holds for all t ∈ [0 , T ] then y solves the given ODE. Unfortunately, it is in general infeasible to solve an ODE exactly and enforce I [ Y ]( t ) = 0 everywhere, which is why numerical ODE solvers typically discretize the time interval and take discrete steps. This leads to the data model used in most probabilistic ODE solvers [45]:

Note that this specific information operator is closely linked to the IVP considered in Eq. (1a). By defining a (slightly) different data model we can also formulate probabilistic numerical IVP solvers for higher-order ODEs or differential-algebraic equations, or encode additional information such as conservation laws or noisy trajectory observations [3, 39].

2.3 Approximate Gaussian inference

The resulting inference problem is described by a Gauss-Markov prior and a Dirac likelihood

with Φ n := Φ( t n +1 -t n ) , Q n := Q ( t n +1 -t n ) , initial value Y (0) = Y 0 , discrete time grid { t n } N n =0 , and zero-valued data Z n = 0 for all n . The solution of the resulting non-linear Gauss-Markov regression problem can then be efficiently approximated with Bayesian filtering and smoothing techniques [37]. Notable examples that have been used to construct probabilistic numerical ODE solvers include quadrature filters, the unscented Kalman filter, the iterated extended Kalman smoother, or particle filters [19, 45, 46]. Here, we focus on the well-established extended Kalman filter (EKF). We briefly discuss the EKF for the given state estimation problem in the following.

Prediction Given a Gaussian state estimate Y ( t n -1 ) ∼ N ( µ n -1 , Σ n -1 ) and the linear conditional distribution as in Eq. (9a), the marginal distribution Y ( t n ) ∼ N ( µ -n , Σ -n ) is also Gaussian, with

Linearization To efficiently compute a tractable approximation of the true posterior, the EKF linearizes the information operator I around the predicted mean µ -n , i.e. I [ Y ]( t n ) ≈ H n Y ( t n ) + b n ,

An exact linearization with Jacobian F y = ∂ y f ( E 0 µ -n , t n ) leads to a semi-implicit probabilistic ODE solver, which we call the EK1 [45]. Other choices include the zero matrix F y = 0 , which results in the explicit EK0 solver [40, 21], or a diagonal Jacobian approximation (the DiagonalEK1 ) which combines some stability benefits of the EK1 with the lower computational cost of the EK0 [24].

Correction step In the linearized observation model, the posterior distribution of Y ( t n ) given the datum Z n is again Gaussian. Its posterior mean and covariance ( µ n , Σ n ) are given by

This is also known as the update step of the EKF.

Smoothing To condition the state estimates on all data, the EKF can be followed by a smoothing pass. Starting with µ S N := µ N and Σ S N := Σ n , it consists of the following backwards recursion:

Result The above computations result in a probabilistic numerical ODE solution with marginals

which, by construction of the state Y , also contains estimates for the ODE solution as y ( t ) = E 0 Y ( t ) . Since the EKF-based probabilistic solver does not compute only the marginals in Eq. (14), but a full posterior distribution for the continuous object y ( t ) , it can be evaluated for times t / ∈ T (also known as 'dense output' in the context of ODE solvers); it can produce joint samples from this posterior; and it can be used as a Gauss-Markov prior for subsequent inference tasks [40, 2, 47].

2.4 Practical considerations and implementation details

To improve numerical stability and preserve positive-semidefiniteness of the computed covariance matrices, probabilistic ODE solvers typically operate on square-roots of covariance matrices, defined by a matrix decomposition of the form M = √ M √ M glyph[latticetop] [26]. For example, the Cholesky factor is one possible square-root of a positive definite matrix. But in general, the algorithm does not require the square-roots to be upper- or lower-triangular, or even square. Additionally, we compute the exact initial state Y 0 from the IVP using Taylor-mode automatic differentiation [9, 26], we compute smoothing estimates with preconditioning [26], and we calibrate uncertainties globally with a quasi-maximum likelihood approach [45, 2].

3 Probabilistic exponential integrators

In the remainder of the paper, unless otherwise stated, we focus on IVPs with a semi-linear vector-field

Assuming N admits a Taylor series expansion around t , the variation of constants formula provides a formal expression of the solution at time t + h :

This observation is the starting point for the development of exponential integrators [31, 15]. By further defining the so-called ϕ -functions

the above identity of the ODE solution simplifies to

In this section we develop a class of probabilistic exponential integrators . This is achieved by defining an appropriate class of priors that absorbs the partial linearity, which leads to the integrated Ornstein-Uhlenbeck processes. Proposition 1 below directly relates this choice of prior to the classical exponential integrators. Proposition 2 demonstrates a direct equivalence between the predictor-corrector form exponential trapezoidal rule and the once integrated Ornstein-Uhlenbeck process. Furthermore, the favorable stability properties of classical exponential integrators is retained for the probabilistic counterparts as shown in Proposition 3.

3.1 The integrated Ornstein-Uhlenbeck process

In Section 2.1 we highlighted the choice of the q -times integrated Wiener process prior, which essentially corresponds to modeling the ( q -1) -th derivative of the right-hand side f with a Wiener process. Here we follow a similar motivation, but only for the non-linear part N . Differentiating both sides of Eq. (15) q -1 times with respect to t yields

Then, modeling d q -1 d t q -1 N ( y ( t ) , t ) as a Wiener process and relating the result to y ( t ) gives

This process is also known as the q -times integrated Ornstein-Uhlenbeck process (IOUP), with rate parameter L and diffusion parameter κ . It can be equivalently stated with the previously introduced notation (Section 2.1), by defining a state Y ( t ) , as the solution of a linear time-invariant (LTI) SDE as in Eq. (3), with system matrices

Remark 1 (The mean of the IOUP process solves the linear part of the ODE exactly) . By taking the expectation of Eq. (20b) and by linearity of integration, we can see that the mean of the IOUP satisfies

This is in line with the motivation of exponential integrators: the linear part of the ODE is solved exactly, and we only need to approximate the non-linear part. Figure 2 visualizes this idea.

t

0

5

y

(

t

)

-3

0

3

a.

t

0

5

b.

t

0

5

c.

IOUP + initial value

3.2 The transition parameters of the integrated Ornstein-Uhlenbeck process

Since the process Y ( t ) is defined as the solution of a linear time-invariant SDE, it satisfies discrete transition densities p ( Y ( t + h ) | Y ( t )) = N ( Φ( h ) Y ( t ) , κ 2 Q ( h ) ) . The following result shows that the transition parameters are intimately connected with the ϕ -functions defined in Eq. (17).

Proposition 1. The transition matrix of a q -times integrated Ornstein-Uhlenbeck process satisfies

Proof in Appendix A. Although Proposition 1 indicates that Φ( h ) may be computed more efficiently than by calling a matrix-exponential on a d ( q +1) × d ( q +1) matrix, this is known to be numerically less stable [41]. We therefore compute Φ( h ) with the standard matrix-exponential formulation.

Directly computing square-roots of the process noise covariance Numerically stable probabilistic ODEsolvers require a square-root, √ Q ( h ) , of the process noise covariance rather than the full matrix, Q ( h ) . For IWP priors this can be computed from the closed-form representation of Q ( h ) via an appropriately preconditioned Cholesky factorization [26]. However, for IOUP priors we have not found an analogous method that works reliably. Therefore, we compute √ Q ( h ) directly with numerical quadrature. More specifically, given a quadrature rule with nodes τ i ∈ [0 , h ] and positive weights w i > 0 , the integral for Q ( h ) given in Eq. (6) is approximated by

with square-roots √ M i = √ w i exp( A ( h -τ i )) B of the summands, which is well-defined since w i > 0 . We can thus compute a square-root representation of the sum with a QR-decomposition

We obtain Q ( h ) ≈ R glyph[latticetop] R , and therefore an approximate square-root factor is given by √ Q ( h ) ≈ R glyph[latticetop] . Similar ideas have previously been employed for time integration of Riccati equations [42, 43]. We use this quadrature-trick for all IOUP methods, with Gauss-Legendre quadrature on m = q nodes.

3.3 Linearization and correction

The information operator of the probabilistic exponential integrator is defined exactly as in Section 2.2. But since we now assume a semi-linear vector-field f , we have an additional option for the linearization: instead of choosing the exact F y = ∂ y f ( EK1 ) or the zero-matrix F y = 0 ( EK0 ), a cheap approximate Jacobian is given by the linear part F y = L . We denote this approach by EKL . This is chosen as the default for the probabilistic exponential integrator. Note that the EKL approach can also be combined with an IWP prior, which will be serve as an additional baseline in the Section 4.

3.4 Equivalence to the classic exponential trapezoidal rule in predict-evaluate-correct mode

Now that the probabilistic exponential integrator has been defined, we can establish an equivalence result to a classic exponential integrator, similarly to the closely-related equivalence statement by Schober et al. [40, Proposition 1] for the non-exponential case.

Proposition 2 (Equivalence to the PEC exponential trapezoidal rule) . The mean estimate of the probabilistic exponential integrator with a once-integrated Ornstein-Uhlenbeck prior with rate parameter L is equivalent to the classic exponential trapezoidal rule in predict-evaluate-correct mode, with the predictor being the exponential Euler method. That is, it is equivalent to the scheme

where Eq. (26a) corresponds to a prediction step with the exponential Euler method, and Eq. (26b) corresponds to a correction step with the exponential trapezoidal rule.

The proof is given in Appendix B. This equivalence result provides another theoretical justification for the proposed probabilistic exponential integrator. But note that the result only holds for the mean, while the probabilistic solver computes additional quantities in order to track the solution uncertainty, namely covariances. These are not provided by a classic exponential integrator.

3.5 L-stability of the probabilistic exponential integrator

When solving stiff ODEs, the actual efficiency of a numerical method often depends on its stability. One such property is A-stability : It guarantees that the numerical solution of a decaying ODE will also decay, independently of the chosen step size. In contrast, explicit methods typically only decay for sufficiently small steps. In the context of probabilistic ODE solvers, the EK0 is considered to be explicit, but the EK1 with IWP prior has been shown to be A-stable [45]. Here, we show that the probabilistic exponential integrator satisfies the stronger L-stability : the numerical solution not only decays, but it decays fast , i.e. it goes to zero as the step size goes to infinity. Figure 1 visualizes the different probabilistic solver stabilities. For formal definitions, see for example [27, Section 8.6].

Proposition 3 (L-stability) . The probabilistic exponential integrator is L-stable.

The full proof is given in Appendix C. The property essentially follows from Remark 1 which stated that the IOUP solves linear ODEs exactly. This implies fast decay and gives L-stability.

3.6 Probabilistic exponential Rosenbrock-type methods

We conclude with a short excursion into exponential Rosenbrock methods [14, 17, 28]: Given a nonlinear ODE ˙ y ( t ) = f ( y ( t ) , t ) , exponential Rosenbrock methods perform a continuous linearization of the right-hand side f around the numerical ODE solution and essentially solve a sequence of IVPs

where J n is the Jacobian of f at the numerical solution estimate ˆ y ( t n ) . This approach enables exponential integrators for problems where the right-hand side f is not semi-linear. Furthermore, by automatically linearizing along the numerical solution the linearization can be more accurate, the Lipschitz-constant of the non-linear remainder becomes smaller, and the resulting solvers can thus be more efficient than their globally linearized counterparts [17].

This can also be done in the probabilistic setting: By linearizing the ODE right-hand side f at each step of the solver around the filtering mean E 0 µ n , we (locally) obtain a semi-linear problem. Then, updating the rate parameter of the integrated Ornstein-Uhlenbeck process at each step of the numerical solver results in probabilistic exponential Rosenbrock-type methods . As before, the linearization of the information operator can be done with any of the EK0 , EK1 , or EKL . But since here the prediction relies on exact local linearization, we will by default also use an exact EK1 linearization. The resulting solver and its stability and efficiency will be evaluated in the following experiments.

̇

Number of steps

1

10

100

Final error

10

⁰

10

⁻

¹

⁰

10

⁻

²

⁰

y

=

-

y

+10

-

1

⋅

y

2

Number of steps

1

10

100

y

=

-

y

+10

-

4

⋅

y

2

Number of steps

1

10

100

y

=

-

y

+10

-

7

⋅

y

2

Number of steps

1

10

100

y

=

-

y

+10

-

10

⋅

y

2

EK0 & IWP(2)

EK1 & IWP(2)

EKL & IWP(2)

EKL & IOUP(2)

̇

̇

̇

4 Experiments

In this section we investigate the utility and performance of the proposed probabilistic exponential integrators and compare it to standard non-exponential probabilistic solvers on multiple ODEs. All methods are implemented in the Julia programming language [1], with special care being taken to implement the solvers in a numerically stable manner, that is, with exact state initialization, preconditioned state transitions, and a square-root implementation [26]. Reference solutions are computed with the DifferentialEquations.jl package [34]. All experiments run on a single, consumerlevel CPU. Code for the implementation and experiments is publicly available on GitHub. 1

4.1 Logistic equation with varying degrees of non-linearity

We start with a simple one-dimensional initial value problem: a logistic model with negative growth rate parameter r = -1 and carrying capacity K ∈ R + , of the form

The non-linearity of this problem can be directly controlled through the parameter K . Therefore, this test problem lets us investigate the IOUP's capability to leverage known linearity in the ODE.

We compare the proposed exponential integrator to all introduced IWP-based solvers, with different linearization strategies: EK0 approximates ∂ y f ≈ 0 (and is thus explicit), EKL approximates ∂ y f ≈ -1 , and EK1 linearizes with the correct Jacobian ∂ y f . The results for four different values of K are shown in Fig. 3. The explicit solver shows the largest error of all compared solvers, likely due to its lacking stability. On the other hand, the proposed exponential integrator behaves as expected: the IOUP prior is most beneficial for larger values of K , and as the non-linearity becomes more pronounced the performance of the IOUP approaches that of the IWP-based solver. Though for large step sizes, the IOUP outperforms the IWP prior even for the most non-linear case with K = 10 .

4.2 Burger's equation

Here, we consider Burger's equation, which is a semi-linear partial differential equation (PDE)

with diffusion coefficient D ∈ R + . Wetransform the problem into a semi-linear ODE with the method of lines [29, 38], and discretize the spatial domain on 250 equidistant points and approximate the differential operators with finite differences. The full IVP specification, including all domains, initial and boundary conditions, and additional information on the discretization, is given in Appendix D.

The results shown in Fig. 4 demonstrate the different stability properties of the solvers: The explicit EK0 with IWP prior is unable to solve the IVP for any of the step sizes due to its insufficient stability, and even the A-stable EK1 and the more approximate EKL require small enough steps ∆ t < 10 -1 . On the other hand, both exponential integrators are able to compute meaningful solutions for a larger range of step sizes. They both achieve lower errors for most settings than their non-exponential

Space

0

1

Time

0

1

a.

ODE solution

-0.5

0.0

0.5

Step size

10 -3

10 -2

10 -1

10 0

Final error

10 -10

10 0

b.

Work-precision diagram

Number of f evaluations

10 1

10 2

10 3

EK0 & IWP(2)

EK1 & IWP(2)

EKL & IWP(2)

EKL & IOUP(2)

EK1 & IOUP(2) (RB)

Space

0

1

Time

0

2

a.

ODE solution

0.0

0.5

1.0

Step size

10 -2

10 -1

10 0

Final error

10 -5

10 0

10 5

b.

Work-precision diagram

Runtime [s]

10 -1

10 0

10 1

EK0 & IWP(2)

EK1 & IWP(2)

EKL & IWP(2)

EKL & IOUP(2)

EK1 & IOUP(2) (RB)

counterparts. The second diagram in Fig. 4 compares the achieved error to the number of vector-field evaluations and points out a trade-off between both exponential methods: Since the Rosenbrock method additionally computes two Jacobians (with automatic differentiation) per step, it needs to evaluate the vector-field more often than the non-Rosenbrock method. Thus, for expensive-to-evaluate vector fields the standard probabilistic exponential integrator might be preferable.

4.3 Reaction-diffusion model

Finally, we consider a discretized reaction-diffusion model given by a semi-linear PDE

where D ∈ R + is the diffusion coefficient and R ( u ) = u (1 -u ) is a logistic reaction term [22]. A finite-difference discretization of the spatial domain transforms this PDE into an IVP with semi-linear ODE. The full problem specification is provided in Appendix D.

Figure 5 shows the results. We again observe the improved stability of the exponential integrator variants by their lower error for large step sizes, and they outperform the IWP-based methods on all settings. The runtime-evaluation in Fig. 5 also visualizes another drawback of the Rosenbrock-type method: Since the problem is re-linearized at each step, the IOUP also needs to be re-discretized and thus a matrix exponential needs to be computed. In comparison, the non-Rosenbrock method only discretizes the IOUP prior once at the start of the solve. This advantage makes the non-Rosenbrock probabilistic exponential integrator the most performant solver in this experiment.

5 Limitations

The probabilistic exponential integrator shares many properties of both classic exponential integrators and of other filtering-based probabilistic solvers. This also brings some challenges.

Cost of computing matrix exponentials The IOUP prior is more expensive to discretize than the IWP as it requires computing a matrix exponential. This trade-off is well-known also in the context of classic exponential integrators. One approach to reduce computational cost is to compute the matrix exponential only approximately [32], for example with Krylov-subspace methods [13, 17]. Extending these techniques to the probabilistic solver setting thus poses an interesting direction for future work.

Cubic scaling in the ODE dimension The probabilistic exponential integrator shares the complexity most (semi-)implicit ODE solvers: while being linear in the number of time steps, it scales cubically in the ODE dimension. By exploiting structure in the Jacobian and in the prior, some filtering-based ODE solvers have been formulated with linear scaling in the ODE dimension [24]. But this approach does not directly extend to the IOUP-prior. Nevertheless, exploiting known structure could be particularly relevant to construct solvers for specific ODEs, such as certain discretized PDEs.

6 Conclusion

We have presented probabilistic exponential integrators, a new class of probabilistic solvers for stiff semi-linear ODEs. By incorporating the fast, linear dynamics directly into the prior of the solver, the method essentially solves the linear part exactly, in a similar manner as classic exponential integrators. We also extended the proposed method to general non-linear systems via iterative re-linearization and presented probabilistic exponential Rosenbrock-type methods. Both methods have been shown both theoretically and empirically to be more stable than their non-exponential probabilistic counterparts. This work further expands the toolbox of probabilistic numerics and opens up new possibilities for accurate and efficient probabilistic simulation and inference in stiff dynamical systems.

Acknowledgments and Disclosure of Funding

The authors gratefully acknowledge financial support by the German Federal Ministry of Education and Research (BMBF) through Project ADIMEM (FKZ 01IS18052B), and financial support by the European Research Council through ERC StG Action 757275 / PANAMA; the DFG Cluster of Excellence Machine Learning - New Perspectives for Science, EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg. Filip Tronarp was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Nathanael Bosch. The authors also thank Jonathan Schmidt for many valuable discussions and for helpful feedback on the manuscript.

References

J. Bezanson, A. Edelman, S. Karpinski, and V. B. Shah. Julia: A fresh approach to numerical computing. SIAM review , 59(1):65-98, 2017.

N. Bosch, P. Hennig, and F. Tronarp. Calibrated adaptive probabilistic ODE solvers. In International Conference on Artificial Intelligence and Statistics . PMLR, 2021.

N. Bosch, F. Tronarp, and P. Hennig. Pick-and-mix information operators for probabilistic ODE solvers. In International Conference on Artificial Intelligence and Statistics . PMLR, 2022.

N. Bosch, A. Corenflos, F. Yaghoobi, F. Tronarp, P. Hennig, and S. Särkkä. Parallel-in-time probabilistic numerical ODE solvers, 2023.

R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems . Curran Associates, Inc., 2018.

J. Cockayne, C. Oates, T. Sullivan, and M. Girolami. Bayesian probabilistic numerical methods. SIAM Review , 61:756-789, 2019.

S. M. Cox and P. C. Matthews. Exponential time differencing for stiff systems. Journal of Computational Physics , 176(2):430-455, 2002.

W. E. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics , 5(1):111, Mar 2017.

A. Griewank and A. Walther. Evaluating Derivatives . Society for Industrial and Applied Mathematics, second edition, 2008.

E. Hairer and G. Wanner. Solving Ordinary Differential Equations II: Stiff and DifferentialAlgebraic Problems . Springer series in computational mathematics. Springer-Verlag, 1991.

P. Hennig, M. A. Osborne, and M. Girolami. Probabilistic numerics and uncertainty in computations. Proceedings. Mathematical, physical, and engineering sciences , 471(2179): 20150142-20150142, Jul 2015.

P. Hennig, M. A. Osborne, and H. P. Kersting. Probabilistic Numerics: Computation as Machine Learning . Cambridge University Press, 2022.

M. Hochbruck and C. Lubich. On Krylov subspace approximations to the matrix exponential operator. SIAM Journal on Numerical Analysis , 34(5):19111925, Oct 1997.

M. Hochbruck and A. Ostermann. Explicit integrators of Rosenbrock-type. Oberwolfach Reports , 3(2):1107-1110, 2006.

M. Hochbruck and A. Ostermann. Exponential integrators. Acta Numerica , 19:209286, 2010.

M. Hochbruck, C. Lubich, and H. Selhofer. Exponential integrators for large systems of differential equations. SIAM Journal on Scientific Computing , 19(5):1552-1574, 1998.

M. Hochbruck, A. Ostermann, and J. Schweitzer. Exponential Rosenbrock-type methods. SIAM Journal on Numerical Analysis , 47(1):786-803, 2009.

G. E. Karniadakis, I. G. Kevrekidis, L. Lu, P. Perdikaris, S. Wang, and L. Yang. Physicsinformed machine learning. Nature Reviews Physics , 3(6):422440, May 2021.

H. Kersting and P. Hennig. Active uncertainty calibration in Bayesian ode solvers. In Proceedings of the 32nd Conference on Uncertainty in Artificial Intelligence (UAI) , pages 309-318, June 2016.

H. Kersting, N. Krämer, M. Schiegg, C. Daniel, M. Tiemann, and P. Hennig. Differentiable likelihoods for fast inversion of Likelihood-free dynamical systems. In International Conference on Machine Learning . PMLR, 2020.

H. Kersting, T. J. Sullivan, and P. Hennig. Convergence rates of Gaussian ODE filters. Statistics and computing , 30(6):1791-1816, 2020.

A. N. Kolmogorov. A study of the equation of diffusion with increase in the quantity of matter, and its application to a biological problem. Moscow University Bulletin of Mathematics , 1:1-25, 1937.

N. Krämer and P. Hennig. Linear-time probabilistic solution of boundary value problems. In Advances in Neural Information Processing Systems . Curran Associates, Inc., 2021.

N. Krämer, N. Bosch, J. Schmidt, and P. Hennig. Probabilistic ODE solutions in millions of dimensions. In International Conference on Machine Learning . PMLR, 2022.

N. Krämer, J. Schmidt, and P. Hennig. Probabilistic numerical method of lines for timedependent partial differential equations. In International Conference on Artificial Intelligence and Statistics . PMLR, 2022.

N. Krämer and P. Hennig. Stable implementation of probabilistic ode solvers, 2020.

J. D. Lambert. Computational Methods in Ordinary Differential Equations . Introductory Mathematics for Scientists And Engineers. Wiley, 1973.

V. T. Luan and A. Ostermann. Exponential Rosenbrock methods of order five construction, analysis and numerical comparisons. Journal of Computational and Applied Mathematics , 255: 417-431, 2014.

N. K. Madsen. The method of lines for the numerical solution of partial differential equations. Proceedings of the SIGNUM meeting on Software for partial differential equations , 1975.

E. Magnani, H. Kersting, M. Schober, and P. Hennig. Bayesian filtering for ODEs with bounded derivatives, 2017.

B. V. Minchev and W. M. Wright. A review of exponential integrators for first order semi-linear problems. Technical report, Norges Teknisk-Naturvetenskaplige Universitet, 2005.

C. Moler and C. Van Loan. Nineteen dubious ways to compute the exponential of a matrix, twenty-five years later. SIAM Review , 45(1):349, Jan 2003.

C. J. Oates and T. J. Sullivan. A modern retrospective on probabilistic numerics. Statistics and Computing , 29(6):1335-1351, 2019.

C. Rackauckas and Q. Nie. DifferentialEquations.jl a performant and feature-rich ecosystem for solving differential equations in julia. Journal of Open Research Software , 5(1), 2017.

C. Rackauckas, Y. Ma, J. Martensen, C. Warner, K. Zubov, R. Supekar, D. Skinner, A. Ramadhan, and A. Edelman. Universal differential equations for scientific machine learning, 2021.

M. Raissi, P. Perdikaris, and G. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics , 378:686707, Feb 2019.

S. Särkkä. Bayesian Filtering and Smoothing , volume 3 of Institute of Mathematical Statistics textbooks . Cambridge University Press, 2013.

W. E. Schiesser. The numerical method of lines: integration of partial differential equations . Elsevier, 2012.

J. Schmidt, N. Krämer, and P. Hennig. A probabilistic state space model for joint inference from differential equations and data. In Advances in Neural Information Processing Systems . Curran Associates, Inc., 2021.

M. Schober, S. Särkkä, and P. Hennig. A probabilistic model for the numerical solution of initial value problems. Statistics and Computing , 29(1):99-122, Jan 2019.

R. B. Sidje. Expokit: A software package for computing matrix exponentials. ACMTransactions on Mathematical Software , 24(1):130156, mar 1998.

T. Stillfjord. Low-rank second-order splitting of large-scale differential Riccati equations. IEEE Transactions on Automatic Control , 60(10):2791-2796, 2015.

T. Stillfjord. Adaptive high-order splitting schemes for large-scale differential Riccati equations. Numerical Algorithms , 78(4):11291151, Sep 2017.

S. Särkkä and A. Solin. Applied Stochastic Differential Equations . Institute of Mathematical Statistics Textbooks. Cambridge University Press, 2019.

F. Tronarp, H. Kersting, S. Särkkä, and P. Hennig. Probabilistic solutions to ordinary differential equations as nonlinear Bayesian filtering: a new perspective. Statistics and Computing , 29(6): 1297-1315, 2019.

F. Tronarp, S. Särkkä, and P. Hennig. Bayesian ODE solvers: The maximum a posteriori estimate. Statistics and Computing , 31(3):1-18, 2021.

F. Tronarp, N. Bosch, and P. Hennig. Fenrir: Physics-enhanced regression for initial value problems. In International Conference on Machine Learning . PMLR, 2022.

C. Van Loan. Computing integrals involving the matrix exponential. IEEE Transactions on Automatic Control , 23(3):395-404, 1978.

Probabilistic Exponential Integrators - Appendix

A Proof of Proposition 1: Structure of the transition matrix

Proof of Proposition 1. The drift-matrix A IOUP ( d,q ) as given in Eq. (21) has block structure

where E q -1 := [0 . . . 0 I d ] glyph[latticetop] ∈ R dq × d . From Van Loan [48, Theorem 1], it follows

which is precisely Eq. (23). The same theorem also gives Φ 12 ( h ) as

Its i th d × d block is readily given by

where the second last equality used the change of variables τ = h (1 -u ) , and the last line follows by definition.

B Proof of Proposition 2: Equivalence to a classic exponential integrator

We first briefly recapitulate the probabilistic exponential integrator setup for the case of the once integrated Ornstein-Uhlenbeck process, and then provide some auxiliary results. Then, we prove Proposition 2 in Appendix B.3.

B.1 The probabilistic exponential integrator with once-integrated Ornstein-Uhlenbeck prior

The integrated Ornstein-Uhlenbeck process prior with rate parameter L results in transition densities Y ( t + h ) | Y ( t ) ∼ N ( Y ( t + h ); Φ( h ) Y ( t ) , Q ( h )) , with transition matrices (from Proposition 1)

where we assume a unit diffusion σ 2 = 1 . To simplify notation, we assume an equidistant time grid T = { t n } N n =0 with t n = n · h for some step size h , and we denote the constant transition matrices simply by Φ and Q and write Y n = Y ( t n ) .

Before getting to the actual proof, let us also briefly recapitulate the filtering formulas that are computed at each solver step. Given a Gaussian distribution Y n ∼ N ( Y n ; µ n , Σ n ) , the prediction step computes

Then, the combined linearization and correction step compute

with observation matrix H = E 1 -LE 0 = [ -L I ] , since we perform the proposed EKL linearization.

B.2 Auxiliary results

In the following, we show some properties of the transition matrices and the covariances that will be needed in the proof of Proposition 2 later.

First, note that by defining ϕ 0 ( z ) = exp z , the ϕ -functions satisfy the following recurrence formula:

See e.g. Hochbruck and Ostermann [15]. This property will be used throughout the remainder of the section.

Lemma B.1. The transition matrices Φ( h ) , Q ( h ) of the once integrated Ornstein-Uhlenbeck process with rate parameter L satisfy

Proof.

where we used Lτϕ 1 ( Lτ ) = ϕ 0 ( Lτ ) -I , and ∂ τ [ τ k ϕ k ( Lτ ) ] = τ k -1 ϕ k -1 ( Lτ ) . It follows that

where we used Lτϕ 2 ( Lτ ) = ϕ 1 ( Lτ ) -I .

Lemma B.2. The prediction covariance Σ -n +1 satisfies

Proof. First, since the observation model is noiseless, the filtering covariance Σ n satisfies

This can be shown directly from the correction step formula:

Next, since the observation matrix is H = [ -L I ] , the filtering covariance Σ n is structured as

This can be shown directly from Eq. (58):

and thus

It follows

Finally, together with Lemma B.1 we can derive the result:

B.3 Proof of Proposition 2

With these results, we can now prove Proposition 2.

Proof of Proposition 2. We prove the proposition by induction, showing that the filtering means are all of the form

where y n , ˜ y n are defined as

This result includes the statement of Proposition 2.

Base case n = 0 The initial distribution of the probabilistic solver is chosen as

This proves the base case n = 0 .

Induction step n → n +1 Now, let

be the filtering mean at step n and Σ n be the filtering covariance. The prediction mean is of the form

The residual ˆ z n +1 is then of the form

(88)

where we used properties of the ϕ -functions, namely Lhϕ 1 ( Lh ) = ϕ 0 ( Lh ) and the commutativity ϕ 0 ( Lh ) L = Lϕ 0 ( Lh ) . With Lemma B.2, the residual covariance S n +1 and Kalman gain K n +1 are then of the form

This gives the updated mean

This proves the first half of the mean recursion:

It is left to show that

Starting from the right-hand side, we have

This concludes the proof of the mean recursion and thus shows the equivalence of the two recursions.

C Proof of Proposition 3: L-stability

We first provide definitions of L-stability and A-stability, following [27, Section 8.6].

Definition 1 (L-stability) . A one-step method is said to be L-stable if it is A-stable and, in addition, when applied to the scalar test-equation ˙ y ( t ) = λy ( t ) , λ ∈ C a complex constant with Re( λ ) < 0 , it yields y n +1 = R ( hλ ) y n , and R ( hλ ) → 0 as Re( hλ ) →-∞ .

Definition 2 (A-stability) . A one-step method is said to be A-stable if its region of absolute stability contains the whole of the left complex half-plane. That is, when applied to the scalar test-equation ˙ y ( t ) = λy ( t ) with λ ∈ C a complex constant with Re( λ ) < 0 , the method yields y n +1 = R ( hλ ) y n , and { z ∈ C : Re( z ) < 0 } ⊂ { z ∈ C : R ( z ) < 1 } .

Proof of Proposition 3. Both L-stability and A-stability directly follow from Remark 1: Since the probabilistic exponential integrator solves linear ODEs exactly its stability function is the exponential function, i.e. R ( z ) = exp( z ) . A-stability and L-stability then follow: Since C -⊂ { z : | R ( z ) | ≤ 1 } holds the method is A-stable. And since | R ( z ) | → 0 as Re( z ) →-∞ the method is L-stable.

D Experiment details

D.1 Burger's equation

Burger's equation is a semi-linear partial differential equation (PDE) of the form

with diffusion coefficient D ∈ R + . We discretize the spatial domain Ω on a finite grid and approximate the spatial derivatives with finite differences to obtain a semi-linear ODE of the form

with N -dimensional y ( t ) ∈ R N , L ∈ R N × N the finite difference approximation of the Laplace operator ∂ 2 x , and a non-linear part F .

More specifically, we consider a domain Ω = (0 , 1) , which we discretize with a grid of N = 250 equidistant locations, thus we have ∆ x = 1 /N . We consider zero-Dirichlet boundary conditions, that is, u (0 , t ) = u (1 , t ) = 0 . The discrete Laplacian is then

The non-linear part of the discretized Burger's equation results from another finite-difference approximation of the term u · ∂ x u , and is chosen as

The initial condition is chosen as

We consider an integration time-span t ∈ [0 , 1] , and choose a diffusion coefficient D = 0 . 075 .

D.2 Reaction-diffusion model

The reaction-diffusion model presented in the paper, with logistic reaction term, has been used to describe the growth and spread of biological populations [22]. It is given by a semi-linear PDE

where D ∈ R + is the diffusion coefficient and R ( u ) = u (1 -u ) is a logistic reaction term. We discretize the spatial domain Ω on a finite grid and approximate the spatial derivatives with finite differences, and obtain a semi-linear ODE of the form

with N -dimensional y ( t ) ∈ R N , L ∈ R N × N the finite difference approximation of the Laplace operator, and the reaction term R is as before but applied element-wise.

We again consider a domain Ω = (0 , 1) , which we discretize on a grid of N = 100 points. This time we consider zero-Neumann conditions, that is, ∂ x u (0 , t ) = ∂ x u (1 , t ) = 0 . Including these directly into the finite-difference discretization, the discrete Laplacian is then

The initial condition is chosen as

The discrete ODE is then solved on a time-span t ∈ [0 , 2] , and we choose a diffusion coefficient D = 0 . 25 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a probabilistic integrator for solving stiff semi-linear ODEs that decouples linear and nonlinear components for improved accuracy.  The approach uses an integrated Ornstein-Uhlenbeck process (IOUP) prior and posterior inference extends classical exponential integrator to exactly solve the linear components.  Nonlinearities are approximated via local linearization similar to the extended Kalman filter (EKF).  The authors present L-stability results and address computation through proposed approximations of the Jacobian that capture only the linear dynamics.\n\nSTRENGTHS:\nThis is a solid paper; well-written, well-motivated, nice execution.  The results aren’t groundbreaking when compared to baseline, but there are clear benefits of the proposed method, particularly when nonlinearities are less significant.  Furthermore, results are presented honestly and limitations are clearly identified.\n\nWEAKNESSES:\nThe paper fails to highlight any benefit of treating the initial value problem (IVP) as a Bayesian inference task.  The posterior distribution *should* characterize uncertainty over the solution, and thus it is possible to quantify and report confidence in the solution.  It is further possible to scrutinize the appropriateness of the IOUP as a prior.  The paper does not address this, but instead prioritizes numerical accuracy and execution time.  While these metrics are certainly important, they are not the primary motivation of probabilistic numerics.  This is not a critical flaw, but perhaps a missed opportunity to demonstrate the benefit of a Bayesian approach to a numerical method.\n\nUnsurprisingly, the proposed method is fairly sensitive to nonlinearity of the ODEs.  The experiment in Sec. 4.1 demonstrates that accuracy degrades as the quadratic term becomes more significant.  In fact, the proposed method seems to be more sensitive to this than baseline IWP methods, albeit with uniformly lower error.  For the simple ODE in Sec. 4.1 all methods achieve very low error (on the order of $10^{-10}$) so the observed improvement with the IOUP is marginal unless the ODE is highly linear.\n\nSome detailed comments:\n  * Fig. 5 should reference subfigures (e.g. (a) is never explicitly mentioned).\n  * L268: Typo \"Since th problem...\"\n\nQUESTIONS:\n* The authors suggest that probabilistic integrators may be preferable for expensive-to-evaluate functions (L259).  How?  Can the method avoid explicit evaluation via prediction?\n\n* The EKF can be quite unreliable in even simple nonlinear dynamical systems, leading to preferred approaches such as the UKF.  Would not such an approach be preferable here to avoid the local linear approximations?  What is the main challenge in a more accurate approximation?\n\nLIMITATIONS:\nLimitations are clearly described in Sec. 5", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-01"}

comments: SUMMARY:
This paper presents probabilistic exponential integrators as a solver for stiff semi-linear ODEs.Their solver can also be extended to solve general non-linear ODEs with iterative re-linearization. The proposed method is shown to be L-stable in theory and empirically more stable than existing methods.

STRENGTHS:
The paper is well written. It tackles the problem of solving stiff systems, which is a great challenge in probabilistic numerics. The authors exploit the properties in semi-linear ODEs and develop a probabilistic solver that integrates the fast, linear dynamics into the prior model. The solver is also extended to general nonlinear ODEs via iterative linearization. Both theoretical and empirical analysis are sound.

WEAKNESSES:
The main difficulty of conventional of stiff ODEs solvers is its sensitivity to step size: small step size is required for stability. Proposition 3 shows the proposed solver is L-stable. However, the experiments mainly focus on final errors across various step sizes. A detailed investigation of stability (boundedness of solutions?) behaviors of solvers under large step size may help readers better understand the proposed approach.

QUESTIONS:
See Weakness

LIMITATIONS:
See Weakness
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: SUMMARY:
This paper proposed probabilistic exponential integrators, which is a new class of probabilistic solvers for stiff semi-linear ODEs. More specifically, the integrated OU process is introduced for a functional prior that directly incorporate the linear part of the dynamics. As a result, the proposed methods can be more stable than their non-exponential probabilistic counterparts. In addition, the authors also provide an extension for general non-linear systems, which can be viewed as a probablistic solver based on exponential Rosenbrock-type methods.

STRENGTHS:
1. The paper is overall clearly written and nicely organized. 
2. The IOUP prior is neat. Also, some theoretical results of the proposed probabilistic exponential integrator have been established (e.g., equivalence to the classic exponential trapezoidal rule and L-stability).

WEAKNESSES:
1. The proposed methods are only for semi-linear ODEs, while other baselines can be applied to more general cases.
2. The IOUP prior is more expensive than the IWP prior.
3. The Kalman filter/smoother based solver has cubic scaling in the ODE dimension. 
4. It seems that the advantage of the method is more significant when the step size is large (error is large), and becomes negligible when the step size is small.

Note that 2 and 3 have been adimitted by the authors.

QUESTIONS:
1. What is the overhead computational time for computing the matrix exponentials? Is it contained in the runtime?
2. The proposed method can be more accurate when the linear part is more dominant. Since EK1 also introduces a linear approximation, would it perform better when combined with the IOUP prior?

LIMITATIONS:
Yes, they do.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-17

comments: SUMMARY:
The paper studies classical exponential integrators using the framework of probabilistic numerics, where Bayesian formalism is used to study numerical approximations of deterministic dynamical systems. The main result seems to be Proposition 1, which establishes exponential stability (called L-Stability for some reason) property for the exponential integrator.

STRENGTHS:
The paper reasonably well introduces probabilistic numerics and studies this in the context of ODEs and tackles important notion of stability of numerical approximation.

WEAKNESSES:
The paper transfers well-understood properties of numerical approximations of ODEs using the formalism of probabilistic numerics. In my view, this is not a sufficiently novel contribution.

QUESTIONS:
- Calling the Ornstein-Uhlenbeck process a Gauss-Markov in eqn 3 is overkill. 
- Notation in eqn 5 of normal distribution not clear what. Why three parameters? 
- Section 2.2 I don't see why \mathcal I [y] = 0  when y is a solution to ODE. This says that y_t-f(y_t,t)=0 not \frac{d}{dt} y_t -f(y_t,t)=0
- Explain the role of Z in 9b
- Paper seems to treat ODEs, but some numerical examples are derived for PDEs. Of course, semi-group theory for PDEs links with exponential integrators, but this seems to go beyond this paper.

LIMITATIONS:
Methodology applies to a particular class of ODES
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-24