Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning

2026-06-16

AI Summary

Neural operators struggle to effectively quantify uncertainty under geometric transformations, limiting their deployment in high-stakes scenarios. This work proposes REEF-GP, a framework that fits a Gaussian process on the frozen embedding space of a pre-trained neural operator to model the residual posterior. It is the first approach to directly leverage the operator’s intrinsic coordinate-feature representation for geometry-aware uncertainty estimation, eliminating the need to learn additional feature mappings. By incorporating spectral normalization projection, heteroscedastic noise modeling, and an efficient subset-based training strategy, REEF-GP achieves markedly improved stability and scalability. Evaluated on five PDE benchmarks with substantial geometric variability, REEF-GP delivers well-calibrated uncertainty estimates matching the accuracy of deep ensembles at substantially lower computational cost, while maintaining robustness under geometric distribution shifts.

Abstract

Neural operators provide fast surrogates for PDEs but their deterministic predictions limit their use in tasks requiring uncertainty quantification (UQ), especially under geometric variability. Existing approaches primarily model uncertainty in network parameters, largely overlooking the geometry-aware representations learned by the operator itself. We propose REEF-GP (Residual on Embedded Features Gaussian Process), a post-hoc UQ framework that fits a GP to the residuals of a frozen neural operator whose internal embeddings define the kernel feature space. Rather than learning a separate feature map, REEF-GP adapts the operator's intrinsic coordinate-feature representations to construct geometry-aware uncertainties. To ensure stability and scalability on unstructured domains, REEF-GP incorporates spectral-normalized projections, heteroscedastic geometry-aware noise, and efficient subset-based training that avoids restrictive low-rank approximations. Across five PDE benchmarks with varying geometries, REEF-GP preserves predictive accuracy while achieving calibrated uncertainty estimates competitive with deep ensembles but at a fraction of their cost. Our approach remains robust under geometric distribution shift, with uncertainty concentrating in physically meaningful regions (e.g., shock fronts). Our results demonstrate that accurate and scalable post-hoc UQ for neural operators can be achieved directly in their learned feature space, offering a practical alternative to parameter-centric approaches.

Problem

• uncertainty quantification
• neural operators
• geometric variability
• post-hoc
• PDEs

Innovation

• geometry-aware uncertainty
• neural operators
• post-hoc UQ
• Gaussian process
• embedded features