Calibrated Uncertainty Quantification for Operator Learning via Conformal Prediction

Ziqi Ma,Kamyar Azizzadenesheli,Anima Anandkumar
2024-02-06
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%.
Machine Learning
What problem does this paper attempt to address?
This paper attempts to solve the problem of reliable uncertainty quantification (UQ) in operator learning. Specifically, the paper focuses on how to provide calibrated uncertainty estimates within the function space to meet the requirements in practical applications. These applications include fields such as weather forecasting, carbon capture, automobile design, and nuclear fusion, which have strict requirements for the uncertainty estimates of model outputs, especially in scenarios where safety is crucial. ### Core Problems of the Paper 1. **Function Space Coverage**: Existing methods are usually only able to perform uncertainty quantification at a single point or a single scalar function, and cannot provide reliable uncertainty estimates over the entire function domain simultaneously. 2. **Distribution - Independent Calibration Guarantee**: Many traditional uncertainty quantification methods rely on the Gaussian assumption, which may not hold in practical applications. Therefore, a method that does not depend on a specific distribution is required to provide calibration guarantees. 3. **Scalability**: Some methods, such as Bayesian inference, are reasonable in principle but have high computational complexity on large - scale datasets and are difficult to scale up for practical applications. ### Solutions The paper proposes a risk - controlling quantile neural operator (UQNO), which combines the generalized quantile loss and the conformal prediction framework to solve the above problems. The specific steps are as follows: 1. **Pre - calibrated Quantile Neural Operator**: - Use an advanced neural operator architecture (such as the Fourier neural operator FNO) to parameterize the uncertainty estimator \( E \). - Train \( E \) through the generalized quantile loss function, which encourages the model output to be close to the \( 1 - \alpha \) quantile of the error magnitude during the training process. 2. **Calibration through Split Conformal Prediction**: - Utilize the conformal prediction framework to select an appropriate scaling factor \( \lambda \) so that the prediction set has good calibration performance on the calibration set. - Define the non - conformity score function \( s \) and select an appropriate value of \( \lambda \) according to the calibration set. 3. **Risk - Controlling Quantile Neural Operator**: - Combine the above two steps to obtain a risk - controlling quantile neural operator that can predict the uncertainty estimate at each point simultaneously when given a test function. - This method can provide advantages such as heteroscedasticity, simultaneous point prediction, and controllability. ### Experimental Results The paper verifies the effectiveness of the method on two tasks: 1. **2D Darcy Flow Problem**: This is a data - rich task with 5,000 training samples. The experimental results show that the UQNO method provides the tightest uncertainty band while meeting the 98% calibration target. 2. **3D Automobile Surface Pressure Prediction**: This is a data - scarce task with only 500 training samples. The experimental results show that UQNO is the only method that meets the 98% calibration target, and its uncertainty band is tighter than that of other methods. ### Main Contributions 1. Provide an uncertainty quantification method within the function space. 2. Provide uncertainty estimates that are calibrated simultaneously over the entire function domain. 3. Provide PAC bounds for uncertainty estimates. Through these contributions, the paper provides a powerful and practical solution for uncertainty quantification in operator learning.