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

2026-01-03

arXiv.org

AI Summary

This work proposes an efficient neural operator framework for solving stochastic differential equations (SDEs) and stochastic partial differential equations (SPDEs). The approach leverages Wiener chaos expansion (WCE) to project stochastic forcing onto a Wick–Hermite orthogonal basis, explicitly decoupling stochastic and deterministic dynamics. By parameterizing the chaos coefficients with neural operators, the method reconstructs full solution trajectories in a single forward pass. This study represents the first integration of classical WCE theory with neural operators, yielding a scalable and general-purpose deep learning solver. Extensive experiments—including SPDE benchmarks, image diffusion sampling, graph interpolation, financial extrapolation, parameter estimation, and flood forecasting—demonstrate high accuracy and strong generalization, confirming the framework’s broad applicability.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Deep Neural Architectures and Foundation Models
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• stochastic differential equations
• stochastic partial differential equations
• solution operators
• Wiener chaos expansion
• neural operators

Innovation

• Wiener chaos expansion
• neural operator
• stochastic differential equations
• Wick-Hermite basis
• solution operator learning