QUTCC🤗: Quantile Uncertainty Training and Conformal Calibration for Imaging Inverse Problems

Cassandra Tong Ye, Shamus Li, Tyler King, Kristina Monakhova

GitHub arXiv

Abstract

While deep learning offers tremendous promise for scientific and medical imaging, any failures and hallucinations (predictions that do not coincide with reality) are hard to pinpoint and can have serious downstream consequences. Uncertainty estimation techniques, such as conformal prediction, can help by predicting statistically valid error bars for a model's prediction. However, popular conformal prediction methods were not designed for high-dimensional image-valued problems and do not take into account spatial correlations within an image during conformal calibration, resulting in larger-than-necessary uncertainty intervals. We propose a practical simultaneous quantile regression method that enables non-linear, spatially-adaptive scaling during conformal calibration. Our method, QUTCC uses a U-Net architecture with a quantile embedding to learn a full conditional quantile distribution during training, and then leverages this non-linear, learned function for spatially-adaptive conformal calibration. At test time, our method can efficiently estimate uncertainty intervals with pixel-marginal coverage guarantees. In addition, QUTCC can also predict pixel-wise conditional probability density estimates without any built-in distributional assumptions. We evaluate our method on several denoising problems, accelerated magnetic resonance imaging, and quantitative phase microscopy. Our method consistently produces tighter uncertainty intervals than prior conformal methods at the same coverage level, can predict plausible conditional distributions for different tasks, and in some cases, high-uncertainty regions can help us locate hallucinations in a model's prediction.

Quantile Regression: Pinball Loss

Instead of predicting a single "best guess," quantile regression predicts a chosen quantile of the target distribution, for example a value the true answer falls below 10% or 90% of the time.

  • q ∈ (0, 1) is the target quantile level
  • x is the true value, xÌ‚ is the prediction
  • Over- and under-estimates are penalized unequally, depending on q
Lq(x, x̂) = {
q · |x − x̂|if x − x̂ ≥ 0
(1 − q) · |x − x̂|otherwise
Pinball loss
Image at each q goes here
QUTCC prediction of an MRI sample at q = 0.50
q = 0.50
Drag the slider to change q

Our Method: QUTCC

We propose Quantile Uncertainty Training and Conformal Calibration (QUTCC🤗), an uncertainty quantification method for imaging inverse problems.

  • One learned network: a single network learns the full range of quantiles, instead of just a fixed upper and lower bound.
  • Spatially Adaptive Calibration: bounds are scaled non-uniformly and nonlinearly during conformal calibration, instead of by one constant factor like prior methods.
  • QUTCC predicts smaller UQ intervals: our method predicts tighter, more informative uncertainty intervals.
  • QUTCC can predict the full conditional distribution: since every quantile is learned, QUTCC can query them all at inference time to estimate each pixel's probability distribution.
  • Tested on 5 inverse problems: Gaussian noise, Poisson noise, real noise, MRI, and quantitative phase imaging (QPI).
Training Calibration Inference
Training
During training, a neural network with a quantile embedding predicts an image as a function of the measurement and quantile, q. The quantile embedding is randomly sampled (q ∈ (0, 1)) and the value of q determines the asymmetry of the pinball loss, enabling the model to learn a range of conditional quantiles.
Calibration
During calibration, the predictive bounds (q_lower, q_upper) are iteratively adjusted on a held-out dataset to satisfy the desired miscoverage level α.
Inference
At test time, the model is queried with q_lower, q_0.5, and q_upper to produce the mean prediction and a corresponding pixel-wise uncertainty map. Querying the full range of quantile values enables the prediction of the pixel-wise distribution.

QUTCC Produces Smaller Uncertainty Intervals

  • For evaluation, we compare MC-Dropout, Deep Ensembles, Im2Im, Im2Im-Asym., K-RCPS, and our method QUTCC on five imaging inverse problems: accelerated MRI, quantitative phase imaging (QPI), and denoising under real noise, synthetic Poisson noise, and Gaussian noise.
  • We compare predicted uncertainty interval lengths stratified by pixel intensity for each task, and QUTCC achieves either the smallest or second-smallest interval lengths across the tasks.
  • When QUTCC ranks second, it is usually behind MC-Dropout or Deep Ensembles, which are far more costly: 28×–47× slower inference and ~10× longer training, respectively.
  • QUTCC achieves the smallest mean interval length on every task among conformal methods, and remains competitive with the posterior-sampling baselines despite using a single model and only two forward passes to construct each bound.

Hover over the chart for mean ± std at each intensity. Data from Table 1 of the paper.

  • QUTCC's smaller uncertainty interval lengths mean it can be more precise about pinpointing areas of error and hallucination.
  • Below, we show predictions from QUTCC and Im2Im-UQ on an undersampled MRI task.
  • Both models highlight regions of high uncertainty that correspond to regions of high error, but QUTCC predicts precise regions.

Pixel-wise Probability Density Functions

Because QUTCC learns every quantile, it can predict a full probability density function (PDF) at each pixel, not just an upper and lower bound.

  • Prior quantile methods like Im2Im-UQ only predict discrete upper and lower bounds, so they can't describe the full distribution at a pixel.
  • QUTCC queries many quantiles from a single network and calibrates them across a range of miscoverage levels α (e.g., 0.1 to 0.9).
  • Together, these calibrated intervals approximate the cumulative distribution function (CDF) at each pixel. Differentiating it gives a conformalized pixel-wise PDF with coverage guarantees at each risk level.
  • Below, we take a noisy Gaussian image (σ = 0.1) and compare the predicted PDFs at three pixels.
Noisy Gaussian measurement with three marked pixels
Gaussian noise, σ = 0.1. Click a circled pixel.

  • We compare the predicted conditional PDFs of 10 representative pixels, recovered from a measurement with different levels of Gaussian and Poisson noise, across QUTCC, Ensemble, and MC-Dropout.
  • Consistent with a Gaussian distribution, as σ increases, the conditional PDF at each pixel becomes broader, reflecting the increased variance of the distribution.
  • In contrast, in the Poisson setting, increasing λ shifts the distribution upward, while the spread increases in proportion to the mean.
Fig. 5: predicted pixel PDFs, Gaussian noise, less noise

Gaussian noise, σ = 0.1. Predicted PDFs at the 10 colored pixels for QUTCC, Deep Ensemble and MC-Dropout. From Fig. 5 of the paper.

Conclusion

We propose QUTCC, a new uncertainty quantification method for imaging inverse problems that uses a learned, nonlinear, and spatially-adaptive conformal calibration. QUTCC trains a single U-Net with a quantile embedding on q ∈ (0, 1) quantiles through simultaneous quantile regression, then adjusts its quantile bound predictions during calibration until the desired risk is satisfied. Across five imaging inverse problems, QUTCC achieves smaller uncertainty intervals than prior conformal methods and is competitive with MC-Dropout and Deep Ensembles at a fraction of their compute cost, all while maintaining the same statistical coverage. QUTCC can also estimate a conformalized pixel-wise conditional PDF, something previous conformal methods for image-to-image regression could not do, and its high-uncertainty regions often coincide with reconstruction errors and hallucinations. While QUTCC requires paired data for training and calibration and offers pixel-wise marginal coverage guarantees, we hope it helps practitioners identify harmful mistakes in scientific and medical imaging.
Bibtex Citation