Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics
Abstract
Stochastic partial differential equations (SPDEs) are the mathematical tool of choice for modelling spatiotemporal PDE-dynamics under the inﬂuence of randomness.
Based on the notion of mild solution of an SPDE, we introduce a novel neural architecture to learn solution operators of PDEs with (possibly stochastic) forcing from partially observed data.
The proposed Neural SPDE model provides an extension to two popular classes of physics-inspired architectures.
On the one hand, it extends Neural CDEs and variants – continuous-time analogues of RNNs – in that it is capable of processing incoming sequential information arriving at arbitrary spatial resolutions.
On the other hand, it extends Neural Operators – generalizations of neural networks to model mappings between spaces of functions – in that it can parameterize solution operators of SPDEs depending simultaneously on the initial condition and a realization of the driving noise.
By performing operations in the spectral domain, we show how a Neural SPDE can be evaluated in two ways, either by calling an ODE solver (emulating a spectral Galerkin scheme), or by solving a ﬁxed point problem.
Experiments on various semilinear SPDEs, including the stochastic Navier-Stokes equations, demonstrate how the Neural SPDE model is capable of learning complex spatiotemporal dynamics in a resolution-invariant way, with better accuracy and lighter training data requirements compared to alternative models, and up to 3 orders of magnitude faster than traditional solvers.
Introduction
Stochastic partial differential equations (SPDEs) are the mathematical formalism used to model many physical, biological and economic systems subject to the inﬂuence of randomness, be it intrinsic (e g. quantifying uncertainty) or extrinsic (e.g. modelling environmental random perturbations).
Notable examples of SPDEs include the Kardar–Parisi–Zhang (KPZ) equation for modelling random interface growth such as the propagation of a forest ﬁre from a burnt region to an unburnt region [12], the Ginzburg-Landau equation describing phase transitions of ferromagnets and superconductors near critical temperature [34], or the stochastic Navier-Stokes equations modelling the dynamics of a turbulent ﬂuid ﬂow under the presence of local random ﬂuctuations [31].
For an introduction to the theory of SPDEs see Hairer [11]; a comprehensive textbook is Holden et al.
Classical numerical approaches for solving SPDEs include ﬁnite difference methods and spectral
Galerkin methods [27] among others.
To ensure accuracy and stability of numerical solutions of complex SPDEs, computations must be carried out at high resolution using ﬁne discretization grids, rendering the resulting schemes computationally intractable.
This limitation motivates the study of data-driven methods that can learn solutions to differential equations from partially observed data.
Related work There has been an increased interest in recent years to combine neural networks and differential equations into a hybrid approach [5, 17, 37].
Neural controlled differential equations (Neural CDEs), as popularised by [3, 18, 32], are continuoustime analogues to recurrent neural networks (RNN, GRU, LSTM etc.).
The input to a Neural CDE model is a multivariate time series interpolated into a continuous path X : [0 , T ] → R d ξ ; the model consists of a matrix-valued feedforward neural network f θ : R d h → R d h × d ξ parameterizing the vector ﬁeld of the following dynamical system (and satisfying some minimal Lipschitz regularity to ensure existence and uniqueness of solutions)
[Equation]
where ‘ θ : R d u → R d h and π θ : R d h → R d u are feedforward neural networks.
The output response u : [0 , T ] → R d u is then fed to a (possibly pathwise) loss function (mean squared, cross entropy etc.) and trained via stochastic gradient descent in the usual way.
In practice, the term “ dX t ” means that the solution z t of the equation (e.g. attention required from a doctor) can change in response to a change of an external stream of information X t (e.g. heart rate of a patient).
Depending on the level of roughness of the control path X , the integral in eq.
(1) can be interpreted in different ways.
In [18], X is assumed differentiable and is obtained in practice via cubic splines interpolation of the original time series.
In this way, the term “ dX s ” can be interpreted as “ ˙ X s ds ” so that eq.
(1) becomes an ODE of the form ˙
z t = f θ ( z t ) ˙
X t that can be evaluated numerically via a call to an ODE solver of choice (Euler, Runge-Kutta, implicit, adaptive stepsize schemes etc.).
More generally, if X is of bounded variation then the integral above can be seen as a classical
Riemann–Stieltjes or Young integral [38].
Neural SDEs [19, 20, 22, 26] are a special subclass of
Neural CDEs where the control is a sample path from a d ξ -dimensional Brownian motion (which is not of bounded variation), and eq.
(1) is understood via stochastic integration (Itô, Stratonovich etc.).
Neural RDEs [32] allow to relax even further the regularity assumptions on X by treating the integral using rough integration [10, 30].
In practice, Neural RDEs are particularly well suited for long time series.
This is due to the fact that the model can be evaluated via a numerical scheme from stochastic analysis (called the log-ODE method [32]) over intervals much larger than what would be expected given the sampling rate of the time series.
However, the space complexity of the numerical solver increases exponentially in the number of channels d ξ , thus model complexity becomes intractable for high dimensional time series.
Despite offering many advantages for modelling temporal dynamics, these models are not designed to process signals varying both in space and in time such as physical ﬁelds described by SPDEs.
In particular, although these models are time-resolution invariant, they are not space-resolution invariant, and are not well suited to capture nonlinear interactions between the various space-time points typically observed in SPDE-dynamics.
Similar to CDEs, the solution u to an SPDE is characterized by an initial condition u 0 and a driving noise X .
However, in the case of
CDEs, ( u 0 , X t , u t ) are vectors, while in the case of SPDEs they are functions.
Neural Operators [21, 23, 24, 28] are generalizations of neural networks capable of modelling mappings between spaces of functions and offer an attractive option for learning with spatiotemporal data [21].
Among all kinds of Neural Operators, Fourier Neural Operators (FNOs)
[25] stand out because of their easier parametrization while demonstrating similar learning performance compared to other Neural Operator models.
However, Neural Operators generally fail to incorporate the effect that an external (possibly random) spatiotemporal signal might have on the system they describe.
In the case of SPDEs, the external signal is indeed random (e.g. sample from a Wiener process) and its presence leads to new phenomena, both at the mathematical and the physical level, often describing more complex and realistic dynamics than the ones arising from deterministic PDEs.
Contributions To overcome the above limitations faced by Neural CDEs and Neural Operators, we introduce the neural stochastic partial differential equation (Neural SPDE) model, capable of learning solution operators of SPDEs from partially observed data by processing, continuously in time and space, incoming sequential information arriving at an arbitrary resolution.
We propose two separate algorithms to evaluate our model: the ﬁrst reduces the Neural SPDE to a system of
ODEs in Fourier space, which can then be solved numerically by means of any ODE solver of choice (emulating a spectral Galerkin scheme); the second rewrites the Neural SPDE as a ﬁxed point problem, which is solved via classical root-ﬁnding schemes.
For both choices of evaluation, the
Neural SPDE model inherits memory-efﬁcient backpropagation capabilities provided by existing adjoint-based and implicit-differentiation-based methods respectively.
Finally, we perform extensive experiments on various semilinear SPDEs, including the stochastic Ginzburg-Landau, Korteweg-De
Vries, Navier-Stokes equations.
The empirical results illustrate several useful aspects of our model: 1) it is space and time resolution-invariant, meaning that even if trained on a lower resolution it can be directly evaluated on a higher resolution; 2) it requires a lower amount of training data to achieve similar or better performance compared to alternative models; 3) its evaluation is up to 3 orders of magnitude faster than traditional numerical solvers.
The outline of the paper is as follows: in Sec.
2 we provide a brief introduction to SPDEs which will help us to deﬁne our Neural SPDE model in Sec. 3, followed by numerical experiments in Sec. 4.
In Appendix A, we provide an overview of the computational aspects of SPDEs used to design the model and solve SPDEs numerically.
Additional experiments can be found in Appendix B.
Background on SPDEs
Let T > 0 and d, d u , d ξ ∈ N .
Let D ⊂ R d be a bounded domain.
Let H u
= { f : D → R d u } and
H ξ = { f : D → R d ξ } be two Hilbert spaces of functions from D to R d u and R d ξ respectively.
We consider a large class of SPDEs of the following type
[Equation]
where W t is either an inﬁnite dimensional Q -Wiener process [27, Def. 10.6] or a cylindrical Wiener process [11, Def. 3.54] with values in H ξ , F : H u
→ H u and G : H u
→ L ( H ξ , H u ) are two continuous operators, L ( H ξ , H u ) is the space of bounded linear operators from H ξ to H u , and L is a linear differential operator generating a semigroup 1 e t L : H u
→ H u .
For further details on Wiener processes see Appendix A.2, and for a primer on semigroup theory see Hairer [11, Section 4].
A function u : [0 , T ] → H u is said to be a mild solution of the SPDE (2) if for any t ∈ [0 , T ] it satisﬁes
[Equation]
where the second integral is a stochastic integral interpreted in the Itô sense [11, Def. 3.57].
Thus, an
SPDE can be informally thought of as an SDE with values in the functional space H u and driven by an inﬁnite dimensional Brownian motion W .
Assuming global Lipschitz regularity on F and G , a mild solution u to (2) exists and is unique [11, Thm. 6.4], at least for short times.
We follow Friz and Hairer [9] and consider a regularization W (cid:15) =
ϕ (cid:15) ∗ W of the driving noise W with a molliﬁer 2 , where ∗ means convolution.
As done in Kidger et al.
[18] for Neural CDEs, we can rewrite the mild solution of the molliﬁed version of eq.
(2) as the following randomly forced PDE
[Equation]
where ξ = ˙ W (cid:15) ,
H u
= L 2 ( D , R d u ) and H ξ =
L 2 ( D , R d ξ ) .
We will refer to ξ as white noise if W is a cylindrical Wiener process and as coloured noise if W is a Q -Wiener process.
In view of machine learning applications, one should think of W as a continuous space-time embedding of an underlying spatiotemporal data stream.
In this paper we are only going to consider W to be a sample path from a Wiener process, but we emphasise that the Neural SPDE model extends, in principle, beyond the scope of SPDEs and could be used for example to process videos in computer vision applications, which we leave as future work.
Next we introduce the Neural SPDE model
( u 0 : 𝒟 → ℝ d u , ξ : [0, T ] → L 2 (𝒟,ℝ d ξ ) )
( z 0 : 𝒟 → ℝ d h , ξ : [0, T ] → L 2 (𝒟,ℝ d ξ ) )
z : [0, T ] → L 2 (𝒟,ℝ d h ) u : [0, T ]
→ L 2 (𝒟,ℝ d u )
L θ : ℝ d u → ℝ d h Π θ : ℝ d h → ℝ d u Φ θ , ξ ( z ) t =
ℱ −1 d ( ℱ −11 ( B ) t ℱ d ( z 0 ) )
+ ℱ −1 d +
( B ℱ d +1 ( 1 ≥0 H θ , ξ ( z ⋅ ) ) )
[Figure]
Fixed Point Solver :
z = Φ θ ,
[-ODE Solver : z t = ℱ −1 d ( v t ) v t = v 0 + ∫ t 
 Av s + ℱ d ( H θ , ξ (ℱ −1 d ( v s ))), v 0 =-][Figure]
