PaperlayerView exact revision on arXiv

Source-linked AI summary

Calibrated Uncertainty Quantification for Operator Learning via Conformal Prediction

Ziqi Ma, Kamyar Azizzadenesheli, Anima Anandkumar

arXiv:2402.01960v2cs.LG

TL;DR

Operator learning needs uncertainty estimates that cover continuous output functions simultaneously without relying on Gaussian assumptions or asymptotic guarantees. The paper proposes UQNO, a distribution-free finite-sample conformal method, and reports calibrated, efficient bands on 2D Darcy flow and 3D car pressure prediction, with only its method meeting the 98% target on the 3D task.

  • Problem

    Operator-learning uncertainty quantification lacks simultaneous function-space coverage, distribution-free calibration, and scalable guarantees for large scientific and engineering problems.

  • Method

    UQNO trains a generalized quantile neural operator for heuristic uncertainty bands, then calibrates scaled prediction sets with conformal prediction.

  • Results

    UQNO provides calibrated, efficient uncertainty bands across 2D Darcy flow and 3D car surface pressure prediction, outperforming baseline methods.

  • Takeaways & Limitations

    The method is the only evaluated approach meeting the 98% target calibration percentage on the 3D car surface pressure task.

  • Takeaways & Limitations

    The method relies on a good heuristic uncertainty estimator for tight bands, and its risk measure averages over the domain rather than targeting worst-case risk.

Abstract

from arXiv · show

Operator learning has been increasingly adopted in scientific and engineering applications, many of which require calibrated uncertainty quantification. Since the output of operator learning is a continuous function, quantifying uncertainty simultaneously at all points in the domain is challenging. Current methods consider calibration at a single point or over one scalar function or make strong assumptions such as Gaussianity. We propose a risk-controlling quantile neural operator, a distribution-free, finite-sample functional calibration conformal prediction method. We provide a theoretical calibration guarantee on the coverage rate, defined as the expected percentage of points on the function domain whose true value lies within the predicted uncertainty ball. Empirical results on a 2D Darcy flow and a 3D car surface pressure prediction task validate our theoretical results, demonstrating calibrated coverage and efficient uncertainty bands outperforming baseline methods. In particular, on the 3D problem, our method is the only one that meets the target calibration percentage (percentage of test samples for which the uncertainty estimates are calibrated) of 98%.

1 Introduction

Operator-learning uncertainty quantification must cover entire output functions, provide distribution-free finite-sample calibration, and scale to large domains. UQNO addresses these challenges with conformal calibration and achieves tight, calibrated uncertainty bands across Darcy-flow and 3D automotive pressure tasks.

  • Motivation: Function-space uncertainty quantification faces three challenges: incomplete domain coverage, missing distribution-free guarantees, and scalability.Prior approaches may focus on single points, rely on Gaussian assumptions or heuristics, and struggle with very large meshes.
  • Motivation: Simultaneous uncertainty estimates are needed across function domains because single-point or aggregate measures can miss structured errors over surfaces.The motivation includes automotive surface-pressure prediction, where whole-surface uncertainty can reveal structured error.
  • Method: UQNO combines a generalized quantile neural operator with conformal prediction to produce distribution-free, finite-sample function-space calibration.The method calibrates heuristic uncertainty bands into prediction sets with a PAC coverage guarantee.
  • Results: 98% target calibration is met by UQNO on the 3D car surface pressure task, where it is the only evaluated method to satisfy the target.The method also outperforms baselines on 2D Darcy flow, with tighter bands at the stated calibration target.
  • Results: On 2D Darcy flow, UQNO bands are 1.52x tighter than MCDropout and 76.1x tighter than Laplace approximation while satisfying 98% target calibration.The experiments cover data-rich and data-scarce regimes, with per-function point samples reaching 177k.
  • Contributions: The framework provides calibrated estimates simultaneously over all domain points and a PAC bound on their coverage percentage.Its empirical evaluation targets both calibration percentage and band tightness.

