\section{Introduction}

Interest in combining differential equations and deep learning has recently given rise to a large variety of neural differential equation architectures such as neural ordinary \citep{chen2018neural}, stochastic \citep{li2020scalable}, controlled \citep{kidger2020neural} and rough \citep{morrill2021neural} differential equations (Neural ODEs, SDEs, CDEs, RDEs). These models can be seen as continuous analogues of some classical neural network architectures (ResNets, RNNs etc.) and are emerging as leading machine learning tools for modelling various kinds of temporal dynamics with finite dimensional state space.

Many dynamical systems though, evolve in an infinite dimensional state space of functions and can best be described using the language of partial differential equations (PDEs). Neural Operators \citep{kovachki2021neural, li2020multipole, li2020neural, li2020fourier} are generalizations of neural networks offering an elegant way to learn mappings between spaces of functions, which makes them an attractive option for modelling PDE-dynamics. However, deterministic PDEs fail to incorporate randomness into the system they describe. In many cases the presence of noise leads to new phenomena, both at the mathematical and the physical level, generally describing more complex and realistic dynamics than the ones arising from deterministic PDEs.

\emph{Stochastic partial differential equations} (SPDEs) are the mathematical tool of choice to model many physical, biological and economic systems subject to the influence of randomness, be it intrinsic (e.g. quantifying uncertainty) or extrinsic (e.g. modelling environmental random perturbations). Examples include the \emph{Kardar–Parisi–Zhang (KPZ) equations} for random interface growth modelling for instance the propagation of a forest fire from a burnt region to an unburnt region \citep{hairer2013solving}, the \emph{$\Phi^4$-models} describing phase transitions of ferromagnets and superconductors near critical temperatures \citep{kleinert2001critical}, or the \emph{stochastic Navier-Stokes equations} modelling the dynamics of a turbulent fluid flow under the presence of local random fluctuations \citep{mikulevicius2004stochastic}. For an introduction to SPDEs see \citep{hairer2009introduction}; a comprehensive textbook is \citep{ holden1996stochastic}.

Whilst Neural Operators can certainly be used to learn operators mapping the initial condition $u_0$ or a known forcing term $\xi$ to the solution $u$ of a PDE, they are not designed to handle SPDEs where solutions depend simultaneously on an initial condition $u_0$ representing prior knowledge of the system and on a stochastic forcing $\xi$. Using the notion of mild solution to an SPDE and some basic semigroup theory, we resolve this issue and introduce the \emph{neural stochastic partial differential equation} (Neural SPDE) model. A Neural SPDE is a flexible model that can be used
\begin{enumerate}
    \item in a supervised learning setting to model the operator $(u_0,\xi) \mapsto u$ (assuming the forcing term $\xi$ is observed);
    \item as a generator in a conditional generative modelling setting to produce randomly sampled solutions matching the distributional behaviour of some observed solutions of an underlying unknown SPDE, conditioned on $u_0$.
\end{enumerate}

We perform experiments on various semilinear SPDEs with additive and multiplicative stochastic forcing, including the stochastic Navier-Stokes equations. The outcomes illustrate three key aspects of our model. One it is resolution-invariant, which means that even if trained on a lower resolution it can be directly evaluated on a higher resolution (zero-shot super-resolution). Two, it does does not require knowledge of the exact form of the underlying SPDE. Three, it is constrained to respect real physical dynamics and consequently requires only a modest  amount  of  training data,  depends  on  a significantly smaller amount of parameters and has better generalization properties compared to all alternative (neural differential equations and operators based) models.

\section{Related work}

In this section we provide a brief introduction to the most popular neural differential equation models for temporal dynamics in finite dimensions as well as to Neural Operators, a recent generalization of neural network to infinite dimensional spaces of functions. Let $T > 0$ and $d,d_u,d_\xi,d_h \in \mathbb{N}$.

\subsection{Neural ordinary differential equations}

A neural ordinary differential equation (Neural ODE), as popularised by \citep{chen2018neural}, is a parametric model with initial input $u_0 \in \mathbb{R}^{d_u}$ consisting of
\begin{itemize}
    \item a lift from the input state space $\mathbb{R}^{d_u}$ to a hidden state $\mathbb{R}^{d_h}$ (with $d_h$ chosen by the user) parameterized via a first neural network $\ell_\theta:\mathbb{R}^{d_u} \to \mathbb{R}^{d_h}$,
    \item a vector field of an ODE evolving in the hidden state modelled via a second neural network $f_\theta:\mathbb{R}^{d_h}\to\mathbb{R}^{d_h}$ (with some minimum Lipschitz regularity),
    \item and a projection from the hidden state back to the original space via a final (usually linear) readout map $\pi_\theta:\mathbb{R}^{d_h} \to \mathbb{R}^{d_u}$.
\end{itemize}
The forward pass of a Neural ODE is as follows
\begin{equation}\label{eqn:neural_ODE}
    z_0 = \ell_\theta(u_0), \quad z_t = z_0 + \int_0^tf_\theta(z_s)ds, \quad  u_t = \pi_\theta(z_t).
\end{equation}
The integral in (\ref{eqn:neural_ODE}) can be evaluated numerically via a call to an ODE solver of choice
\begin{equation*}
    z_t \approx \text{ODESolve}(z_0, f_\theta, t)
\end{equation*}
The output $u_t$ is then fed to a loss function (mean squared, cross entropy etc.) and trained via stochastic
gradient descent in the usual way. As described in \cite{chen2018neural}, backpropagation can be done efficiently via adjoint methods \cite{pontryagin1987mathematical} by solving additional ODEs for the gradients of the loss instead of backpropagating through the operations of the ODE solver. Once the parameters $\theta$ have been learnt, the solution of \cref{eqn:neural_ODE} it's determined by the initial condition $u_0$, and there
is no direct mechanism for adjusting the trajectory based on data that arrives at later times \citep{ciccone2018nais}.

\subsection{Neural controlled differential equations}

