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

TL;DR

This paper introduces a neural operator framework based on Wiener-chaos expansions for efficiently learning solutions to stochastic differential equations and stochastic partial differential equations, enabling fast and accurate trajectory reconstructions.

Contribution

It develops a novel neural operator architecture leveraging Wiener-chaos expansions, explicitly modeling chaos coefficients for SDEs and SPDEs, and demonstrates its effectiveness across multiple applications.

Findings

Achieves competitive accuracy on SPDE benchmarks

Enables single-pass trajectory reconstruction from noise

Scalable approach for diverse stochastic modeling tasks

Abstract

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 NOs to parameterize the resulting chaos coefficients, enabling reconstruction of full trajectories from noise in a single forward pass. We also make the underlying WCE structure explicit for multi-dimensional SDEs and semilinear SPDEs by showing the coupled deterministic ODE/PDE systems governing these coefficients. 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. These results suggest WCE-based neural operators are a practical and scalable approach to learning SDE/SPDE solution operators across domains.