( z ) ODE Solver : z t = ℱ
[Figure]
( v t ) v
[Figure]
Av
[Figure]
[Figure] [Figure]
[Figure]
Physical Space
Latent Space
Loss ℒ( u
[Figure] [Figure]
Neural SPDEs
For a large class of differential operators L , the action of the semigroup e t L can be written as an integral against a kernel function K t : D × D → R d u × d u such that
[Equation]
for any h ∈ H u , any x ∈ D and t ∈ [0 , T ] , and where µ t is a Borel measure on D .
As in Kovachki et al.
[21], here we take µ t to be the Lebesgue measure on R d but other choices can be made, for example to incorporate prior information.
We assume that K is stationary so that eq.
(3) can be rewritten in terms of the spatial convolution ∗
[Equation]
For a large class of SPDEs of the form (2), both F and G are local operators acting on a function h ∈ H u .
In other words, the evaluations F ( h )( x ) and G ( h )( x ) at any point x ∈ D only depend h ( x ) , and not on the evaluation h ( y ) at some other point y ∈ D in the neighbourhood of x .
3.1 The model
Let H h
=
L 2 ( D , R d h ) for some latent space dimension
d h > d u .
Let
[Equation]
be four feedforward neural networks.
For any differentiable control ξ : [0 , T ] → H ξ , deﬁne the map
H θ,ξ : H u
→ H u so that for any h ∈ H u , x ∈ D , t ∈ [0 , T ]
[Equation]
A Neural SPDE is deﬁned as follows
[Equation]
We note that globally Lipschitz conditions can be imposed by using ReLU or tanh activation functions in the neural networks F θ and G θ .
In secs.
3.2 and 3.3 we propose two distinct algorithms to evaluate the Neural SPDE model (4) which are based on two different parameterization of the kernel K . a
Levereraging the convolution theorem , we can rewrite the integral in eq.
(4) as follows
[Equation]
where F , F − 1 are the d -dimensional Fourier transform (FT) and its inverse (see Def. A.1).
If one further assumes that L is a polynomial differential operator, it can be shown that there exists a map
A : C d → C d h × d h such that F ( K t )( y ) =
e tA ( y ) (see appendix B.1 for a derivation).
It follows that
[Equation]
where v t : C d → C d h is the solution of the following ODE
[Equation]
Hence, z t can be obtained by applying the inverse FT to the output of an ODE solver on [0 , t ] with initial condition F ( z 0 ) , vector ﬁeld
[Equation]
=
[Equation]
This approach can naturally be seen as a “neural version” of the classical spectral Galerkin method for SPDEs as described in Appendix A.3.
We note that this numerical evaluation of the Neural SPDE model (4) allows to inherit memoryefﬁcient adjoint-based backpropagation capabilities as in the evaluation of a Neural CDE.
For further details on adjoint-based backpropagation we refer the reader to Chen et al.
[5], Kidger [17].
3.3 Evaluating the model by solving a ﬁxed point problem
In our second approach of model evaluation, we make use of three different versions of the FT: the time-only FT F 1 and its inverse F − 1 1 , the space-only FT F d and its inverse F − 1 d , and the space-time
FT F d +1 and its inverse F − 1 d +1 (see Def.
A.1 for details).
Denoting by ? the space-time convolution, the integral in eq.
(4) can be rewritten as
[Equation]
where 1 ≥ 0 is the indicator function restricting the temporal domain to the positive real line.
Using again the convolution theorem we obtain
[Equation]
where all multiplications are matrix-vector multiplications.
Using the trick introduced in [25], one can parameterize F d +1 ( K )( y ) directly in Fourier space as a complex tensor B , so that the solution of eq.
(4) can be obtained by solving the ﬁxed point problem z
= Φ θ,ξ ( z ) with
[Equation]
and where we used the fact that F d ( K t ) =
F − 1 1 ( F d +1 ( K )) t .
This can be solved numerically using classical root-ﬁnding schemes (e.g. by Picard’s iteration)
[Equation]
Analogously to adjoint-based backpropagation for the evaluation approach mentioned in Sec. 3.2, there is a mechanism that leverages the implicit function theorem allowing to backpropagate through the operations of a ﬁxed point solver in a memory-efﬁcient way.
See Bai et al.
[2] for further details.
We note that the FTs are numerically approximated using the discrete Fourier transform (DFT) and selecting a maximum number of frequency modes 3 (see Appendix A.1 for details)
3.4 Space-time resolution-invariance
As depicted in Fig.
1, the input to a Neural SPDE corresponds to a (possibly irregularly sampled) time-indexed sequence of (possibly partially-observed) spatial observations recorded on a space-time grid.
The data is then interpolated into a continuous spatiotemporal signal ξ and initial condition u 0 .
By construction, a Neural SPDE operates in continuous time and space on the tuple of functions ( u 0 , ξ ) and produces a spatiotemporal response u , which is also continuous in space and time; the function u can then be evaluated at an arbitrary space-time resolution, possibly different from the one used during training.
In the next section we will demonstrate empirically that even if trained on a coarser resolution, a Neural SPDE can be evaluated on a ﬁner resolution without sacriﬁcing performance, a property known as zero-shot super-resolution .
3.5 Comparison of the two evaluation methods
Number of parameters For a ﬁxed dimension d h of the latent space, the majority of trainable parameters in the ODE parameterization lies in the complex tensor A , which consists of k 1 max ...
k d max d 2 h parameters, where each k i max is the maximum number of selected frequencies in the Fourier domain.
Regarding the Fixed Point parameterization, the bulk of the parameters is in the complex tensor B , which consists of k 1 max ...
k d +1 max d 2 h , where the additional frequency is due to the fact that we are taking the FFT in space-time rather than just in space as done in the ODE Solver approach.
Hence, the latter would in principle have the advantage of using a lower number of parameters than the former; however, to achieve similar performance, we found that the dimensionality of the latent space has to be roughly 20 times higher, which offsets the aforementioned advantage.
Time complexities
The time complexity of the ODE Solver approach is O ( NN x log( N x )) , where
N is the number of time steps taken by the ODE solver and N x is the number of points on the spatial grid, while the complexity of the Fixed Point approach is O ( IN x N t (log( N x ) + log( N t )))
where
I is the number of Picard iterations and N t is the number of points on the temporal grid.
In our experiments we choose N t ≈ N x and IN t ≈ N , making the two complexities comparable.
Speed of computation We found that the ODE approach is approximately 10 times slower than the Fixed Point approach.
We believe this is largely an implementation issue of the torchdiffeq library, while the FFT is a highly optimised transform in Pytorch.
3.6 Considerations about convergence
We follow Friz and Hairer [9] and consider a regularization W (cid:15) =
ϕ (cid:15) ∗ W of the driving noise W with a compactly supported smooth molliﬁer ϕ (cid:15) .
It is a classical result (Wong-Zakai [35]) from rough path theory [10, 30] that, for the case of SDEs, the sequence of random ODEs driven by the molliﬁcation of Brownian motion converges in probability to a limiting process that does not depend on the choice of molliﬁer and agrees with the Stratonovich solution of the SDE.
Furthermore, the solution map ( u 0 , W ) 7→ u is continuous in an appropriate rough path topology.
This result nicely extends to the setting of SPDEs driven by a ﬁnite dimensional noise [9, Thm. 1.3]: if u (cid:15) denotes the random
PDE solutions driven by ˙ W (cid:15) dt (instead of ◦ dW t ), then u (cid:15) converges in probability to a limiting process corresponding to the Stratonovich solution of the SPDE.
In our setting though, the driving noise is inﬁnite dimensional and the resulting integral cannot be interpreted in the Stratonovich sense because otherwise the corresponding Itô-Stratonovich correction would be inﬁnite.
Nonetheless,
Hairer and Pardoux [14, Thm. 1.1] show that, in the case of the heat operator and under appropriate renormalization and drift correction, the random PDE solution u (cid:15) converges in probability to the Itô solution of the SPDE, and that the solution map is continuous in an appropriate regularity structures topology.
We note that extending this result to a generic differential operators would require a similarly rigorous proof, which goes beyond the scope of this article and that we leave as future work.
Experiments
In this section, we run experiments on three semilinear SPDEs: the stochastic Ginzburg-Landau equation in 4.1, the stochastic Korteweg-De Vries equation in 4.2, and the stochastic Navier-Stokes equations in 4.3.
We note that although the assumption of globally Lipschitz vector ﬁelds might be violated for the following SPDEs, well-posedness (i.e. existence of global solutions) can be shown using equation-speciﬁc arguments.
We consider three supervised operator-learning settings:
• u 0 7→ u , assuming the noise ξ is not observed; • ξ 7→ u , assuming the noise ξ is observed, but the initial condition u 0 is ﬁxed across samples; • ( u 0 , ξ ) 7→ u , assuming the noise ξ is observed and u 0 changes across samples.
We note that learning the operator u 0 7→ u of an SPDE without observing the driving noise ξ unavoidably yields poor results for all considered models as only partial information about the system is provided as input.
However, we ﬁnd it informative to include the performances obtained in this setting, as this provides a sanity check that emphasizes the importance of the noise in all the experiments we consider in this paper.
Moreover, the ability to process the initial condition u 0 on its own (in absence of noise) testiﬁes that Neural SPDEs can also be used to learn deterministic PDEs.
We provide an example on the deterministic Navier-Stokes equations in Appendix B.5.
Neural CDE, Neural RDE, FNO and DeepONet [28, 29] will be the main benchmark models.
In addition, we also propose an additional baseline Neural CDE-FNO, which is a hybrid model consisting of a Neural CDE where the drift is modelled by an FNO and the diffusion by a feedforward neural network.
The motivation for using a FNO to represent the drift comes from the universal approximation properties of FNOs studied in Kovachki et al.
[21, Thm. 4].
An interesting line of work to tackle SPDE-learning is provided in [7, 16].
The authors construct a set of features from the pair ( u 0 , ξ ) following the deﬁnition of a model from the theory of regularity structures [13].
They then perform linear [7] and nonlinear [16] regression from these features to the solution of the SPDE at a single time point.
Therefore, these models would have to be retrained for any new prediction.
In addition, both [7, 16] assume knowledge of the differential operator L governing the dynamics, while Neural SPDE learns a representation of L via the parametrization of the associated kernel.
For these reasons these recent models are not included in our benchmark.
For all the experiments, the loss function is the relative pathwise L 2 error.
The hyper-parameters for all the models are selected by grid-search (see Appendix B.2 for further experimental details).
Experiments are run on a Tesla P100 NVIDIA GPU.
The code for the experiments is provided in the supplementary material.
Additional experiments may be found in Appendix B.
4.1 Stochastic Ginzburg-Landau equation
We start with the stochastic Ginzburg-Landau equation, a reaction diffusion equation in 1D given by
[Equation]
This equation is also known as the Allen-Cahn equation in 1 -dimension and is used for modeling various physical phenomena like superconductivity [34].
Here ξ denotes space-time white noise with sample paths generated using classical sampling schemes for Wiener processes detailed in A.2.
[Table]
We consider two data-regimes: a low data regime where the total number of training observations is N = 1 000 , and a large data regime where N = 10 000 .
In both cases, the response paths are generated by solving the SPDE along each sample path of the noise ξ using a ﬁnite difference scheme described in Appendix A.3 using 128 evenly distanced points in space and time and step size ∆ t = 10 − 3 .
Following the same setup as in Chevyrev et al.
[7, eq.
(3.6)], we solve the SPDE until
T = 0 .
05 resulting in 50 time points .
We choose as initial condition u 0 ( x ) =
x (1 − x )+ κη ( x ) , with η ( x ) = a
+
P k =10 k = − 10 a k / (1 + | k | 2 ) sin ( kπx ) where a k ∼ N (0 , 1) .
We take κ = 0 and κ = 0 .
1 to generate a dataset where the initial data is either ﬁxed or varies across samples.
We provide extra experiments on this SPDE for larger time horizons T and multiplicative forcing in Appendix B.3.
We report the results in Table 1.
The Neural SPDE model (NSPDE) yields the lowest relative error for all tasks, reaching one order of magnitude improvement on the main task ( u 0 , ξ ) 7→ u in the large data regime compared to all the applicable benchmark models (NCDE, NRDE, NCDE-FNO).
In all settings, even with a limited amount of training samples ( N = 1 000 ), NSPDE achieves ∼ 1% error rate, and marginally improves to < 1% error when N = 10 000 .
4.2 Stochastic Korteweg–De Vries equation
Next, we consider the stochastic Korteweg–De Vries (KdV) equation, a higher order SPDE given by ∂ t u + γ∂ 3 x u = 6 u∂ x u + ξ, u ( t, 0) =
u ( t, 1) , u (0 , x ) = u 0 ( x ) , ( t, x ) ∈ [0 , T ] × [0 , 1] .
This equation is used to describe the propagation of nonlinear waves at the surface of a ﬂuid subject to random perturbations (another wave equation is studied in Appendix B.4).
We refer the reader to Wazwaz [36] for an overview on the KdV equation and its relations to solitary waves.
The stochastic forcing is given by ξ = ˙ W for W being a partial sum approximation of a Q-Wiener process as per Example 10.8 in Lord et al.
[27] with λ j ∼ j − 5+ ε and φ j ( x ) = sin ( jπx ) (see eq.
(6) in
Appendix A.2).
Taking small ε > 0 guarantees that W t is twice differentiable in space for every t ≥ 0 .
To generate the datasets
[Table]
we solve the SPDE with γ = 0 .
1 until T = 0 .
[Table]
The stochastic forcing is simulated using 128 evenly distanced points in space and a time step ∆ t ref = 10 − 3 .
We then approximate realizations of the solution of the KdV equation using a time step ∆ t = 10 − 2 until T = 0 .
Here, the initial condition is given by u 0 ( x ) = sin(2 πx ) + κη ( x ) , where η is deﬁned as in Sec. 4.1.
Similarly to Sec.
4.1 we either take κ = 0 or κ =to generate datasets where the initial condition is either ﬁxed or varies across samples.
Each dataset consists of N = 1 000 training observations.
As reported in Table 2a,
Neural SPDEs outperforms the second best model FNO by a full order of magnitude in the task ξ 7→ u and the second best model
NCDE-FNO by almost two orders of magnitude in the task ( u 0 , ξ ) 7→ u .
We also perform the same tasks for a larger time horizon T = 1 and report the results of a comparison against FNO in Table 2b.
Neural SPDEs are able to process signals that are irregularly sampled both in space and in time by interpolating between observations.
Yet, the ability of a model to process irregular data does not guarantee its robustness when some observations are dropped.
Robustness can only be guaranteed if the signal is regular enough so that replacing dropped observations by interpolation results in a new signal that is close, in some suitable norm, to the original signal.
To illustrate this point, we run two additional experiments where we drop uniformly at random 1) 10% of the data in time and 2) 50% of the data in space.
As it can be observed in Table 2c, the performance of Neural SPDE remains roughly unchanged when data is dropped in space but decreases when data is dropped in time, which is to be expected since the driving signal is a Q-Wiener process, which is rough in time, but smoother in space.
We also note that to ensure a good approximation of the FT by the FFT, the interpolation must translate the irregular data to a (possibly ﬁner) regular grid.
4.3 Stochastic Navier-Stokes equations in 2D
Finally, we consider the vorticity form of the Navier-Stokes equations for an incompressible ﬂow
[Equation] where u is the unique divergence free ( ∇ · u = 0 ) velocity ﬁeld such that w = ∇ × u .
These equations describe the motion of an incompressible ﬂuid with viscosity ν subject to external forces [34].
The deterministic forcing f , deﬁned as in Li et al.
[25], is a function of space only.
The stochastic forcing ξ is given by ξ = ˙ W for W being a Q-Wiener process which is colored in space and rescaled by σ = 0 .
(see Appendix A.2).
The initial condition is generated according to w 0 ∼ N (0 , 3 3 / 2 ( − ∆+49 I ) − 3 ) with periodic boundary conditions.
The viscosity is set to ν = 10 − 4 .
[Figure]
For each realization of the Q-Wiener process (sampled according to the scheme in Appendix A.2)
we solve eq.
(5) with a pseudo-spectral solver described in Appendix A.3, where time is advanced with a Crank–Nicolson update.
We solve the SPDE on a 64 × 64 mesh in space and use a time step of size 10 − 3 .
For the tasks u [-
 (cid:55)→-]0 7→ u and ξ 7→ u , we generate the datasets by solving the SPDE up to time T = 1 and downsample the trajectories by a factor of 10 in time (resulting in 100 time steps) and 4 in space (resulting in a 16 × 16 spatial resolution).
