Neural Stochastic Partial Differential Equations

Stochastic partial differential equations (SPDEs) are the mathematical tool of choice to model dynamical systems evolving under the influence of randomness. By formulating the search for a mild solution of an SPDE as a neural fixed-point problem, we introduce the Neural SPDE model to learn solution operators of PDEs with (possibly stochastic) forcing from partially observed data. Our model provides an extension to two classes of physics-inspired neural 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 even when the latter evolves in an infinite dimensional state space. 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 $(u_0,\xi) \mapsto u$ of SPDEs depending simultaneously on the initial condition $u_0$ and a realization of the driving noise $\xi$. A Neural SPDE is resolution-invariant, it may be trained using a memory-efficient implicit-differentiation-based backpropagation and, once trained, its evaluation is up to 3 orders of magnitude faster than traditional solvers. Experiments on various semilinear SPDEs, including the 2D stochastic Navier-Stokes equations, demonstrate how Neural SPDEs capable of learning complex spatiotemporal dynamics with better accuracy and using only a modest amount of training data compared to all alternative models.