Expanding the Chaos: Neural Operator for Stochastic (Partial) Differential Equations
Dai Shi, Lequan Lin, Andi Han, Luke Thompson, José Miguel Hernández-Lobato, Zhiyong Wang, Junbin Gao
TL;DR
The paper introduces a Wiener Chaos Expansion (WCE) based neural-operator framework (NO) to learn solution operators for SPDEs and SDEs by projecting noise onto Wick–Hermite features and learning deterministic propagator components. The authors derive explicit propagator systems: ODEs for SDE chaos coefficients and PDEs for SPDE chaos coefficients, enabling one-shot trajectory reconstruction via neural operators such as FNO and UNet/ GNN backbones. Empirically, the approach, via F-SPDENO and SDENO variants, achieves competitive accuracy across SPDE benchmarks (e.g., dynamic Φ^4_1, Navier–Stokes), diffusion-based sampling in images, graph interpolation, financial extrapolation (Heston model), parameter estimation, and manifold SDEs, while preserving the one-shot evaluation advantage. This framework demonstrates scalable, discretization-invariant learning of stochastic solution operators across diverse domains and time horizons, with potential extensions to singular SPDEs and intrinsic manifolds. Overall, WCE-based neural operators provide a principled, flexible, and broadly applicable tool for fast, accurate learning of SDE/SPDE solution maps.
Abstract
Stochastic differential equations (SDEs) and stochastic partial differential equations (SPDEs) are fundamental tools for modeling stochastic dynamics across the natural sciences and modern machine learning. Developing deep learning models for approximating their solution operators promises not only fast, practical solvers, but may also inspire models that resolve classical learning tasks from a new perspective. In this work, we build on classical Wiener chaos expansions (WCE) to design neural operator (NO) architectures for SPDEs and SDEs: we project the driving noise paths onto orthonormal Wick Hermite features and parameterize the resulting deterministic chaos coefficients with neural operators, so that full solution trajectories can be reconstructed from noise in a single forward pass. On the theoretical side, we investigate the classical WCE results for the class of multi-dimensional SDEs and semilinear SPDEs considered here by explicitly writing down the associated coupled ODE/PDE systems for their chaos coefficients, which makes the separation between stochastic forcing and deterministic dynamics fully explicit and directly motivates our model designs. On the empirical side, we validate our models on a diverse suite of problems: classical SPDE benchmarks, diffusion one-step sampling on images, topological interpolation on graphs, financial extrapolation, parameter estimation, and manifold SDEs for flood prediction, demonstrating competitive accuracy and broad applicability. Overall, our results indicate that WCE-based neural operators provide a practical and scalable way to learn SDE/SPDE solution operators across diverse domains.