Controlled differential equations (CDEs) \citep{lyons1998differential, lyons2014rough, lyons2007differential} allow to resolve this issue by incorporating incoming information into the dynamics of an ODE. Given an initial condition $u_0 \in \mathbb{R}^d$ and some time-dependent information (usually a time series) interpolated into a continuous path $\xi: [0,T] \to \mathbb{R}^{d_\xi}$ of bounded variation, a Neural CDE \citep{kidger2020neural} consists of the same neural networks $\ell_\theta, f_\theta, \pi_\theta$ as in a Neural ODE, as well as an additional matrix-valued neural network $g_\theta : \mathbb{R}^{d_h} \to \mathbb{R}^{d_h \times d_\xi}$ (with minimal Lipschitz regularity)\footnote{In \citep{kidger2020neural} the control path $\xi$ is "augmented-with-time" to form a new path $\hat\xi_t = (t,\xi_t) \in \mathbb{R}^{d_\xi+1}$ so that the update in \cref{eqn:neural_CDE} can be rewritten more compactly as $z_t = z_0 + \int_0^th_\theta(z_s)d\hat\xi_s$, where $h_\theta : \mathbb{R}^{d_h}\to\mathbb{R}^{d_h \times (d_\xi+1)}$ is another neural network.}. The forward pass of a Neural CDE is as follows
\begin{equation}\label{eqn:neural_CDE}
    z_0 = \ell_\theta(u_0), \quad
    z_t = z_0 + \int_0^tf_\theta(z_s)ds +  \int_0^tg_\theta(z_s)d\xi_s, \quad
    u_t = \pi_\theta(z_t).
\end{equation}
where the integral is interpreted in the Riemann–Stieltjes sense. In \cite{kidger2020neural}, the control $\xi$ is obtained using cubic spline interpolation on the original data so that the term "$d\xi_s$" can be interpreted as "$\dot \xi_s ds$". Therefore \cref{eqn:neural_CDE} simply becomes an ODE and forward and backward passes can be carried out as in a Neural ODE model. In \citep{morrill2021neural} the authors discuss how to relax the differentiability assumption on $\xi$ and propose the neural rough differential equation (Neural RDE) model as an extension to Neural CDEs.

\subsection{Neural stochastic differential equations}\label{sec:neural_sde}

Neural stochastic differential equations (Neural SDEs) are special classes of Neural CDEs where the control $\xi$ is a sample path from a $\mathbb{R}^{d_\xi}$-dimensional Brownian motion and where the integral is a stochastic integral understood in the Stratonovich sense $\int_0^tg_\theta(z_s)\circ d\xi_s$. Neural SDEs can be used as generative models (trained either as VAEs or GANs) for path-valued random variables \citep{kidger2021neural, kidger2021efficient}.

Neural ODEs, CDEs, SDEs, RDEs are elegant continuous models for describing dynamics of systems where inputs $u_0, \xi$ and outputs $u$ take their values in finite dimensional spaces $\mathbb{R}^{d_u}, \mathbb{R}^{d_\xi}$. However, they can't be used in situations where the state spaces of the control and of system are infinite dimensional, as it is the case when one studies the dynamics described by partial differential equations (PDEs).

\subsection{Neural Operators}

Neural Operators \citep{kovachki2021neural, li2020multipole, li2020neural} are generalizations of neural networks offering an elegant way to model mappings between spaces of functions. Contrary to the previous differential equation models, in Neural Operators the initial condition is treated as a function $u_0:\mathcal{D}\to \mathbb{R}^d$ from a spatial domain $\mathcal{D}$ to $\mathbb{R}^d$. A Neural Operator can then be used to learn the mapping $u_0 \mapsto u$ or the mapping $\xi \mapsto u$, where $u:[0,T] \times \mathcal{D} \to \mathbb{R}^{d_u}$ is the solution of an underlying unknown PDE with initial condition $u_0:\mathcal{D}\to\mathbb{R}^{d_u}$ and spatio-temporal forcing term $\xi:[0,T] \times \mathcal{D} \to \mathbb{R}^{d_\xi}$. To work with them numerically, the user has access to $u_0$ via its pointwise evaluations on a discretization $D$ of the spatial domain $\mathcal{D}$ and to the forcing term $\xi$ through its pointwise evaluations on a discretized spatio-temporal grid $\mathcal{T}\times D$, where $\mathcal{T}=\{0=t_0<t_1<...<t_n=T\}$.

Contrary to prior work attempting to merge PDEs and deep learning, Neural Operators do not require knowledge of the exact form of the underlying PDE, they are resolution-invariant\footnote{Neural CDEs, SDEs, RDEs are time-resolution invariant but clearly not space-resolution invariant.} so even if trained on a courser grid they can be evaluated on a finer grid without giving up on accuracy, and they achieve superior performance compared to previous deep-learning-based PDE solvers.

Given an input function $u^{\text{in}}:[0,T]\times \mathcal{D} \to \mathbb{R}^{d_u}$, a Neural Operator with $M$ layers is defined as an iterative architecture such that for any $t \in [0,T]$, $x \in \mathcal{D}$, $k=1,...,M$
\begin{equation}\label{eqn:NO}
    z^0(t,x) = L_\theta(u^{\text{in}}(t,x)), \quad
    z^k = N_\theta^k(z^{k-1}), \quad
    u^{\text{out}}(t,x) = \Pi_\theta(z^M(t,x)),
\end{equation}
where $L_\theta : \mathbb{R}^{1+d_u}\to \mathbb{R}^{d_h}$ and $\Pi_\theta : \mathbb{R}^{d_h}\to \mathbb{R}^{d_u}$ are feedforward neural networks and each $N_\theta^k$ is a parametric operators defined for any function $z:[0,T]\times\mathcal{D} \to \mathbb{R}^{d_h}$ as
\begin{equation}\label{eqn:NO_layer}
    N^k_\theta(z)(t,x) := \sigma\left(A^k_\theta z(t,x) + b^k_\theta
    + \int_0^t\int_{\mathcal{D}} \mathcal{K}^k_\theta(t,s,x,y)z(s,y)dsdy \right)
