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

arXiv: https://arxiv.org/abs/2601.01021 · HTML: https://mlsanity.com/paper/2601.01021/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 76 references that are in the map.
- First posted on arXiv 3 Jan 2026.

## Builds on (most cited first, 28 of 28)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering](https://mlsanity.com/paper/1606.09375/) (2016)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Fast Sampling of Diffusion Models with Exponential Integrator](https://mlsanity.com/paper/2204.13902/) (2022)
- [Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed](https://mlsanity.com/paper/2101.02388/) (2021)
- [UNet 3+: A Full-Scale Connected UNet for Medical Image Segmentation](https://mlsanity.com/paper/2004.08790/) (2020)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [Riemannian Score-Based Generative Modelling](https://mlsanity.com/paper/2202.02763/) (2022)
- [Fast Sampling of Diffusion Models via Operator Learning](https://mlsanity.com/paper/2211.13449/) (2022)
- [Riemannian Continuous Normalizing Flows](https://mlsanity.com/paper/2006.10605/) (2020)
- [Estimating the Optimal Covariance with Imperfect Mean in Diffusion Probabilistic Models](https://mlsanity.com/paper/2206.07309/) (2022)
- [Diffusion Models for Time Series Applications: A Survey](https://mlsanity.com/paper/2305.00624/) (2023)
- [Moser Flow: Divergence-based Generative Modeling on Manifolds](https://mlsanity.com/paper/2108.08052/) (2021)
- [Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics](https://mlsanity.com/paper/2110.10249/) (2021)
- [Physics-constrained polynomial chaos expansion for scientific machine learning and uncertainty quantification](https://mlsanity.com/paper/2402.15115/) (2024)
- [Neural Operator with Regularity Structure for Modeling Dynamics Driven by SPDEs](https://mlsanity.com/paper/2204.06255/) (2022)
- [Generalized energy and gradient flow via graph framelets](https://mlsanity.com/paper/2210.04124/) (2022)
- [Polynomial Chaos Expansion for Operator Learning](https://mlsanity.com/paper/2508.20886/) (2025)
- [Deterministic Mean-field Ensemble Kalman Filtering](https://mlsanity.com/paper/1409.0628/) (2014)
- [Solving stochastic partial differential equations using neural networks in the Wiener chaos expansion](https://mlsanity.com/paper/2411.03384/) (2024)
- [Topological Schr"odinger Bridge Matching](https://mlsanity.com/paper/2504.04799/) (2025)
- [Functional SDE approximation inspired by a deep operator network architecture](https://mlsanity.com/paper/2402.03028/) (2024)

## Built on by (most cited first, 5 of 5)

- [Latent-Variable Learning of SPDEs via Wiener Chaos](https://mlsanity.com/paper/2602.11794/) (2026)
- [Wiener Chaos Expansion based Neural Operator for Singular Stochastic Partial Differential Equations](https://mlsanity.com/paper/2603.08219/) (2026)
- [Generative Neural Operators through Diffusion Last Layer](https://mlsanity.com/paper/2602.04139/) (2026)
- [Martingale Neural Operators: Learning Stochastic Marginals via Doob-Meyer Factorization](https://mlsanity.com/paper/2605.15806/) (2026)
- [S2-PINN: Stochastic Separable Physics-Informed Neural Networks](https://mlsanity.com/paper/2610.03303/) (2026)

## Abstract (arXiv)

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.