Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning

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.

• The paper introduces REEF-GP, which uses latent geometry-aware features of frozen neural operators to model residual discrepancies via Gaussian processes.
• It integrates deep kernel learning with spatial coordinates to achieve state-of-the-art calibration, efficient compute cost, and interpretable uncertainty localization.
• Experimental results across five PDE benchmarks demonstrate robust accuracy, reliable uncertainty quantification, and improved performance over parameter-centric approaches.

Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning: An Expert Review

Introduction and Problem Formulation

Neural operator architectures, especially their transformer-based variants, continue to yield impactful advances for surrogate modeling of parametrized PDEs. While neural operators such as Transolver deliver state-of-the-art accuracy, their deterministic predictions fundamentally limit their deployment in UQ-centric tasks and risk model misuse in scientific workflows where reliable uncertainty characterization under geometric variability is essential. Existing UQ protocols for neural operators are primarily parameter-centric, operating either via train-time Bayesian approximations (deep ensembles, MC dropout) or post-hoc localizations such as Laplace approximations in parameter or final-layer space. However, such approaches largely discard or underexploit the nontrivial geometric structure encoded in the latent representations of the neural operator itself.

"Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning" (2606.17513) rigorously challenges the parameter-centric hegemony, proposing REEF-GP (Residual on Embedded Features Gaussian Process): a framework for post-hoc UQ which systematically leverages the internal geometry-aware features of a frozen neural operator. REEF-GP eschews direct parameter perturbations, instead fitting a Gaussian process to the operator’s residual discrepancy, with the GP kernel operating in the concatenated space of coordinates and selected internal representations. This paradigm shift enables direct geometry-aware UQ without necessitating network retraining or architectural modifications.

Methodology: Geometry-Aware Discrepancy Modeling and Deep-Kernel GP Construction

The core mechanism underlying REEF-GP is rooted in the Kennedy–O'Hagan discrepancy framework, reframed for operator learning. Rather than treating the pretrained neural operator $\mathcal{G}_\theta$ as a black-box surrogate, its hidden representations are requisitioned as intrinsic coordinates defining a deep kernel GP prior for the residual discrepancy. This allows local model inadequacy to be expressed functionally as:

$$v(x) = \mathcal{G}_\theta(u, a)(x) + \delta(u, x, \mathbf{h}(u, a, x)) + \epsilon,$$

where $\mathbf{h}$ is the augmented state derived from selected internal layers. The discrepancy is further decomposed into a latent GP-modeled function and a heteroscedastic noise model. Selection of informative layers utilizes a CKA-based approach, ensuring that the concatenated feature space spans the representational hierarchy while controlling computational cost.

The latent space for the kernel is generated via spectral-normalized MLP projections, enforcing Lipschitz continuity and preventing feature collapse. The resulting kernel is constructed as the product of coordinate, function, and geometry-embedding factors, thereby inducing a nonstationary, physically-aware correlation structure. This ensures that regions of physical complexity (e.g., shock fronts) induce metric space separation, allowing uncertainty mass to localize meaningfully.

Figure 1: The REEF-GP architecture leverages the frozen neural operator's hidden geometric embeddings to build a discrepancy kernel, enabling geometry-aware UQ concentrated at physically salient error loci.

For computational tractability over datasets where the number of training points often reaches $10^6$, REEF-GP employs stochastic subset optimization during training, minimizing mini-batch negative log marginal likelihood. At inference, efficient generalized Product-of-Experts (gPoE) aggregation over multiple independent GP experts enables scalable and memory-bounded posterior computations. All hyperparameters and architectural details (e.g., kernel, subset, and expert sizes) are fixed across diverse benchmarks, further evidencing robustness.

Empirical Results: Probabilistic Calibration, Cost Analysis, and OOD Robustness

Experiments are conducted across five point-cloud PDE benchmarks (Airfoil, Pipe, Elasticity, ShapeNet Car, Ahmed Body), with strong baseline comparisons encompassing deep ensembles, MC dropout, PNO, input perturbations, Laplace approximations (LUNO-LA), and deep Vecchia ensembles.

Predictive Accuracy: REEF-GP consistently preserves or matches the deterministic neural operator's base accuracy, in contrast to stochastic train-time methods and input perturbation schemes which degrade rL2. Only deep ensembles can marginally surpass the base operator's accuracy, but at significantly elevated computational expense.