\end{equation}
where $A^k_\theta \in \mathbb{R}^{d_h \times d_h}$, $b^k_\theta \in \mathbb{R}^{d_h}$, $\mathcal{K}^{k}_\theta:[0,T]^2 \times \mathcal{D}^2 \to \mathbb{R}^{d_h \times d_h}$ is a matrix-valued kernel and $\sigma$ is a non-linear activation function.

\begin{remark}
An alternative architecture (one-step-ahead Neural Operator) attempts to model the mapping $u^{\text{in}}(t_{i-1},\cdot) \to u^{\text{out}}(t_i, \cdot)$ for $i=1,...,n$, rather than modelling $u^{\text{in}} \to u^{\text{out}}$ on the whole time interval $[0,T]$. The architecture is very similar to the above, with the exception that all functions are only evaluated at a spatial input $x \in \mathcal{D}$, the kernel $\mathcal{K}^k_\theta : \mathcal{D}^2 \to \mathbb{R}^{d_h \times d_h}$ acts only on space, and the integral is only in space.
\end{remark}

\subsection{Fourier Neural Operators}\label{sec:FNO}

Among all kinds of Neural Operators \cite{kovachki2021neural, li2020multipole, li2020neural}, Fourier Neural Operators (FNO) stand out because of their easier parametrization whilst maintaining similar computational efficiency and accuracy performance for learning complex PDE operators \citep{li2020fourier}. In a FNO, the kernel $\mathcal{K}^k_\theta$ is assumed to be stationary so that $\mathcal{K}^k_\theta(t,s,x,y) = \mathcal{K}^k_\theta(t-s,x-y)$. With this assumption, the double integral in \cref{eqn:NO_layer} becomes a space-time convolution, denoted by $*_{d+1}$, so that
\begin{equation}
    N^k_\theta(z)(t,x) = \sigma\left(A^k_\theta z(t,x) + b^k_\theta + \left(\mathcal{K}^k_\theta *_{d+1} z\right)(t,x) \right)
\end{equation}
The main trick proposed \citep{li2020fourier} consists in using the convolution theorem to rewrite \cref{eqn:NO_layer} in terms of the $(d+1)$-dimensional Fourier Transform $\mathcal{F}_{d+1}$ and its inverse $\mathcal{F}_{d+1}^{-1}$ (see definitions in \Cref{appendix:FT})
\begin{equation}\label{eqn:NO_layer_Fourier}
    N^k_\theta(z)(t,x) = \sigma\Big(A^k_\theta z(t,x) + b^k_\theta + \mathcal{F}^{-1}_{d+1}\left(\mathcal{F}_{d+1}(\mathcal{K}^k_\theta)\mathcal{F}_{d+1}(z)\right)(t,x) \Big).
\end{equation}
As argued in \cite{li2020fourier}, assuming that the kernel $\mathcal{K}$ is periodic on $[0,T]\times\mathcal{D}$, instead of working with the Fourier Transform $\mathcal{F}_{d+1}(\mathcal{K})$ we can consider its Fourier series expansion and truncate it at a maximal number of frequency modes $\omega^{\text{max}}_1, \omega^{\text{max}}_2, ..., \omega^{\text{max}}_{d+1} \in \mathbb{N}$. This allows one to parameterize the evaluation $\mathcal{F}_{d+1}(\mathcal{K}^k_\theta)(\omega)$ at any frequency mode $\omega \in [0,T]\times\mathcal{D}$ directly as a complex $(2\omega^{\text{max}}_1 \times ...\times 2\omega^{\text{max}}_{d+1} \times d_h \times d_h)$-tensor $K^k_\theta$ and rewrite \cref{eqn:NO_layer_Fourier} as follows
\begin{equation}
    N^k_\theta(z)(t,x) = \sigma\left(A^k_\theta z(t,x) + b^k_\theta + \mathcal{F}^{-1}_{d+1}\left(K^k_\theta\mathcal{F}_{d+1}(z)\right)(t,x) \right),
\end{equation}
where the multiplication $K^k_\theta\mathcal{F}_{d+1}(z)$ is a matrix-vector multiplication\footnote{For any frequency mode $\omega\in[0,T]\times\mathcal{D}$ one has $\mathcal{F}_{d+1}(z)(\omega) \in \mathbb{C}^{d_h}$ and $\mathcal{F}_{d+1}(K^k_\theta)(\omega) \in \mathbb{C}^{d_h\times d_h}$.}. Assuming uniform sampling of observations across space and time, we can approximate the Fourier Transform by the Fast Fourier Transform, as done in \citet[Section 4]{li2020fourier}.

Despite the ability of FNOs to learn operators mapping either the initial condition $u_0$ or a forcing term $\xi$ to the solution $u$ of a PDE, there is, however, no natural mechanism to adapt their architecture to learn solution operators mapping the pair $(u_0,\xi)$ to the solution $u$, which is crucially needed to characterize SPDEs. Furthermore, as highlighted in the conclusion of \cite{li2020fourier}, for learning complex PDEs (such as SPDEs), FNOs require a large amount of training samples, which is hard to satisfy as data generation is generally a very expensive procedure. In addition, the number of learnable parameters in an FNO is usually very high, mainly because the weights are not shared among the different layers, making the model memory consuming and hard to train.

We note that an interesting line of work to tackle SPDE-learning is provided in the recent paper \cite{chevyrev2021feature}. The authors construct a set of features from the pair $(u_0, \xi)$ following the definition of a model from the theory of regularity structures \cite{hairer2014theory}. They then perform linear regression from this set of featuresto the solution of the SPDE at a specific point $(t,x)$ in space-time. However, contrary to all other benchmarks considered in this paper, one would have to fit a different regression function for each space-time point. For this reason we do not include this recent line of work as one of our benchmarks.

In the \Cref{sec:background} we provide a brief introduction to the theory of SPDEs which will help us defining our Neural SPDE model in \Cref{sec:model}.

