Structure-Aware Epistemic Uncertainty Quantification for Neural Operator PDE Surrogates

TL;DR

This work proposes a structure-aware epistemic UQ scheme that exploits the modular anatomy common to modern NOs (lifting-propagation-recovering) and yields more reliable coverage, tighter bands, and improved residual-uncertainty alignment compared with common baselines, while remaining practical in runtime.

Abstract

Neural operators (NOs) provide fast, resolution-invariant surrogates for mapping input fields to PDE solution fields, but their predictions can exhibit significant epistemic uncertainty due to finite data, imperfect optimization, and distribution shift. For practical deployment in scientific computing, uncertainty quantification (UQ) must be both computationally efficient and spatially faithful, i.e., uncertainty bands should align with the localized residual structures that matter for downstream risk management. We propose a structure-aware epistemic UQ scheme that exploits the modular anatomy common to modern NOs (lifting-propagation-recovering). Instead of applying unstructured weight perturbations (e.g., naive dropout) across the entire network, we restrict Monte Carlo sampling to a module-aligned subspace by injecting stochasticity only into the lifting module, and treat the learned solver dynamics (propagation and recovery) as deterministic. We instantiate this principle with two lightweight lifting-level perturbations, including channel-wise multiplicative feature dropout and a Gaussian feature perturbation with matched variance, followed by standard calibration to construct uncertainty bands. Experiments on challenging PDE benchmarks (including discontinuous-coefficient Darcy flow and geometry-shifted 3D car CFD surrogates) demonstrate that the proposed structure-aware design yields more reliable coverage, tighter bands, and improved residual-uncertainty alignment compared with common baselines, while remaining practical in runtime.

Figures (13)

• Figure 1: Overview of neural operator architectures and our structure-aware uncertainty quantification (UQ) scheme.
• Figure 2: Sampling from other parameter groups can lead to uninformative and low-quality samples.
• Figure 3: Visualizations of the residual and uncertainty band fields. Compared with other methods, our approach provides a more accurate characterization of the residual fields with reliable coverage.
• Figure 4: Comprehensive comparisons with dropout-based baselines. Notably, for each choice of $p$ and $T$, MC Dropout must be recalibrated since the uncertainty bands are altered by the newly introduced stochasticity, whereas our method A can share a single constant calibration scale across all $p$ and $T$ settings. Note that, to obtain the case-wise Avg. B.W., one should multiply the reported values by the grid size $N$, which typically ranges from a few hundred to tens of thousands.
• Figure 5: Uncertainty quantification (UQ) results on two challenging benchmarks.