Neural Stochastic Partial Differential Equations

Stochastic partial differential equations (SPDEs) are the mathematical tool of choice for modelling dynamical systems evolving under the influence of randomness. We introduce a novel neural architecture to learn solution operators of PDEs with (possibly stochastic) forcing from partially observed data. The proposed \emph{Neural SPDE} model provides an extension to two popular classes of physics-inspired architectures. On the one hand, it extends Neural CDEs, SDEs, RDEs -- continuous-time analogues of RNNs -- in that it is capable of processing incoming sequential information arriving at an arbitrary resolution, both in space and in time. On the other hand, it extends Neural Operators -- generalizations of neural networks to model mappings between spaces of functions -- in that it can be used to learn solution operators of SPDEs (a.k.a. It\^o maps) depending simultaneously on the initial condition and a realization of the driving noise. By transferring some of its operations to the spectral domain, we show how a Neural SPDE can be evaluated either calling an ODE solver or solving a fixed point problem, inheriting in both cases memory-efficient backpropagation capabilities for training provided by existing adjoint-based or implicit-differentiation-based methods. Experiments on various semilinear SPDEs (including stochastic Navier-Stokes) demonstrate how our model is capable of learning complex spatiotemporal dynamics with better accuracy and using only a modest amount of training data compared to all alternative models, and its evaluation is up to 3 orders of magnitude faster than traditional solvers.