UQ Calibration: On all benchmarks, REEF-GP produces negative log-likelihood (NLL) and CRPS scores comparable to or better than deep ensembles and always superior to other post-hoc alternatives. For instance, on 3D ShapeNet Car, REEF-GP yields the lowest NLL ($2.84$) and is highly competitive on Ahmed Body. The method produces spatially meaningful uncertainty maps that concentrate at regions of physical error, as visualized in both 2D and 3D test cases.

Figure 2: Predictive standard deviation fields on Airfoil and ShapeNet Car; REEF-GP concentrates uncertainty at shock fronts (airfoil) and front bumpers (cars), corresponding to error maxima.

Compute Cost: REEF-GP achieves training overheads one order of magnitude lower than deep ensembles. Evaluation time and train memory scale comparably to other post-hoc methods, with higher, but manageable, evaluation memory due to per-expert kernel matrices.

Figure 3: Compute cost profiling demonstrates that REEF-GP hits a favorable trade-off, dominating all scalable post-hoc baselines in both speed and calibration, with a manageable memory increase due to expert subsetting.

Ablation Studies: Extensive ablations confirm that (i) the CKA-based three-layer selection matches concatenating all layers in calibration, (ii) uncertainty predicted by REEF-GP inflates consistently with operator error when training set size is reduced, (iii) the method is insensitive to both subset size and number of experts, and (iv) heteroscedastic noise modeling yields sharper, more localized uncertainty fields.

Geometric OOD Robustness: By stratifying test samples via their MMD distance from the training distribution, the authors demonstrate that REEF-GP maintains stable calibration (NLL) even as deterministic errors (rL2) grow, outperforming other post-hoc baselines that collapse under geometric shift.

Analysis of the Feature Geometry and Kernel Space

Layerwise t-SNE analyses and CKA similarity trends reveal that Transolver’s intermediate representations nontrivially deform the geometry into a latent manifold with geometric awareness at multiple scales. The MLP warping within REEF-GP further reshapes this space for UQ, yielding embeddings in which physically meaningful discontinuities (e.g., shocks) are topologically split in the kernel metric, facilitating targeted uncertainty assignment.

Figure 4: t-SNE visualizations track the deformation flow across Transolver layers, showing progressive geometry-aware encoding culminating in physically organized latent representations.

Figure 5: Feature map projections illustrate how the final kernel space teases apart regions with strong discrepancy, furnishing a geometry-consistent uncertainty prior.

Discussion: Implications and Future Work

REEF-GP substantively demonstrates that UQ in neural operator regimes should be coordinate-feature-centric. By harnessing embeddings already shaped by empirical risk minimization, REEF-GP achieves physically aware, geometry-localized uncertainty quantification unattainable by parameter-space perturbations. This positions the framework as a theoretically robust and computationally efficient alternative to deep ensembles and inducing point GPs, which become brittle or intractable under geometric variability and infinite-dimensional function spaces.

The theoretical underpinning connects recent developments in kernel flows and deep kernel learning, but REEF-GP distinguishes itself by (1) capitalizing on fixed learned feature flows and (2) sidestepping the sparsity/low-rank compromises that limit standard DKL or GP scalability. This architecture unlocks new UQ possibilities for downstream automation (adaptive experimental design, robust surrogate-based optimization) in industrial and scientific PDE applications with arbitrarily complicated geometries.

Limitations include current reliance on architectures with pointwise spatial correspondence (as in Transolver), and the restriction to stationary, single-output PDEs; extension to time-dependent, vector-valued, or spectral (FNO-class) backbones would require nontrivial adaptation. Finally, while evaluation memory scales quadratically with subset size per expert, in all reported settings, this remained tractable on modern GPUs.

Conclusion

"Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning" (2606.17513) provides an authoritative protocol for post-hoc, scalable, geometry-centric UQ in neural operators. By directly leveraging feature-coordinate spaces in GPs conditioned on latent embeddings, the method delivers robust calibration, cost-efficiency, and interpretable uncertainty fields across complex geometric and parametric variability. This work shifts the paradigm for future research towards embedding-aware, post-hoc uncertainty frameworks and invites further exploration into extending these architectures for broader operator learning classes (e.g., coupled multiphysics, time-dependent systems, and high-dimensional outputs).