Expanding the Chaos: Neural Operator for Stochastic (Partial) Differential Equations

Stochastic differential equations (SDEs) and stochastic partial differential equations (SPDEs) are fundamental for modeling stochastic dynamics across the natural sciences and modern machine learning. Learning their solution operators with deep learning models promises fast solvers and new perspectives on classical learning tasks. In this work, we build on Wiener--chaos expansions (WCE) to design neural operator (NO) architectures for SDEs and SPDEs: we project driving noise paths onto orthonormal Wick--Hermite features and use NO backbones to parameterize the corresponding chaos coefficients governed by deterministic equations, enabling reconstruction of S(P)DE solutions from the observed noise trajectories. Empirically, we achieve competitive accuracy across several tasks, including standard SPDE benchmarks and SDE-based diffusion one-step image sampling, topological graph interpolation, financial extrapolation, parameter estimation, and manifold SDE flood forecasting. We further explore a law-matching setting where the Brownian trajectories paired with training solutions are unobserved, and show that the proposed framework remains effective by sampling auxiliary Brownian paths from the prescribed noise prior. These results suggest WCE-based neural operators are a practical and scalable approach to learning SDE/SPDE solution operators across domains. Our source code can be found in https://github.com/EEthanShi/SDENO.

Randomness plays an important role in many real-world systems, from fluid motion and climate-related processes to financial markets and modern generative AI. These systems are often described by stochastic differential equations, which are mathematical models for dynamics influenced by random noise. However, accurately simulating such systems can be slow, especially when many possible random outcomes need to be considered. This paper proposes a neural operator framework for learning the solution behavior of stochastic differential equations and stochastic partial differential equations. The main idea is to separate the randomness from the deterministic structure of the system. We represent the driving random noise through a set of structured features, and then train neural networks to learn how these features determine the evolution of the system. Once trained, the model can predict full solution trajectories in a single forward pass, instead of repeatedly solving the equation step by step. This framework has potential value for scientific and industrial applications where uncertainty matters. In scientific computing, it can accelerate simulations of stochastic physical systems such as fluid dynamics. In generative AI, it provides a way to view diffusion sampling through the lens of stochastic solution operators. In quantitative research and finance, where stochastic models are widely used to describe volatility, risk, and market dynamics, this type of approach may offer a faster and more flexible tool for stochastic analysis and forecasting. Experiments across several stochastic systems show that the proposed method is accurate, efficient, and broadly applicable.