# Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning

arXiv: https://arxiv.org/abs/2606.17513 · HTML: https://mlsanity.com/paper/2606.17513/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 57 references that are in the map.
- First posted on arXiv 16 Jun 2026.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [On the Spectral Bias of Neural Networks](https://mlsanity.com/paper/1806.08734/) (2018)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Fourier Neural Operator with Learned Deformations for PDEs on General Geometries](https://mlsanity.com/paper/2207.05209/) (2022)
- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Layer by Layer: Uncovering Hidden Representations in Language Models](https://mlsanity.com/paper/2502.02013/) (2025)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere](https://mlsanity.com/paper/2306.03838/) (2023)
- [Scalable Variational Gaussian Process Classification](https://mlsanity.com/paper/1411.2005/) (2014)
- [Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP)](https://mlsanity.com/paper/1503.01057/) (2015)
- [Uncertainty Quantification in Scientific Machine Learning: Methods, Metrics, and Comparisons](https://mlsanity.com/paper/2201.07766/) (2022)
- [Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometries](https://mlsanity.com/paper/2502.02414/) (2025)
- [Generalized Product of Experts for Automatic and Principled Fusion of Gaussian Process Predictions](https://mlsanity.com/paper/1410.7827/) (2014)
- [Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs](https://mlsanity.com/paper/2403.12553/) (2024)
- [Kernel Methods are Competitive for Operator Learning](https://mlsanity.com/paper/2304.13202/) (2023)
- [Kernel Flows: from learning kernels from data into the abyss](https://mlsanity.com/paper/1808.04475/) (2018)
- [A Simple Approach to Improve Single-Model Deep Uncertainty via Distance-Awareness](https://mlsanity.com/paper/2205.00403/) (2022)
- [Approximate Bayesian Neural Operators: Uncertainty Quantification for Parametric PDEs](https://mlsanity.com/paper/2208.01565/) (2022)
- [Active Learning for Neural PDE Solvers](https://mlsanity.com/paper/2408.01536/) (2024)
- [Error Analysis of Kernel/GP Methods for Nonlinear and Parametric PDEs](https://mlsanity.com/paper/2305.04962/) (2023)
- [Operator Learning with Gaussian Processes](https://mlsanity.com/paper/2409.04538/) (2024)
- [Probabilistic neural operators for functional uncertainty quantification](https://mlsanity.com/paper/2502.12902/) (2025)
- [GP+: A Python Library for Kernel-based learning via Gaussian Processes](https://mlsanity.com/paper/2312.07694/) (2023)

## Built on by (most cited first, 1 of 1)

- [Optimal Transport Dropout for Structured Predictive Uncertainty](https://mlsanity.com/paper/2609.33377/) (2026)

## Abstract (arXiv)

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.