\section{Background on SPDEs}\label{sec:background}

Let $\mathcal{D} \subset \mathbb{R}^d$ be a bounded domain. Let $\mathcal{H}_u=\{f:\mathcal{D}\to \mathbb{C}^{d_u}\}$ be a Hilbert space of functions on $\mathcal{D}$ with values in $\mathbb{C}^{d_u}$ and let $\mathcal{H}_\xi=\{f:\mathcal{D}\to \mathbb{C}^{d_\xi}\}$ be a Hilbert space of functions on $\mathcal{D}$ with values in $\mathbb{C}^{d_\xi}$.

We consider a large class of SPDEs of parabolic type
\begin{equation}\label{eqn:SPDE}
    du = [-\mathcal{L}u + F(u)]dt + G(u)d\xi_t
\end{equation}
where $\xi_t$ is an infinite dimensional $Q$-Wiener process \citep[Definition 10.6]{lord2014introduction} or a cylindrical Wiener process \citep[Definition 3.54]{hairer2009introduction}, $F : \mathcal{H}_u \to \mathcal{H}_u$ and $G : \mathcal{H}_u \to L(\mathcal{H}_\xi,\mathcal{H}_u)$ are two continuous functions, with $L(\mathcal{H}_\xi,\mathcal{H}_u)$ denoting the space of bounded linear operators from $\mathcal{H}_\xi$ to $\mathcal{H}_u$, and where $-\mathcal{L}$ is a linear differential operator generating a \emph{semigroup} $S_t$. For a primer on semigroup theory see \citep[Section 4]{hairer2009introduction}.

\begin{definition}\label{def:semigroup}
A strongly continuous semigroup $S$ on $\mathcal{H}_u$ is a family of bounded linear operators $S = \left\{S_t : \mathcal{H}_u \to \mathcal{H}_u\right\}_{t\geq0}$ with the properties that
\begin{enumerate}
    \item $S_0=\text{Id}$, the identity operator on $\mathcal{H}_u$,
    \item $S_t \circ S_s = S_{t+s}$, for any $s,t \geq 0$,
    \item the function $t \mapsto S_tu$ is continuous from $[0,T]$ to $\mathcal{H}_u$, for any $u \in \mathcal{H}_u$.
\end{enumerate}
\end{definition}

\medskip

A function $u:[0,T]\to\mathcal{H}_u$ is a \emph{mild solution} to the SPDE (\ref{eqn:SPDE}) if for any $t \in [0,T]$
\begin{equation}\label{eqn:mild_solution}
    u_t = S_tu_0 + \int_0^tS_{t-s}F(u_s)ds + \int_0^tS_{t-s}G(u_s)d\xi_s,
\end{equation}
where the integral is a stochastic integral \citep[Definition 3.57]{hairer2009introduction}. An SPDE can be informally thought as a SDE with infinite dimensional state space $\mathcal{H}_u$ and driven by an infinite dimensional "Brownian motion".

Assuming global Lipschitz regularity on the functions $F$ and $G$, \citep[Theorem 6.4]{hairer2009introduction} guarantees the existence and uniqueness of a solution $u:[0,T]\to\mathcal{H}_u$ to \Cref{eqn:SPDE}. The proof relies on the Banach fixed point theorem (also known as the contraction mapping theorem) and works in almost exactly the same way as the usual proof of uniqueness of solution for ODEs with Lipschitz coefficients. The central argument in that proof, which is crucial for constructing our model in the next section, shows that under these conditions, the operator $\Phi^\xi : C([0,T],\mathcal{H}_u) \to C([0,T],\mathcal{H}_u)$ depending on $\xi$ and defined for any $t \in [0,T]$ as
\begin{equation}\label{eqn:Phi}
    \Phi^\xi(u)_t := S_tu_0 + \int_0^tS_{t-s}F(u_s)ds + \int_0^tS_{t-s}G(u_s)d\xi_s
\end{equation}
is a contraction, making the solution $u$ of \cref{eqn:SPDE} a fixed point for the operator $\Phi^\xi$, satisfying the identity $u = \Phi^\xi(u)$.

\section{Neural SPDEs}\label{sec:model}

Let $m,n \in \mathbb{N}$. Let $D = \{x_1,...,x_m\} \subset \mathcal{D}$ be a $m$-points discretization of the spatial domain $\mathcal{D}$, and let $\mathcal{T} = \{t_0, ..., t_n\} \subset [0, T]$ be a $(n+1)$-points discretization of the time interval $[0, T]$ with  $t_0< ...< t_n$. Suppose the initial condition $u_0 \in \mathcal{H}_u$ is observed on $D$ and the forcing term $\xi : [0,T]\to \mathcal{H}_\xi$ is observed on $D$ at each time $t_i \in \mathcal{T}$.

Let $\xi^\epsilon = \delta^\epsilon * \xi$ be the regularisation of $\xi$ with a compactly supported smooth mollifier $\delta^\epsilon$ scaled by $\epsilon$ in space and $\epsilon^2$ in time. From now on we take $\mathcal{H}_u=L^2(\mathcal{D},\mathbb{R}^{d_u})$ and $\mathcal{H}_\xi = L^2(\mathcal{D},\mathbb{R}^{d_\xi})$. Replacing $\xi$ by $\xi^\epsilon$ in \cref{eqn:Phi} we get
\begin{align}
    \Phi^{\xi^\epsilon}(u)_t &= S_tu_0 + \int_0^tS_{t-s}F(u_s)ds + \int_0^tS_{t-s}G(u_s)\dot\xi^\epsilon_sds \nonumber \\
    &= S_tu_0 + \int_0^t S_{t-s}H^{u,\xi^\epsilon}(s,\cdot) ds, \label{eqn:Phi_H}
\end{align}
where $H^{u,\xi^\epsilon} : [0,T] \times \mathcal{D} \to \mathbb{C}^{d_u}$ is defined as
\begin{equation}\label{eqn:H}
    H^{u,\xi^\epsilon}(t,x) := \left(F(u_t) + G(u_t)\dot\xi^\epsilon_t\right)(x).
