🤖 machine learning

Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning

This paper introduces REEF-GP, a scalable post-hoc uncertainty quantification framework that leverages a neural operator's intrinsic geometry-aware embeddings to fit a Gaussian Process on prediction residuals, achieving calibrated and robust uncertainty estimates across varying PDE geometries at a fraction of the cost of deep ensembles.

Original authors: Oriol Vendrell-Gallart, Nima Negarandeh, Ramin Bostanabad

Published 2026-06-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Oriol Vendrell-Gallart, Nima Negarandeh, Ramin Bostanabad

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Fast but Blind" Solver

Imagine you have a super-smart, super-fast robot that can predict how wind flows over a car, how stress moves through a bridge, or how heat spreads through a machine. This robot is a Neural Operator. It's incredibly fast and accurate, but it has a major flaw: it is deterministic.

Think of this robot as a confident chef who always serves the exact same dish. If you ask, "How will this soup taste?" the robot says, "It will taste exactly like this." It never says, "It might be a little too salty," or "If the ingredients were slightly different, it might be bland." In the real world, we need to know how sure the robot is, especially when the shape of the object (the geometry) changes in weird ways.

The Problem: Why Current "Confidence" Checks Fail

Scientists have tried to fix this by teaching the robot to be unsure while it is learning (Train-Time methods). But this is like trying to teach a chef to be unsure while they are still learning to cook. It's slow, expensive, and often makes the chef worse at actually cooking.

Other methods try to add a "confidence meter" after the robot is already trained (Post-Hoc methods). However, most of these existing meters are like generic weather reports. They say, "It might rain," but they don't look at the specific shape of the mountains or the local wind patterns. They miss the fact that uncertainty often happens in very specific, physically meaningful places (like right where a shockwave hits a wing).

The Solution: REEF-GP (The "Smart Detective")

The authors propose a new method called REEF-GP. Think of REEF-GP as a smart detective that investigates the robot's mistakes after the robot has finished its job.

Here is how it works, step-by-step:

1. The "Frozen" Robot

First, they take the fast, accurate robot and freeze it. They don't change its brain or retrain it. They just let it do its job and make a prediction.

2. The "Residual" (The Mistake)

The detective looks at the difference between what the robot predicted and what actually happened. This difference is called the residual.

Analogy: If the robot predicts a bridge will hold 100 tons, but it actually holds 90, the "residual" is that missing 10 tons.

3. Using the Robot's Own "Secret Notes"

This is the clever part. Most detectives would just look at the final mistake. But REEF-GP looks at the robot's internal notes (its hidden layers) while it was solving the problem.

Analogy: Imagine the robot is a painter. Instead of just looking at the final painting to see if it's good, REEF-GP looks at the sketches and brushstrokes the robot made while painting.
The paper claims these internal notes already contain a "map" of the geometry. The robot has already learned how the shape of the object deforms and changes. REEF-GP reuses these notes to build a custom map for uncertainty.

4. The "Geometry-Aware" Map

Because REEF-GP uses the robot's own internal notes, it knows exactly where to look for trouble.

Analogy: If the robot is predicting airflow over a wing, a generic meter might say, "Uncertainty is everywhere." But REEF-GP looks at the robot's internal notes and says, "Ah, the robot's internal map shows a tear right at the shockwave. That's where the uncertainty should be."
The result is a confidence map that lights up exactly where the physics are tricky (like shock fronts), rather than just guessing randomly.

5. The "Subset" Trick (Scalability)

Usually, calculating this kind of detailed confidence is so heavy it crashes computers (like trying to count every grain of sand on a beach).

Analogy: Instead of counting every grain of sand, REEF-GP picks a random handful of sand, studies it, and uses that to make a very good guess about the whole beach. It does this over and over again with different handfuls. This makes it fast enough to run on standard computers, even for huge 3D shapes.

The Results: Fast, Cheap, and Accurate

The paper tested REEF-GP on five different physics problems (like air flowing over a wing or water in a pipe).

It doesn't slow down the robot: Adding this confidence meter didn't make the robot's predictions worse.
It's as good as the "Gold Standard": The most accurate way to get uncertainty is to train 5 different robots and compare them (Deep Ensembles). This is very expensive. REEF-GP achieved similar accuracy to this expensive method but cost a tiny fraction of the computing power.
It handles new shapes: When the robot was tested on shapes it had never seen before (like a weirdly shaped car), REEF-GP correctly realized, "I'm not sure about this," and its uncertainty grew exactly where the errors were. Other methods got confused and gave false confidence.

Summary

REEF-GP is a way to add a "confidence meter" to a fast AI physics solver without retraining it. It works by looking at the solver's own internal "sketches" to understand the shape of the problem, allowing it to pinpoint exactly where the AI might be wrong. It's like giving a fast, confident robot a smart assistant that knows exactly where to double-check the work, saving time and money while keeping safety high.

Technical Summary: Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning

Problem Statement

Neural operators have emerged as powerful surrogates for solving partial differential equations (PDEs) across domains like fluid dynamics and solid mechanics. However, their standard deterministic nature limits their utility in scientific applications requiring rigorous uncertainty quantification (UQ), particularly when dealing with geometric variability. Existing UQ approaches generally fall into two categories:

Train-time methods: Techniques like Deep Ensembles or Bayesian approximations (e.g., MC Dropout) that model uncertainty in network parameters. These are often computationally prohibitive for large-scale operator learning on unstructured meshes.
Post-hoc methods: Approaches like Laplace approximation that add probabilistic components to a frozen model. While more efficient, many rely on local weight-space approximations or generic feature maps that fail to capture the specific geometry-aware representations learned by the operator.

