Diffusion-Based Stochastic Operator Networks for Uncertainty Quantification in Stochastic Partial Differential Equations

Abstract: We introduce a novel framework for uncertainty quantification of solution operators associated with stochastic partial differential equations (SPDEs). Although SPDEs play a central role in modeling complex physical systems under uncertainty, their practical use typically requires specifying the magnitude and structure of model uncertainties that are often unknown and difficult to infer from noisy measurements. To address this challenge, we develop a stochastic operator-learning framework that learns directly from noisy data and outputs both a mean solution field and a quantification of uncertainty. The proposed method, namely the Stochastic Operator Network (SON), is constructed by combining the structure of the Deep Operator Network (DeepONet) with Stochastic Neural Networks (SNNs) to model stochasticity and enable probabilistic prediction. The training procedure is carried out by minimizing a Hamiltonian-type loss and optimizing the resulting objective using the Stochastic Maximum Principle. Numerical experiments on benchmark SPDEs under multiple uncertainty sources demonstrate the accuracy and robustness of the proposed method in capturing solution structure and quantifying predictive uncertainty.

• The paper presents a probabilistic operator-learning framework (SON) that combines stochastic neural networks with DeepONet to quantify uncertainties in SPDEs.
• It employs a two-phase training strategy and a Hamiltonian loss based on the stochastic maximum principle to robustly capture mean solution structures and uncertainty bands.
• Empirical results on reaction-diffusion, advection-diffusion, heat, and Burgers' equations demonstrate SON’s superior accuracy and computational efficiency relative to existing methods.

Diffusion-Based Stochastic Operator Networks for Uncertainty Quantification in Stochastic Partial Differential Equations

Motivation and Problem Scope

Stochastic partial differential equations (SPDEs) are fundamental models for complex systems under uncertainty, encompassing randomness in initial/boundary conditions, coefficients, and external forcing. Quantifying predictive uncertainty is essential when dealing with noisy measurements and intrinsic stochasticity; however, traditional intrusive/non-intrusive UQ methods (e.g., polynomial chaos, stochastic finite elements) can incur prohibitive computational costs in high-dimensional, nonlinear, or time-dependent settings. Recent advances in neural operators, especially DeepONet architectures, have enabled the learning of solution operators that generalize efficiently across functional inputs. Yet, conventional neural operator methods yield deterministic predictions and are sensitive to noise, often resulting in overfitting or poor generalization in stochastic regimes. Existing UQ-enabled neural operator methods (e.g., Bayesian DeepONet, ensemble schemes, randomized priors) facilitate uncertainty estimation but often compromise computational tractability and scalability.

Methodological Contribution

The paper introduces the Stochastic Operator Network (SON), a probabilistic operator-learning framework targeting uncertainty quantification in SPDE solution operators. SON integrates stochastic neural networks (SNNs) as the branch component within Decoder-DeepONet, inheriting the universal operator approximation capabilities but augmenting deterministic neural operators via trainable diffusion mechanisms. The SNN branch network is formulated by discretizing neural ODE dynamics with an additive Brownian diffusion term, parameterized via layer-wise neural networks for both drift and diffusion. The resulting forward SDE trajectory is merged with the trunk network output (e.g., spatial-temporal evaluation points) through a flexible decoder architecture, supporting high-resolution and unaligned grids as in Decoder-DeepONet [Chen2024].

SON is trained by minimizing a Hamiltonian-type loss derived from the stochastic maximum principle (SMP), necessitating sample-wise forward SDE simulation and backward SDE adjoint computation. The optimization exploits layer-wise stochastic gradient descent with projections onto admissible parameter sets, facilitated by sample-wise backpropagation—addressing the curse of dimensionality commonly encountered in conditional expectation evaluation.

A significant innovation is the two-phase training strategy:

