Task-Driven Uncertainty Quantification in Inverse Problems via Conformal Prediction

TL;DR

The paper tackles uncertainty quantification in ill-posed imaging inverse problems by focusing on downstream task outputs and using conformal prediction to bound the true task value with probability at least $1-\alpha$. It introduces three conformal predictors—Absolute Residuals, Locally Weighted Residuals, and Conformalized Quantile Regression—and a multi-round acquisition protocol to stop data collection when task uncertainty is acceptably low, with intervals adaptively shrinking via posterior sampling. The approach is demonstrated on accelerated knee MRI for soft-output meniscus-tear classification, showing maintained coverage while achieving shorter, locally adaptive intervals and reduced data requirements. This task-centered UQ framework provides a principled, practically useful tool for uncertainty quantification in imaging with downstream decision implications, and it includes deployable code for replication.

Abstract

In imaging inverse problems, one seeks to recover an image from missing/corrupted measurements. Because such problems are ill-posed, there is great motivation to quantify the uncertainty induced by the measurement-and-recovery process. Motivated by applications where the recovered image is used for a downstream task, such as soft-output classification, we propose a task-centered approach to uncertainty quantification. In particular, we use conformal prediction to construct an interval that is guaranteed to contain the task output from the true image up to a user-specified probability, and we use the width of that interval to quantify the uncertainty contributed by measurement-and-recovery. For posterior-sampling-based image recovery, we construct locally adaptive prediction intervals. Furthermore, we propose to collect measurements over multiple rounds, stopping as soon as the task uncertainty falls below an acceptable level. We demonstrate our methodology on accelerated magnetic resonance imaging (MRI): https://github.com/jwen307/TaskUQ.

Figures (10)

• Figure 1: High-level overview of our approach: For true image $x$, measurement $y=h(x)$, recovery $\widehat{x}=g(y)$, and task output $\widehat{z}=f(\widehat{x})$, we use conformal prediction to construct an interval $\mathcal{C}(\widehat{x};d_{\mathsf{cal}})\subset{\mathbb{R}}$ that is guaranteed to contain the true task output $z=f(x)$ in the sense that $\mathbb{P}(Z\in\mathcal{C}(\widehat{X};D_{\mathsf{cal}}))\geq 1-\alpha$ for some chosen error-rate $\alpha$.
• Figure 2: Detailed overview of our approach: For true image $x$, measurement $y=h(x)$, reconstructions $\{\widehat{x}^{ (j)}\}_{j=1}^p$, and task outputs $\widehat{z}^{ (j)}=f(\widehat{x}^{ (j)})$, we use conformal prediction with a calibration set $d_{\mathsf{cal}}=\{(\{\widehat{x}_i^{ (j)}\}_{j=1}^p,z_i)\}_{i=1}^n$ to construct an interval $\mathcal{C}(\{\widehat{x}^{ (j)}\};d_{\mathsf{cal}})=[b_l,b_h]$ that is guaranteed to contain the true task output $z=f(x)$ in the sense that $\mathbb{P}(Z\in\mathcal{C}(\{\widehat{X}^{ (j)}\};D_{\mathsf{cal}}))\geq 1-\alpha$ for some chosen error-rate $\alpha$.
• Figure 3: Proposed multi-round measurement protocol. In each round, measurements are collected and reconstructions and conformal intervals are computed. If the length of the interval falls below a user-set threshold $\tau$, the procedure stops. Otherwise, more measurements are collected, and the process repeats until the threshold has been met.
• Figure 4: a) Average mean interval length versus acceleration $R$ with $p=32$ samples. b) Mean interval length versus $p$ with acceleration $R=16$. All results use error-rate $\alpha=0.05$ and $T=10000$ trials.
• Figure 5: For the AR, LWR, and CQR conformal methods, each subplot shows the histograms of the empirical and theoretical empirical-coverage samples $\{\mathop{\mathrm{EC}}\nolimits[t]\}_{t=1}^T$ across $T=10000$ Monte-Carlo trials using $\alpha=0.05$, $R=8$, and $p=32$. The subplots are also labelled with the empirical mean of $\{\mathop{\mathrm{EC}}\nolimits[t]\}_{t=1}^T$, which is very close to the target value of $1-\alpha=0.95$.