\end{equation}
Renormalization is an important component of the theory of SPDEs which is very technical and goes beyond the scope of this article. The interested reader is referred to \citep[Section 15.5]{friz2020course} for a detailed account of how this procedure works in the case of the KPZ equation. As far as the present article is concerned, the only thing to keep in mind is that for the solution $u$ to exists in the limit when $\epsilon \to 0$ one needs to introduce a (possibly infinite) normalization constant $C_\epsilon \in \mathbb{R}^{d_u}$ -- \citep[Theorem 15.2]{friz2020course} -- to the right-hand-side of \cref{eqn:Phi_H} to get our final version of $\Phi$
\begin{equation}\label{eqn:Phi_final}
    \Phi^\xi(u)_t = C + S_tu_0 + \int_0^t S_{t-s}H^{u,\xi}(s,\cdot) ds
\end{equation}
where we have dropped the dependency on $\epsilon$ to ease notation. Hence, we can assume we are able to find a fine $\epsilon$-approximation to treat the SPDE (\ref{eqn:SPDE}) are an appropriately-renormalized PDE with forcing term $\xi^\epsilon$.

\subsection{Modelling the action of the semigroup $S$}

For a large class of SPDEs, the action of the semigroup $S$ can be written as a contraction against a kernel function $\mathcal{K}_t:\mathcal{D}\times\mathcal{D} \to \mathbb{C}^{d_u \times d_u}$ such that for any $h \in \mathcal{H}_u$
\begin{equation}\label{eqn:K}
    (S_th)(x) = \int_{ \mathcal{D}}\mathcal{K}_t(x,y)h(y)\mu_t(dy)
\end{equation}
where $\mu_t$ is a Borel probability measure on $\mathcal{D}$. Here, we take $\mu_t$ to be the Lebesgue measure on $\mathbb{R}^d$ but other choices can be made, for example to speed up computations or facilitate the learning process by incorporating prior information.

Several ways to parameterize the kernel functions $\mathcal{K}$ have been introduced in the literature, for example using graph neural networks \citep{li2020neural} or using the Fourier Transform \citep{li2020fourier}. We follow the latter Fourier parameterization for the reasons mentioned in \Cref{sec:FNO}. For this, we assume that the kernel $\mathcal{K} : [0,T] \times \mathcal{D} \to \mathbb{C}^{d_u \times d_u}$ is stationary in space and time so that  $\mathcal{K}_{t,s}(x,y)=\mathcal{K}_{t-s}(x-y)$.

Hence, as observed in \Cref{sec:FNO}, equation (\ref{eqn:Phi_final}) can be rewritten in terms of the convolutions $\star_d$ and $\star_{d+1}$
\begin{equation}\label{eqn:Phi_conv}
    \Phi^\xi(u)_t = C + \mathcal{K}_t *_d u_0 + \left(\mathcal{K} *_{d+1} H^{u,\xi}\right)(t,\cdot)
\end{equation}
By the convolution theorem, \cref{eqn:Phi_conv} can be rewritten as
\begin{equation}\label{eqn:Phi_Fourier}
    \Phi^\xi(u)_t = C + \mathcal{F}^{-1}_d\left(\mathcal{F}_d(\mathcal{K}_t) \mathcal{F}_d(u_0)\right) + \mathcal{F}^{-1}_{d+1}\left(\mathcal{F}_{d+1}(\mathcal{K}) \mathcal{F}_{d+1}(H^{u,\xi})\right)(t,\cdot)
\end{equation}
where both multiplications are matrix-vector multiplication.

Hence, as done for a single layer of a FNO, we parameterize the (evaluation of the) kernel $\mathcal{K}$ directly in Fourier space as a complex tensor $K_\theta$. Doing so, \cref{eqn:Phi_Fourier} becomes
\begin{equation}\label{eqn:Phi_theta}
    \Phi^\xi_\theta(u)_t = C_\theta + \mathcal{F}^{-1}_{d}\left(\mathcal{F}_1^{-1}(K_\theta)_t \mathcal{F}_d(u_0)\right) +  \mathcal{F}_{d+1}^{-1}\left(K_\theta \mathcal{F}_{d+1}(H^{u,\xi}_\theta)\right)(t,\cdot)
\end{equation}
where we used the fact that $\mathcal{F}_d(\mathcal{K}_t) = \mathcal{F}^{-1}_1(\mathcal{F}_{d+1}(\mathcal{K}))(t)$, for any $t \in [0,T]$, and used a subscript $\theta$ to denote all elements of $\Phi^\xi_\theta$ parameterized by learnable parameters.

It remains to discuss how to parameterize the function $H^{u, \xi}_\theta$.

\subsection{Modelling the action of the vector fields $F$ and $G$}

Recall that the function $H^{u, \xi}$ is defined in \cref{eqn:H} in terms of the vector fields $F : \mathcal{H}_u\to\mathcal{H}_u$ and $G : \mathcal{H}_u\to L(\mathcal{H}_\xi, \mathcal{H}_u)$, assume to satisfy some minimal Lipsitchz conditions. For all the most common SPDEs arising from physics and finance, both $F$ and $G$ are local operators depending on $u_t$ and its first spatial derivatives $\partial u_t$, in the sense that the evaluations $F(u_t,\partial u_t)(x)$ and $G(u_t,\partial u_t)(x)$ at a point $x \in \mathcal{D}$ only depend on the evaluations $u_t(x)$ and $\partial u_t(x)$ at $x$, and not on their evaluations at some other point $y \in \mathcal{D}$ in the neighbourhood of $x$. Typical examples of $F$ and $G$ are polynomials in $u_t(x)$ and $\partial u_t(x)$.

Thanks to the properties of the Fourier Transform $\mathcal{F}_d(u_t)$, it is easy to augment $u_t$ with its first derivatives to $[u_t,\partial u_t]$
\begin{equation}
    \partial u_t = \mathcal{F}_d^{-1}(\alpha\omega\mathcal{F}_d(u_t))