A critical gap exists in leveraging the intrinsic, physics-aware latent representations of a pretrained neural operator to construct geometry-aware uncertainties without retraining the model or imposing misspecified priors in infinite-dimensional spaces.

Methodology: REEF-GP

The authors propose REEF-GP (Residual on Embedded Features Gaussian Process), a post-hoc UQ framework that models the discrepancy between a frozen neural operator and the true solution directly in the operator's learned feature space.

Core Framework

The method treats the pretrained neural operator $G_\theta$ as a deterministic computer model. The true solution $v(x)$ at a query coordinate $x$ is modeled as:
$v(x) = G_\theta(u, a)(x) + \delta(u, x, h(u, a, x)) + \epsilon$
where $\delta$ represents the model inadequacy (residual) and $\epsilon$ is observational noise. Unlike standard KOH (Kennedy–O'Hagan) frameworks that treat the model as a black box, REEF-GP exploits the internal structure of $G_\theta$ .

Key Technical Components

Geometry-Aware Deep Kernel:
- Instead of learning a feature map from scratch, REEF-GP utilizes the hidden layers of the frozen operator (specifically Transolver) which already encode deformed, physics-aware representations of the input geometry.
- It selects a subset of informative layers (early, middle, and late) based on Centered Kernel Alignment (CKA) to construct an augmented state $h(u, a, x)$ .
- A deep kernel is defined over the product of spatial coordinates, input functions, and these latent embeddings: $k_\phi(\bar{z}, \bar{z}') = k_{space} \cdot k_{fun} \cdot k_{latent}$ . The latent kernel uses a spectral-normalized MLP to warp the embeddings, creating a non-stationary prior that adapts to local physical regimes.
Probabilistic Discrepancy Modeling:
- The residual $\delta$ is modeled as a zero-mean Gaussian Process (GP) with a heteroscedastic noise term.
- The noise variance is input-dependent, parametrized by a spectral-normalized dense network to prevent feature collapse and ensure stability.
Scalable Training and Inference:
- Stochastic Subset Optimization: To handle the $O((MN)^3)$ complexity of exact GP training (where $M$ is the number of samples and $N$ is the number of discretization points), REEF-GP employs a stochastic subset strategy inspired by Kernel Flows (KF). It optimizes hyperparameters using mini-batches of residuals, decoupling memory constraints from the full mesh resolution.
- Generalized Product of Experts (gPoE): For inference, the method conditions $K$ independent GP experts on random subsets of the data. The final predictive distribution is obtained by aggregating these experts in precision space, bounding the test-time memory footprint to $O(N_s^2)$ per expert rather than the full dataset size.

Key Contributions

Geometry-Aware Deep Kernel: The design of a kernel that leverages the internal, frozen representations of a pretrained neural operator, avoiding the need to learn a separate feature map or impose a prior that may be misspecified for the operator's regularity.
Scalable Post-Hoc UQ: The development of training and inference procedures based on stochastic subset optimization and gPoE, making GP-based UQ practical for operator learning datasets containing millions of points.
Empirical Validation: Demonstration of competitive calibration across five 2D and 3D PDE benchmarks (including Airfoil, ShapeNet Car, and Ahmed Body), showing robustness under geometric distribution shifts.

Results

The paper evaluates REEF-GP against six baselines: Deep Ensembles, MC Dropout, Probabilistic Neural Operator (PNO), Input Perturbation, Laplace Approximation (LUNO-LA), and Deep Vecchia Ensemble (DVE-spatial).

Predictive Accuracy: REEF-GP preserves the predictive accuracy of the base deterministic operator (Transolver) across all datasets, with relative $L_2$ errors identical to the baseline. In contrast, train-time methods (MC Dropout, PNO) and input perturbation incur significant accuracy penalties.
Uncertainty Calibration: On 3D benchmarks (ShapeNet Car), REEF-GP achieves Negative Log-Likelihood (NLL) scores competitive with Deep Ensembles (e.g., 2.84 vs. 4.16 for Ensembles) and significantly outperforms other post-hoc methods like LUNO-LA (33.55).
Spatial Coherence: The uncertainty maps generated by REEF-GP are spatially meaningful, concentrating around physically significant regions of high error, such as shock fronts on airfoils and the front bumper of car geometries.
Geometric OOD Robustness: Under geometric distribution shifts (measured by Maximum Mean Discrepancy), REEF-GP maintains stable calibration. While other methods like LUNO-LA and Perturbation show significant NLL degradation as test geometries diverge from the training distribution, REEF-GP's NLL remains relatively stable.
Computational Cost: REEF-GP incurs a training overhead an order of magnitude lower than Deep Ensembles. While its evaluation memory is higher than some baselines due to kernel matrices, it remains within commodity GPU limits.

Significance and Claims

The authors claim that REEF-GP demonstrates that accurate and scalable post-hoc UQ for neural operators can be achieved directly in the operator's learned feature space. By reusing the geometry-aware embeddings of a frozen operator, the method offers a practical alternative to parameter-centric approaches (like Deep Ensembles) that are too expensive for large-scale problems.

The paper emphasizes that the geometry a neural operator learns while solving a PDE inherently encodes much of the information needed to quantify its own uncertainty. This approach opens a complementary axis to parameter-based UQ, providing a natural starting point for uncertainty quantification on increasingly large, geometry-aware operators without the prohibitive cost of retraining or ensemble methods. The authors note limitations regarding backbones that do not preserve pointwise spatial correspondence (e.g., standard FNOs) and the current focus on stationary, single-output problems, identifying these as natural directions for future work.