Calibrated Uncertainty Quantification for Operator Learning via Conformal Prediction

TL;DR

This work tackles calibrated uncertainty quantification for operator learning, where outputs are functions and require domain-wide calibration. It introduces UQNO, a distribution-free, finite-sample framework that learns a base operator $\hat{\mathcal{G}}$ and a per-point uncertainty proxy $E(a)(x)$, then forms $C_\lambda(a)(x)$ as a per-point ball around $\hat{\mathcal{G}}(a)(x)$ and calibrates it with split conformal prediction to achieve an $(\alpha,\delta)$-risk-controlling guarantee. The approach combines a generalized quantile loss for the base operator with conformal calibration, enabling simultaneous, pointwise uncertainty across the whole domain and providing a PAC bound on calibration coverage. Empirical results on 2D Darcy flow and 3D car surface pressure demonstrate calibrated, tight uncertainty bands that outperform baselines and, in the 3D case, meet a target calibration of $98\%$ where others fail. This framework offers principled, function-space uncertainty for safety-critical PDE-informed tasks with finite data, and opens avenues for extensions to mixed discretizations and uncertainty decomposition.

Abstract

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%.

Figures (7)

• Figure 1: Overall schematic of UQNO, a risk-controlling quantile neural operator. In operator learning, the learned neural operator outputs a function (red dots sampled at grid points). We train a residual operator with generalized quantile loss and then calibrate with conformal prediction, which yields simultaneous pointwise uncertainty estimates with a PAC guarantee on calibration coverage---the expected percentage of true value (black) that lies within our predicted uncertainty bands (green=upper bound and yellow=lower bound). In this example, since the output is 1D, we output a pointwise uncertainty band. In higher dimensions, we output a pointwise heterogeneous uncertainty ball.
• Figure 2: Uncertainty quantification comparison across methods on 2D Darcy flow problem. The leftmost heatmap plots true error. The top panels show the predicted pointwise uncertainty, and the bottom panels show the coverage (i.e. true error less than predicted error) for each point on the domain---yellow points are covered by our predicted uncertainty bands, and purple points are uncovered. The coverage percentage for each method is shown above the bottom panels. Our method of UQNO predicts uncertainty that corresponds well with true error while providing $99.1\%$ domain coverage. MCDropout does not capture the uncertain region well. Laplace approximation greatly overestimates uncertainty---on the scale of $50\times$ larger than true error, and thus appears all yellow in the heatmap.
• Figure 3: Bandwidth vs. calibration percentage comparison across methods on 2D Darcy flow problem. We see a clear advantage of UQNO (purple dot), providing $1.52\times$ tighter band than MCDropout and $76.1\times$ tighter band than Laplace approximation. This plot is generated with $\alpha=0.1$, target calibration percentage $98\%$ for UQNO, and $N=3$ principal components for posterior principal component method.
• Figure 4: Bandwidth vs. calibration percentage trade-off of UQNO on 2D Darcy flow. Each curve shows the bandwidth vs. calibration percentage trade-off for a fixed domain threshold ($\alpha$ in Equation \ref{['eq:rcps']}), demonstrating the flexibility to achieve a higher calibration percentage with wider bands. Overall bandwidth increases as the domain threshold becomes more stringent (smaller $\alpha$).
• Figure 5: Actual vs. target calibration percentage of UQNO. Different $\alpha$ values are plotted with different colors. We see the guarantee is always satisfied.