\end{equation}
where $\alpha \in \mathbb{C}$ is a complex constant.

Hence, we can model the action of $F$ and $G$ as neural networks $F_\theta : \mathbb{C}^{d_u(d+1)} \to \mathbb{C}^{d_u}$ and $G : \mathbb{C}^{d_u(d+1)} \to \mathbb{C}^{d_u \times d_\xi}$ acting on $[u_t,\partial u_t](x) \in \mathbb{C}^{d_u(d+1)}$, i.e.
\begin{align*}
    F(u_t)(x) &\approx F_\theta([u_t,\partial u_t](x))\\
    (G(u_t)\dot \xi_t)(x) &\approx G_\theta([u_t,\partial u_t](x))\dot \xi_t(x)
\end{align*}
In summary
\begin{align}
    H^{u,\xi}_\theta(t,x) &= F_\theta\left(\left[u_t, \mathcal{F}_d^{-1}(\alpha\omega\mathcal{F}_d(u_t)\right](x)\right) \nonumber \\
    &+ G_\theta\left(\left[u_t, \mathcal{F}_d^{-1}(\alpha\omega\mathcal{F}_d(u_t)\right](x)\right)\dot \xi_t(x)\label{eqn:H_theta}
\end{align}

\subsection{A neural fixed point problem}

Our final Neural SPDE model consists of three components: a lift $L_\theta : \mathbb{R}^{d_u+1} \to \mathbb{R}^{d_h}$ to a hidden state $\mathbb{R}^{d_h}$, an operator $\Phi^\xi_\theta$ given by combining expressions (\ref{eqn:Phi_theta}) and (\ref{eqn:H_theta}), and a projection $\Pi_\theta : \mathbb{R}^{d_h} \to \mathbb{R}^{d_u}$ back to the orginal space $\mathbb{R}^{d_u}$.

The operator $\Phi^\xi_\theta$ is used to solve a fixed point problem $z = \Phi^\xi_\theta(z)$ whose solution $z$ is a mild solution of an SPDE. Algorithmically, we proceed by Picard's iteration
\begin{equation}\label{eqn:neural_fixed_point}
    z^i_t(x) = L_\theta(u^{\text{in}}_t(x)), \quad
    z^{i+1}_t(x) = \Phi^\xi_\theta(z^i)_t(x), \quad
    u_t^{\text{out}}(x) = \Pi_\theta(z_t^M(x))
\end{equation}
for any $t \in [0,T]$, $x \in \mathcal{D}$ and $i=0,...,M-1$.

The crucial difference between the Neural SPDE model of \cref{eqn:neural_fixed_point} and a Neural SDE model of \Cref{sec:neural_sde} is that in the former the solution $u_t$ and the noise $\xi_t$ take their values in functional spaces $\mathcal{H}_u, \mathcal{H}_\xi$ respectively, while in the latter $u_t$ and $\xi_t$ are vector-valued. We note that a Neural SDE could be used in the context of SPDE-learning by considering the evaluations of $u_t$ and $\xi_t$ on a discretized grid $D$ as state space of the SDE dynamics. However this would result in a procedure that is not resolution-invariant. Furthermore, as shown in the experimental section below, Neural SDEs in such high dimensional state spaces are very hard to train and yield very poor performance compared to all other considered models.

The crucial difference between the Neural SPDE model of \cref{eqn:neural_fixed_point} and the FNO model of \cref{eqn:NO} is between the Neural SPDE operator $\Phi_\theta$ and the FNO operator $N_\theta^M \circ ...\circ N_\theta^1$. Firstly, by construction, $\Phi_\theta$ operates on the pair $(z,\xi)$, whilst the FNO operator is built to act on $z$ alone (or possibly on $\xi$ alone) and so it can't be used to handle solution operators for SPDEs. Secondly, each FNO layer $N^i_\theta$ consists of a non-linear image of an affine transformation involving a local operation and a global operation, which doesn't represent any physical mechanism (resulting in poor performance when learning complex PDE dynamics, as shown in the next section), whilst $\Phi_\theta$ encodes real physical dynamics as it is precisely the (mild) solution operator of an SPDE. Thirdly, the number of learnable parameters in an FNO is approximately $M$ times larger than in a Neural SPDE, which is much less expensive as weights are naturally shared across iterations.

\section{Experiments}

In this section we demonstrate how our Neural SPDE model can be used in three different supervised learning settings:
\begin{enumerate}
    \item learn the operator $u^\text{in} \mapsto u^\text{out}$ (assuming the underlying noise $\xi$ is not observed),
    \item learn the operator $\xi \mapsto u^\text{out}$ (assuming the underlying noise $\xi$ is observed but the prior data $u^\text{in}$ is fixed across samples),
    \item learn the operator $(u^\text{in},\xi) \mapsto u^\text{out}$ (assuming the underlying noise $\xi$ is observed and that the prior data $u^\text{in}$ changes across samples).
\end{enumerate}
Later on in the section we showcase how the Neural SPDE model can also be used as a generator in a conditional generative modelling setting to produce randomly sampled solutions $\{u^\text{out}_i\}_{i=1}^n$ matching the distributional behaviour of some observed solutions $\{u^\text{true}_i\}_{i=1}^m \sim \mathbb{P}$ of an underlying unknown SPDE, conditioned $u^{\text{in}}_i$.

\subsection{Supervised learning}

For every supervised learning problem we use the un-normalized pathwise L2 loss. We use the Adam optimizer to train for $500$ epochs with an initial learning rate of $0.001$ that is halved every $100$ epochs. For all models the readin and readout maps as well as the vector fields are modelled as feedforward neural networks. The feedforward neural networks for the readin and the readout maps have one and two layers respectively. In the $\Phi^4_1$ example, for all models we use $d_h=16$ hidden units. For FNO and Neural SPDE we use the $32$ maximal frequency modes both in space and time. For FNO we stack four integral operator layers as in \cite{li2020fourier}. For Neural SPDE we use $4$ Picard's iterations. For the Navier-Stokes example, we use the maximal $16$ frequencies in space, the maximal $10$ frequencies in time and $d_h=8$.

\subsubsection{Dynamic $\Phi^4_1$ model }

We start with the dynamic $\Phi^4$ model in 1D which is denoted by $\Phi^4_1$ and is given by,
\begin{align}
    \partial_t u(t,x) - \Delta u(t,x) &= 3u(t,x) -u(t,x)^3 + u(t,x)\dot{\xi}(t,x), && (t,x)\in[0,T]\times\mathbb{T} \nonumber\\

    u(0,x) &= u_0(x), && x\in\mathbb{T}
\end{align}
where $\mathbb{T}$ denotes the $1$-dimensional torus, i.e. $[0,1]$ with endpoints identified ($\mathbb{T}=\mathbb{R}/\mathbb{Z}$). Here $\xi$ is a space-time white noise, and we use the approximation given in \citet[Example 10.31]{lord2014introduction} to sample realizations from this process. We then solve the SPDE for each sample path using a finite difference method with $128$ evenly distanced points in space and a time step $\Delta t=10^{-3}$. We solve the SPDE until $T=0.05$ resulting in $50$ time points. In the case where the initial data $u^{\text{in}}(t,x)=u_0(x)$ is fixed across samples, we take $u_0(x)=x(1-x)$. In the case where it varies across samples, we add a random term $u_0(x)=x(1-x)+\eta(x)$ where
\begin{align*}
    \eta(x)=0.1\sum_{k=-10}^{k=10}\frac{a_k}{1+|k|^2}\sin\left(\lambda^{-1}k\pi(x-0.5)\right), \ \ \text{with} \ \  a_k\sim\mathcal{N}(0,1)
\end{align*}
similar to \citet[Equation (3.6)]{chevyrev2021feature} with $\lambda=2$.

\begin{figure}[h]
    \centering
    \includegraphics[scale=0.23]{fig/Phi41.pdf}
    \caption{Ground truth and predictions on three test instances for the $\Phi^4_1$ example.}
    \label{fig:Phi41}
\end{figure}

We report the results in Table \ref{table:phi41}. The Neural SPDE model (NSPDE) obtains the lowest relative error for all tasks, and obtains nearly one order of magnitude lower relative
error compared to Neural CDE models (NCDE). We note that the Neural SPDE is the only model that can be applied to the three considered supervised learning settings.

\begin{table*}[h]
\caption{\textbf{Dynamic $\Phi^4_1$ model}. We report the relative L2 error on the test set. The symbol $\times$ indicates that the model is not applicable. In the first task ($u_0\mapsto u$) only partial information ($u_0$) is provided as input; the underlying noise $\xi$ is used to generate the dynamics, it varies across samples, but is not provided as an input to the models. This explains why the performance of the applicable models (FNO and NSPDE) is poorer than for all other settings. In the second setting, ($\xi\mapsto u$) the initial condition $u_0$ is kept fixed. In the third setting, the initial condition $u_0$ and the noise $\xi$ change and are both provided as input to the model. In these last two settings, even with a limited amount of training samples ($N=1\,000$), the NSPDE model is able to achieve $\sim1\%$ error rate, and even marginally improves to $<1\%$ error when $N=10\,000$. This indicates that our model is able to capture complex dynamics even in a small data regime.}
\begin{center}
\begin{tabular}{lccc|ccc}
\toprule
   \multirowcell{2}[0pt][l]{Model} &  \multicolumn{3}{c}{$N=1\,000$} &  \multicolumn{3}{c}{$N=10\,000$}\\
    \cmidrule(lr){2-7}
 & $u_0\mapsto u$ & $\xi\mapsto u$ & $(u_0,\xi)\mapsto u$ & $u_0\mapsto u$ & $\xi\mapsto u$ & $(u_0,\xi)\mapsto u$\\
\midrule
NCDE-linear & x & 0.112 & 0.127 & x & 0.056 & 0.072\\
NCDE-spline & x & 0.110 & 0.128 & x &  0.056 & 0.070\\
NRDE (depth 2) & x & x & x & x & x & x\\
FNO & 0.128 & 0.032 & x & 0.126 & 0.027 & x\\
\midrule
NSPDE (Ours)  &  0.128 & \textbf{0.009} & \textbf{0.012} & 0.126 & \textbf{0.006} & \textbf{0.006} \\
\bottomrule
\end{tabular}
\label{table:phi41}
\end{center}
\end{table*}

\subsubsection{Stochastic Navier-Stokes in 2D}

We consider the vorticity equation for an incompressible two-dimensional flow,
\begin{align}\label{eq:SNS}
    \partial_t w(t,x)- \nu \Delta w(t,x)&= - u(t,x)\cdot \nabla w(t,x) + f(x) +\sigma\dot{\xi}(t,x), && (t,x) \in[0,T]\times \mathbb{T}^2\nonumber\\

    w(0,x) &= w_0(x), && x \in \mathbb{T}^2
\end{align}
where $\mathbb{T}^2$ denotes the $2$-dimensional torus, i.e. $[0,1]^2$ with opposite sides identified ($\mathbb{T}^2=\mathbb{R}^2/\mathbb{Z}^2$). The system is subject to both deterministic and random forces. The deterministic force $f$ is a function of space only and is defined as in \cite{li2020fourier}. The additive stochastic forcing term $\xi$ depends on both time and space and is a Q-Wiener process which is colored in space and rescaled by $\sigma=0.05$. We take the Q-Wiener process defined in \citet[Example 10.12]{lord2014introduction} with $\alpha=0.005$.
The initial condition $w_0$ is generated according to $w_0\sim\mathcal{N}(0, 3^{3/2}(-\Delta+49I)^{-3})$
with periodic boundary conditions. The viscosity parameter is set to $\nu=10^{-4}$.

We use \citet[Algorithm 10.6]{lord2014introduction} to approximate sample paths from the Q-Wiener process. For each sample path we solve (\ref{eq:SNS}) with a pseudo-spectral solver (adapted from \cite{li2020fourier}) where time is advanced with a Crank–Nicolson update. We solve the SPDE on a $64\times64$ mesh in space and use a time-step of $10^{-3}$ for the Crank–Nicolson scheme. We generate $10$ long trajectories of $15\,000$ steps each  from $t=0$ to $t=15$. We partition each of these $10$ trajectories into consecutive sub-trajectories of $500$ time-steps using a rolling window. This procedure yields a total of $2\,000$ input-output pairs. We chose to split the data into shorter sequences of $500$ time steps so that one batch training could fit in memory on a Tesla P100 NVIDIA GPU. However, the model shows good performance (see Figure \ref{fig:mySNS}) even when evaluated on a $5\,000$ time step period. To showcase zero-shot super-resolution, we train the Neural SPDE model on a downsampled version of the data with $16\times16$ spatial resolution. After training on this coarser resolution we evaluate the model on the original resolution as shown in Figure \ref{fig:mySNS}. We can see that even if trained on a lower resolution the model is capable of learning the stochastic Navier-Stokes dynamics on a finer resolution.

\begin{figure}
    \centering
    \includegraphics[scale=0.32]{fig/SNS.pdf}
    \caption{\textbf{Top panel:} Solution of the vorticity equation for one realisation of the stochastic forcing between the $500^{\text{th}}$ and the $5\,000^{\text{th}}$ time steps.  \textbf{Bottom panel:} Predictions with the Neural SPDE model given the initial condition at the $500^{\text{th}}$ time step and the forcing between the $500^{\text{th}}$ and the $5\,000^{\text{th}}$ time steps. The model has been trained on lower resolution data, with a $16\times 16$ mesh and is evaluated on a $64\times 64$ mesh.}
    \label{fig:mySNS}
\end{figure}

\subsection{Generative model}

In this last section we discuss how one can use the Neural SPDE model as generator to model random solutions of an SPDE. A generative adversarial network (GAN) learns to generate new data with the same "statistics" as the data observed in the training set. A GAN has two main components: 1) a generator trying to minimize a distance between the generated and true datasets and 2) a discriminator trying to maximize this distance, hence acting adversarially.