The number of training samples is
N = 1 000 .
To generate the training set for the task ( u 0 , ξ ) 7→ u , we generate 10 long trajectories of 15 000 steps each up to time T = 15 .
We partition each trajectory into consecutive sub-trajectories of 500 time-steps using a rolling window.
This yields a total of 2 000 input-output pairs.
We split the data into shorter sequences of 500 time steps so that one batch ﬁts in memory of the used GPU.
[Table]
As shown in Table 3, Neural SPDEs marginally outperforms FNO on the task ξ 7→ u , but with a signiﬁcantly larger gap on the task ( u 0 , ξ ) 7→ u from the second best model NCDE-FNO.
Fig. 2 indicates that our model is capable of zero-shot super-resolution in space-time, achieving good performance even when evaluated on a larger time horizon and on an upsampled spatial grid.
Finally, we report in Table 4 some run time statistics indicating that NSPDEs can be up to 3 orders of magnitude faster than traditional numerical solvers.
Conclusion
[Table]
We introduced Neural SPDEs, a model capable of learning solution operators of PDEs with (possibly stochastic) forcing from partially observed data.
Our model provides an extension to two classes of physics-inspired models.
It extends Neural CDEs in that it is resolution-invariant both in space and in time, and it extends Neural Operators as it can be used to learn solution operators of SPDEs depending simultaneously on the initial condition and driving noise.
We performed extensive experiments illustrating how the model achieves superior performance while requiring a lower amount of training data compared to other models, and its evaluation is up to 3 orders of magnitude faster than traditional numerical solvers.
Limitations and future work Similarly to other neural operator models, parameterising the kernel in Fourier space is by no means the only available option; other parameterisations could mitigate some of the disadvantages of the FFT (irregular grids, aliasing effect ...), see for example [23].
A question we leave to future work is how to construct a discrepancy between probability measures supported on spatiotemporal signals, generalizing for example the signature kernel MMD in [33].
Neural SPDEs paired with such discrepancy would allow the design of new generative models for spatiotemporal signals.
Another research direction will be to assess whether Neural SPDEs can be used in computer vision to process videos at arbitrary resolution.
Acknowledgments and Disclosure of Funding
This project was supported by G-Research and by DataSig under the grant EP/S026347/
[-+
-]
[-( B D d +1 ( f | D ×T )) where f = 1 ≥ 
-]