MNN-BasisONet: A moment-based operator learning framework for uncertainty quantification

Abstract Neural operators extend the application of neural networks to problems in infinite-dimensional spaces and have demonstrated excellent prediction capabilities. However, to handle the inherent randomness in real-world systems, it is crucial to incorporate mechanisms for quantifying uncertainty. A major paradigm in uncertainty quantification methods is the sampling-based approach, which uses statistical sampling to quantify uncertainty but comes with a high computational cost. This paper presents MNN-BasisONet, a novel framework integrating moment neural networks (MNN) with BasisONet to address uncertainty quantification in neural operators. Our method incorporates stochastic differential equations within neurons and approximates their stochastic behavior using statistical moments, which substantially reduces computational complexity while preserving performance. MNN-BasisONet requires only a few extra parameters to quantify uncertainty in a single forward pass without sampling. The efficiency of MNN-BasisONet is demonstrated through theoretical derivations and numerical experiments, which include noiseless and noisy PDE problems, as well as real-world black-box modeling. Our model achieved comparable performance in terms of accuracy and effective uncertainty quantification compared to sampling-based methods while significantly reducing computational cost. Furthermore, we validate the MNN approach as a general uncertainty quantification method by comparing it with its stochastic counterpart and applying it to other neural operator models.