2 Problem Formulation

The formulation represents operator-learning uncertainty as a calibrated set-valued mapping over every domain point. It uses pointwise output-space balls whose radii are scaled from heuristic uncertainty estimates and calibrated by a risk-controlling PAC criterion.

  • 2.1 Uncertainty Quantification for Operator Learning: Operator learning maps an input function a in A to an output function u in U over a compact domain D.The learned operator minimizes expected loss, while a point estimate alone does not quantify predictive uncertainty.
  • 2.1 Uncertainty Quantification for Operator Learning: Uncertainty quantification replaces a point estimate with an uncertainty mapping C that assigns each input function and domain point a prediction set.For each a, C(a) maps x in D to a set in the output co-domain.
  • 2.2 Risk-Controlling Prediction Set for Operator Learning: The risk-controlling prediction set requires that the probability of excessive violating-point coverage risk is at most δ.The criterion controls the expected percentage of domain points whose true values fall outside the predicted uncertainty sets.
  • 2.3 Parameterizing Uncertainty Set: The formulation uses pointwise balls in the output co-domain, centered at the base operator prediction and scaled by a nonnegative uncertainty radius.This parameterization converts function-space uncertainty into calibrated pointwise set radii.
  • 2.3 Parameterizing Uncertainty Set: The uncertainty estimator E is intentionally heuristic before calibration, using a residual neural operator for flexibility and discretization convergence.Calibration is applied to the scaled heuristic sets rather than assuming the initial uncertainty estimates are already valid.
  • 2.3 Parameterizing Uncertainty Set: For input a and location x, Cλ(a)(x) contains output vectors within 2-norm distance λE(a)(x) of the base prediction.E(a)(x) supplies the heuristic radius and λ supplies the calibration scale.

3 Methods

The method combines generalized quantile operator learning with split conformal prediction to produce distribution-free, functionally calibrated uncertainty balls over all domain points. It supports heteroscedastic and irregularly discretized inputs while providing a conservative finite-sample calibration guarantee.

  • The method combines generalized quantile loss with conformal prediction to construct a risk-controlling quantile neural operator and calibrated prediction set.The quantile operator estimates a heuristic uncertainty band, while conformal prediction selects the scaling parameter λ for calibration.
  • Split conformal prediction uses an exchangeable calibration set to obtain a finite-sample, distribution-free confidence set without requiring training data to share the test distribution.Split conformal prediction avoids recalibration for each new test point.
  • The base uncertainty estimator uses neural operators, enabling heteroscedastic uncertainty estimates for inputs dissimilar to the training data.FNO is used for 2D Darcy flow, while a geometric-informed neural operator is used for 3D pressure prediction.
  • The algorithm trains the quantile neural operator on a separate dataset, calibrates nonconformity scores, and predicts uncertainty balls B(G(a)(x), λ̂E(a)(x)) for any domain point.The resulting prediction is defined over a test input function and can be queried pointwise across the domain.
  • The calibration guarantee is an upper bound on risk, so actual coverage is no lower than the target; sparser domain sampling makes the uncertainty band more conservative.The formulation also accommodates different discretizations and irregular grid geometries as exchangeable approximations of continuous functions.

4 Experiments

