Physics-Informed Variational State-Space Gaussian Processes

Thank you for your review. We address the points you raised below.

W1: Problem statement.
Our work contributes to the growing field of data-driven physics-informed models. These are hybrid models that aim to exploit physical inductive biases and data observations. Recently there have been principled ways to approach this problem through GPs, where the current state of the art includes AUTOIP and HELMHOLTZ-GP. However, these approaches are cubic in the number of data points limiting them to small-scale studies. In this work we provide a unifying view of such GP approaches and derive state-space algorithms that are linear in the temporal dimension, bringing computational complexity closer to standard ODE solvers, and enable the application to large scale spatio-temporal problems.

W2: Virtual observations.
Please, see Q1 below.

W3: Uncertainty quantification.
We agree that providing uncertainty metrics is valuable. In Table 1, we have already provided negative log predictive distributions (NLPDs) however we will report them and CRPS for all experiments. For the Diffusion-Reaction experiment they are

and for the Ocean Current experiment we achieve an NLPD of $-0.52$ and CRPS of $0.078$.

W4: Typo.
Thank you for pointing out this typo, we have addressed this now in the manuscript.

Q1: Uncertainty.
In this paper, we assume that the data observations and mechanistic descriptions align (although we can simply extend to handle missing physics, see review zAxE Q2). In this setting, we can view the uncertainty as representing our beliefs on 'how well we have solved the system'. Since we can only ever place a finite number of colocation points there will always be some error in our solution, and uncertainty quantification is vital in representing that. Moving forward having capable (both computationally and predictively) uncertainty-equipped models will be vital for handling misspecification (perhaps in prior or the likelihood) and for parameter learning.

Q2: Collocation points?
Please, see the discussion in reply to reviewer zAxE Q3.

Q3: GP as a limitation.
Under non-linear differential equations the true resulting posterior (of the generative model in Eqs. 5-6) will be non-Gaussian, and with highly non-linear equations (and or chaotic systems) approximating the true posterior with a Gaussian could lead to underestimation of the uncertainty (Turner et al., 2011). However, for many problems, the Gaussian assumption is effective as a Gaussian system is much simpler to solve [54,22].

• Richard Turner and Maneesh Sahani (2011). Two problems with variational expectation maximisation for time-series models. Bayesian Time Series Models.

Thank you for your review and positive view of our work. We address the points you raised below.

W1: Clarity.
To aid understanding we will provide an additional notation section and nomenclature table in the appendix as well as the proposed additions mentioned in the reply to reviewer 4ngK. To clarify, the boundary values are incorporated as observations (as defined in Eq. 6). At the inference stage this is incorporated through an additional likelihood. This means for PHYSS-VGP the likelihood term in Eq. (13) can be written as $p(Y_n | \bar{f}) = p(Y^\mathcal{O} | H_\mathcal{O} \, F_n) \, p(O | g(F_n)) \, p(Y^\mathcal{B} | H_\mathcal{B} \, F_n)$ where the differential equation is used to define $g$. We will add this discussion into Sec. 3.

W2: Novelty.
Please, see the response to reviewer X3gM W1.

Q1: Decomposition of kernels.
The primary motivation for the decomposition of the kernel into space and time is computation. This is a standard assumption for GPs and in this paper, we show that we can exploit this separability to derive state-space inference algorithms that are linear in the temporal dimension. One would expect that for linear and autonomous differential equations (time-invariant in this case) this will be a reasonable assumption, however, this is an interesting avenue for future research.

Q2: Unknown physics.
Thank you for this excellent question. This paper has focused on modelling solutions to differential equations; however, we can simply extend our framework (as in [35]) to model missing physics through GPs. This is similar to the latent force model of [66] except we now jointly model the latent force and the solution as GPs. To demonstrate this we construct 300 observations from a non-linear pendulum $\frac{d^2 \theta}{dt^2} + \sin(\theta) = 0$. We will consider the $\sin(\theta)$ term as missing, which we would like to learn. We now define $g=\frac{d^2 \, f_1}{dt^2} + f_2(t) = 0$. We achieve an RMSE of $0.068$ indicating we have recovered the latent force/unknown physics well. Going beyond this simple example is an exciting avenue for future work. We will add this discussion to the main paper with this example in the appendix.

Q3: PDE residual.
This is a common modelling assumption for collocation based methods. Ideally, we would model the residual as a noise-free observation (arising from a delta likelihood) and enforce it exactly. Treating the residual as arising from a Gaussian likelihood is a relaxation of this. We will add this discussion to the main paper.