\paragraph{Generator} Here we use a Neural SPDE as generator to map sample paths from space-time white noise $\xi$ to sample solutions $u^{\text{out}}$ conditioned on prior data $u^{\text{in}}$.

\paragraph{Discriminator} There are several variants of GAN consisting of different choices of discriminator: Wasserstein-GAN \cite{arjovsky2017wasserstein} uses the Wasserstein distance, MMD-GAN \cite{li2017mmd} uses the maximum mean discrepancy (MMD). The signature kernel $k_{\text{sig}}$ \cite{salvi2021signature} provides a natural kernel for high dimensional paths and the resulting signature-MMD
\begin{equation}
    D_{k_{\text{sig}}}(\mathbb{P}, \mathbb{Q}) = \mathbb{E}_{\mathbb{P}}[k_{\text{sig}}(x,x')] - 2 \mathbb{E}_{\mathbb{P}, \mathbb{Q}}[k_{\text{sig}}(x,x')] + \mathbb{E}_{\mathbb{Q}}[k_{\text{sig}}(y,y')]
\end{equation}
is a natural metric for comparing distributions of path space \cite{salvi2021higher}. We use the latter as discriminator to train our generative model by considering $u^{\text{in}}, u^{\text{out}}, u^{\text{true}}$ as paths\footnote{Note that this procedure is resolution-invariant. Also the complexity of computing the sig-MMD is linear in the dimensionality of the input sample paths.} evolving on $D \times \mathbb{R}^{d_u}$.

For simplicity let's assume we do not condition on prior data. So we assume we are given $n$ solution samples $\{(u^{\text{true}}_i)\}_{i=1}^n \sim \mathbb{P}_{\text{true}}$. Then we simulate $m$ space-time white noise sample $\{\xi_i\}_{i=1}^m \sim \mathbb{P}_{\text{noise}}$ and feed them to a first Neural SPDE to obtain $\mathbb{P}_{\text{out}} = \text{NSPDE}^1_\theta(\mathbb{P}_{\text{noise}})$. We then use a second Neural SPDE ($\text{NSPDE}^2_\phi$) to help the discriminator and learn by playing a two-player game where $\text{NSPDE}^2_\phi$ tries to distinguish between $\mathbb{P}_{\text{out}}$ and $\mathbb{P}_{\text{true}}$, whilst $\text{NSPDE}^1_\theta$ tries to confuse $\text{NSPDE}^2_\phi$. In other words we use the following objective loss
\begin{equation}
    \min_\theta \max_\phi D_{k_{\text{sig}}\circ \text{SPDE}_1^\phi}(\mathbb{P}_{\text{out}}, \mathbb{P}_{\text{true}})
\end{equation}

\begin{figure*}
    \centering
    \includegraphics[scale=0.12]{fig/KPZ.pdf}
    \caption{\textbf{Top panel:} Sample solutions of the KPZ equation obtained obtained with the py-pde numerical solver.  \textbf{Bottom panel:} Samples produced by the Neural SPDE-GAN.}
    \label{fig:my_label}
\end{figure*}
\subsubsection{KPZ equations}
We consider the Kardar-Parisi-Zhang (KPZ) equations,
\begin{align}
   \partial_t u(x,t) &= \nu \nabla^2u(x,t)+
            \frac{\lambda}{2} \left(\nabla u(x,t)\right)^2  + \sigma\dot{\xi}(x, t), && t \in [0,T], x \in (0,1) \nonumber\\

    u(x,0) &= u_0(x), && x \in (0,1). \nonumber
\end{align}

We use the py-pde Python package \citep{zwicker2020py} to generate solutions of the KPZ equations with $(\nu,\lambda, \sigma)=(0.5, 3, 0.2)$. We use the $32$ maximal frequencies in both space and time, $d_h=32$ hidden channels and $2$ Picard's iterations. The batch size is $128$.