Experiments compare UQNO with MCDropout, Laplace approximation, and neural posterior principal components on 2D Darcy flow and 3D car pressure prediction. UQNO consistently provides calibrated, efficient uncertainty bands, with especially strong results in the data-scarce car task.

  • 3D Car Surface Pressure Prediction: UQNO is the only method providing calibrated results for the data-scarce 3D car pressure task, while its band satisfies the 98% target.The comparison covers high- and low-domain-threshold settings; other baselines fail to provide calibrated estimates.
  • Darcy Flow: 99.1% domain coverage is obtained by UQNO in a Darcy-flow visualization, where predicted uncertainty matches the true-error structure better than the baselines.MCDropout misses the uncertain region, whereas Laplace approximation overestimates uncertainty by roughly 50×.
  • Experimental Setup: The experiments use a data-rich Darcy-flow setting with 5000 training samples and a data-scarce car-pressure setting with 500 training samples on a 3586-point 3D mesh.Darcy flow uses 421 × 421 resolution, while the car task predicts pressure at every surface mesh point.
  • Darcy Flow: UQNO delivers the tightest band among methods satisfying calibration on Darcy flow, reaching a 98.8% calibration percentage.The neural posterior principal components method has no finite bandwidth because it outputs a subspace; its explained error variance is 15%.
  • Darcy Flow: 1.52× tighter bands than MCDropout and 76.1× tighter bands than Laplace approximation are achieved by UQNO on 2D Darcy flow at the 98% target.Figure 3 compares bandwidth against calibration percentage with α = 0.1.
  • Darcy Flow: The bandwidth–calibration trade-off shows that higher calibration percentages require wider bands, while stricter domain thresholds increase bandwidth.UQNO can be calibrated to different target percentages and domain coverage thresholds.

5 Theoretical Guarantee

The framework constructs a risk-controlling confidence set whose expected domain coverage satisfies the desired calibration bound. The guarantee follows by combining calibration with a Hoeffding bound and extends to mixed discretizations.

  • The confidence set is constructed with the minimum possible λ, beginning with m sampled domain points per calibration function.
  • P[Ex[A(x)] < 1 −α] ≤δ, establishing the desired expected coverage guarantee.The proof combines calibration with a Hoeffding bound.
  • The same theoretical analysis extends to mixed discretizations in which calibration functions use different discretizations.

6 Related Work

Prior uncertainty methods for operator learning are limited by Gaussian assumptions, heuristic calibration, scalability, or pointwise rather than function-space coverage. This work positions conformal prediction as a distribution-free, finite-sample alternative for functional uncertainty.

  • Conformal Prediction: Conformal prediction provides a distribution-free and finite-sample foundation that has been extended across regression, forecasting, imaging, and functional data analysis.
  • Approximate Gaussian Methods: Gaussian-based methods can yield heuristic uncertainty estimates and may overestimate errors, while high-dimensional Bayesian approaches face scalability challenges.
  • Uncertainty in Operator Learning: Existing operator-learning methods include Gaussian, variance-based, or loss-focused approaches, whereas this work provides a bound simultaneously for all domain points.

7 Conclusion

The paper introduces a risk-controlling neural operator that uses conformal prediction and quantile loss for principled uncertainty estimates across function domains. It proves a calibration guarantee and reports tight, calibrated bands on 2D Darcy flow and 3D car pressure prediction, while noting important estimator and risk-measure limitations.

  • Method: The risk-controlling neural operator combines conformal prediction with quantile loss to provide principled uncertainty estimates simultaneously across the domain.
  • Empirical Results: The method consistently provides the tightest band while satisfying target calibration percentages across 2D Darcy flow and 3D car pressure prediction.In the 3D problem, it is the only method satisfying the 98% target.
  • Theoretical Guarantee: The paper provides a proof of the method's calibration percentage guarantee in addition to its empirical evaluation.
  • Limitations: Tight uncertainty bands still depend on a good heuristic estimator E; an uninformative estimator can produce very wide bands.
  • Limitations: The analysis uses expected domain risk, while worst-case risk, uncertainty decomposition, error representations, and efficient sampling remain open directions.

A Mixed Discretization Theoretical Guarantee

For mixed discretizations, the conservative construction uses the minimum number of sampled points across calibration functions. Finer discretization tightens concentration, but shared quantities require the minimum discretization level for the PAC guarantee.

  • Finer discretization provides more independent samples per function and a tighter concentration bound, but shared quantities require choosing the minimum m.
  • The conservative mixed-discretization construction uses m̄ = min{m1, ...mn} and still satisfies the (α, δ) PAC bound.
Loading 2402.01960v2…