Q4: Assumptions on $g$.
In general, for the probabilistic model to be well defined we require that g is measurable. To ensure that we are modelling something 'useful' we require that there be a unique well-defined solution to the differential equation. For ODEs, one can appeal to the Picard-Lindelöf theorem (see Thm. 36.4 in [22]) which gives general conditions for unique solutions to exist and requires that the non-linear operator be locally Lipsitz continuous. However, we are unaware of any results for general PDEs. Extending to problems that have multiple solutions would be an interesting avenue for future work. We will add this discussion to the main paper.

We thank you for your review and for raising questions that have helped improve the presentation of the paper.

We understand your concern related to the complicated notation. To aid understanding we will use the additional page in the camera-ready stage to include a notation section that will clarify the dimensions of all quantities. Additionally, we will provide a table in the appendix reporting the same. However, we would like to emphasize that throughout the paper we have attempted to use standard notation.

W1: Definition of Q/P.
Q/P are scalars denoting the number of latent functions and outputs. Q as the number of latent functions is established notation in multi-output GPs (see, e.g., [67, 37]). Additionally in the App. A.3 we have already provided an expanded form which provides the dimensions of all the relevant quantities. We will explicitly state these dimensions in Sec. 3 of the main paper.

W2: Abrupt switch in s/t.
We are unsure what you mean by this. Throughout we have been clear that we are operating in the spatio-temporal setting (explicitly stated on line 58/59) and throughout 't' refers to time and 's' to space.

W3: Collocation point definition.
The collocation points are defined in Eq. (6). Intuitively we would want the function that we learn (which we have placed a GP prior over) to coincide exactly with the solution of the differential equation at hand, i.e. that the residual between them is zero. In practice, we can only ever enforce this at a finite set of locations, and these are called the collocation points.

W4: Definition of $g$.
The function $g$ is computing the residual described in [Collocation point definition] above. It measures the (pointwise) error between the current function and the solution to the differential equation. Equivalently if $f$ followed the differential equation exactly, then, by definition, its first time derivate would equal $N_\theta(f)$ and $g(f)$ would be zero.

Q1: Definitions.
All quantities $N_s, N_t, x_{t,s}, y_{t, s}$ are defined on line 59. Again this is in accordance with the state-space GP literature (see [18]).

Q2: Multi-output.
As explained on line 79, Eq. (3) is a multi-output prior over a latent function f and its space/time derivates. Here we use established notation (see [47]) where a sample from $\bar{f}$ is of dimension $N \cdot D$ and corresponds to the multiple-outputs being stacked. We will clarify this further in the main paper.

Q3: Dimensions.
Yes, the equation is correct. As described on line 65/66 this is a $d$-dimensional vector (where $d$ is the number of time-derivates) and $\bar{f}$ is the corresponding state vector. See for example [48].

Q4: Notation.
Throughout we use $\bar{f}$ to denote $f$ together with its spatial and/or temporal derivates. On lines 66-71 we highlight that standard state-space GPs construct a state over $f$ and its time derivates, and hence use the notation $\bar{f}$. Specifically on line 66 we discuss timeseries models, and hence the state only depends on time. On lines 70/71 we discuss how spatio-temporal GPs can be represented as a state-space model, where now the state is modelling the temporal dynamics at each spatial point and hence the state is constructed over the spatial points (see [48, Sec. 12.5]). On line 78 we are constructing a GP prior over $\bar{f}$ at all the input locations $X$, as defined in Eq. (3). We will further clarify this in Sec. 2 of the main paper.

Q5: Higher-order derivatives.
The non-linear operator can be dependent on any order of spatial or temporal derivatives. The only requirement is that the latent multi-output GP is defined over those derivatives as well (which will require sufficiently smooth temporal and spatial kernels). We will add this discussion to Sec. 3 to clarify.

Q6: Dimensionality.
As discussed in (1-1) above this follows standard multi-output GP notation. Here $f_q$ are (stacked) multi-ouput GPs. At a single input location a sample will have dimension $D$ (as defined on line 73). $F_n$ is defined by stacking all outputs from the $f_q$ GPs, which for a single input location will have dimension $Q \times D$. This is then linearly mixed to create $P \times D$ outputs. To make this clearer we will write the stacking explicitly so that $F_n (X_n) = W \, \text{vec}[f_1(X_n), \cdots, f_Q(X_n)]^\top$ and state all these dimensions in Sec. 3 of the main paper.

Q7: What is $g$ doing?
The function $g: \mathbb{R}^{PD} \to \mathbb{R}$ computes the residual between the latent process and the differential equation (see W4. in 4ngK). This follows the notation set out in [22, 42, 35].

Q8: What is z?
The full magnetic field is defined in a 3-dimensional space on which we place our spatio-temporal model PHYSS-GP. Here 'z' is simply a label to denote the third dimension and can be thought of as 'depth'. We will be more explicit about this.

Thank you for your review and positive view of our work. We address the points you raised below.