• Phase I: Train an expressive deterministic operator (e.g., Decoder-DeepONet) to approximate mean solution structure.
• Phase II: Feed deterministic outputs through an SNN, focusing on uncertainty modeling via stochasticity calibration and diffusion learning. This architectural decoupling enables flexible allocation of computational resources—dedicated model capacity for mean operator learning and lightweight, stable refinement networks for drift/diffusion components.

The construction addresses challenges associated with dimension-matching in operator learning propagation (especially when decoder operations modify feature map shapes) using generalized adjoint initialization for sample-wise backward SDE recursion.

Numerical Results

Empirical validation spans multiple SPDE classes:

• Reaction-diffusion equations (additive uniform noise): SON produces sample means closely matching reference solutions (maximum prediction errors $< 0.035$) and accurately calibrates output uncertainties (estimated std within $1\%$ of injected noise amplitude).
• Advection-diffusion equations (space-dependent noise in the right-hand side and velocity field): SON outperforms deterministic DeepONet and Bayesian DeepONet in mean and std predictions. In velocity field noise scenarios, SON delivers lower sample mean and std errors, robustly quantifying correlated uncertainties in the presence of spatially coupled stochasticity. Cross-sectional and domain-wide analyses confirm consistent calibration and resolution of uncertainty bands.
• Heat equations (space-time dependent random perturbations): SON accurately captures temporal dynamics of mean trajectories and uncertainty evolution, outperforming deterministic baselines in both spatial and temporal predictions.
• Two-dimensional Burgers' equations (additive noise in solution and nonlinear fluxes): SON demonstrates strong accuracy and uncertainty quantification even as stochasticity propagates nonlinearly, supporting dynamically increasing diffusion parameterization. Prediction errors in sample means and stds are consistently small across multiple test cases and temporal slices.

Comparative studies highlight that Bayesian DeepONet architectures (under equal epoch counts) fail to efficiently learn uncertainty structure and mean profiles, validating the computational and statistical efficiency of SON relative to ensemble-based surrogate and Bayesian alternatives.

Theoretical and Practical Implications

SON advances operator learning for stochastic dynamical systems by:

• Providing a principled operator-learning approach for probabilistic solution characterization in SPDEs
• Enabling efficient uncertainty calibration in high-dimensional, multiscale, and nonlinear PDE contexts
• Generalizing operator approximation theory to stochastic architectures—supported by SMP and Hamiltonian loss formulations

Practically, SON supports rapid evaluation of solution operators conditioned on deterministic input functions, yielding full predictive distributions and robust confidence bands. The two-phase modular design is adaptable to diverse operator learning tasks and can accommodate independent customization of deterministic and stochastic components. Empirical robustness across benchmark SPDEs and diverse uncertainty sources substantiates its utility for real-world quantitative UQ in scientific ML settings.

Speculation on Future Directions

Further research may focus on:

• Extension to non-Gaussian, heavy-tailed, or multiscale random phenomena via generalized diffusion models
• Integration with physics-informed constraints, data assimilation, and hybrid domain-agnostic architectures
• Scaling to multi-operator/multi-output settings and learned hierarchical uncertainty propagation in coupled physical systems
• Application to uncertainty-aware surrogate modeling in real-time control, simulation, and experimental design under noisy measurements

The methodology can guide a new class of operator networks for rigorous, scalable UQ in SPDE-driven scientific ML pipelines, supporting both post-hoc analysis and active learning in uncertainty-sensitive domains.

Conclusion

The Stochastic Operator Network (SON) framework enables efficient operator learning and uncertainty quantification for stochastic PDEs by integrating diffusion-driven stochastic neural networks with DeepONet architectures. Through a two-phase training protocol, SON achieves computational efficiency, superior predictive accuracy, and robust uncertainty calibration relative to conventional deterministic and Bayesian neural operator approaches. Comprehensive empirical results across stationary and time-dependent SPDEs substantiate SON's flexibility and reliability for UQ-aware operator learning. This framework establishes a versatile foundation for stochastic operator learning in complex, uncertainty-dominated scientific applications (2605.17107).