W1: Typos.
Thank you for highlighting these typos, we have addressed them all. However, line 167 is not a typo; the computational complexity is $O(N (N_s \, d_s \, d)^3)$ because the expected log-likelihood decomposes across all spatio-temporal datapoints (as shown by the first term on Eqn 17), and hence is linear in $N$. This motivates the spatial-minibatching approximation where this is reduced to be $O(N_t \, (N_s \, d_s \,d)^3)$. We will add this explanation to line 167.

W2: Increasing RMSE.
This is an artefact of the stochastic nature of the optimisation algorithms used for inference. We have rerun our Non-linear Damped Pendulum experiment across 5 datasets constructed with different random seeds. We report the results below.

Q1: Uncertainty in PINNs.
PINNs have become a popular method for solving differential equations and amount to a highly complicated optimisation problem that requires specialised training regimes (see Wang et al., 2024). Current approaches to quantifying uncertainty (UQ) in PINNs are based on dropout (see Papamarkou et al.) and conformal predictions (see Podina et al., 2024). In recent years UQ for deep learning has received much attention however is limited by its computational cost (see Papamarkou et al., 2024). Exciting avenues of work could be to explore combinations of state-space algorithms and PINNs, to achieve linear time complexities but with the flexibility of PINNs. We will add this discussion to the related work section of the main paper.

• Podina et al. (2024). Conformalized physics-informed neural networks. In ICLR 2024 Workshop on AI4Differential Equations In Science.
• Papamarkou et al. (2018). Position: Bayesian deep learning is needed in the age of large-scale AI. arXiv preprint.
• Zhang et al. (2018). Quantifying total uncertainty in physics-informed neural networks for solving forward and inverse stochastic problems. Journal of Computational Physics.
• Wang et al. (2024). Respecting causality for training physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering.

Q2: Application to AUTOIP.
The derivation of the state-space algorithms hinges on the approximate posterior being represented as a posterior with likelihoods/sites that decompose across time, which is guaranteed within the natural gradient framework (as shown in Eq. 16). AUTOIP has no such guarantees since it optimises the natural parameters in Euclidean space with optimizers like Adam, and uses a whitened representation. However, if one drops the requirements of deriving a state-space algorithm, all of the approaches for reducing the spatial computational complexity in Sec. 4 can equally be applied to AUTOIP. Indeed, this is simple to do within our codebase (that will be released on publication). For example, on the Non-linear damped pendulum, we run this extension of AUTOIP with whitening and 50 inducing points for $C=1000$ and achieve an RMSE of $0.06 \pm 0.001$ and running time of $158.16 \pm 0.34$, clearly improving the running time against AUTOIP. This is still slower than our methods as it is cubic in the number of inducing derivatives/points. We will add this example to the appendix and this discussion into the main paper. The benefit of our proposed approach is that we remain linear in the temporal dimension which is vital for applications that are highly structured in time.

Thank you for your positive review. We address your two comments below.

W1: Novelty of the proposed approach.
Whilst we agree that some special cases of our framework are known, the 'beauty and originality of the proposed framework is its unifying property' (reviewer z1FQ). Not only do we provide a variational state-space framework unifying AUTOIP, HELMHOLTZ-GP, PMM, and EKS (Thm. 1), but we additionally extend this framework in three ways - Spatial-temporal derivative inducing points which are necessary for data sets that do not align to a grid and for reducing the cubic cost associated with the filtering state - A structured VI approximation in which the state-space prior is only defined over the space-time points and temporal derivatives, which significantly reduces the size of the state - Spatial minibatching which is necessary for large spatial data sets

Each of these points is novel, and when used in conjunction alleviates many of the limitations of previous methods. For example, on our ocean currents experiment, running any competitor (AUTOIP, HELMHOLTZ, PMM) is infeasible due to the sheer size of the data set ($N=42243$).

W2: Relation to existing works.
Thank you for the additional literature, they are indeed related works, and we will include them in our literature review. We briefly reiterate their connection here:

• Both Krämer and Hennig (2021) and Schmidt et al. (2021) are interesting works that are limited to the ODE setting.

• Krämer et al. (2022): This work derives a spatio-temporal PDE solver that can be encapsulated in our framework by seeing it as an extension of the EKF prior described in Ex. 3.1. As such we do not see it as an alternative to our work but a method that can be combined with ours. This could lead to interesting directions where finite differences can be used in space whilst used in conjunction with our variational approximations enabling the application to large spatio-temporal datasets. We will add this discussion to our related work and future work section.

• Duffin et al. (2021) build on the StatFEM work of Girolami et al. (2021) and Conrad et al. (2017) by deriving an extended Kalman filter where space is discretised through finite-elements. Like Krämer et al., this would be an interesting approach that could be used in conjunction with our work.