Task-Driven Uncertainty Quantification in Inverse Problems via Conformal Prediction

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.

1. Introduction

In imaging inverse problems, one seeks to recover an image $x$ from measurements $y=h\left(x\right)$ that mask, distort, and/or corrupt $x$ with noise [7]. Linear inverse problems, where $y=Ax+ϵ$ with noise $ϵ$ and known forward operator $A$, include deblurring, super-resolution, inpainting, colorization, computed tomography (CT), and magnetic resonance imaging (MRI) [26]. Non-linear inverse problems include phase-retrieval, de-quantization, and image-to-image translation. These problems are generally ill-posed, in that it is impossible to perfectly infer $x$ from $y$.

Most image recovery methods provide a single “point estimate” $\widehat{x}$ from measurement $y$ [7]. From $\widehat{x}$ alone, it is difficult to determine accuracy with respect to the true $x$. That is, $\widehat{x}$ does not quantify the uncertainty [1] in measurement-and-reconstruction that arises from the ill-posed nature of the inverse problem. This is problematic in safety-critical applications like medical imaging, where hallucinations or degradations in $\widehat{x}$ can result in costly misdiagnoses [8].

Several approaches have been proposed to provide uncertainty quantification (UQ) within the image recovery process. One approach is to utilize Bayesian Neural Networks (BNNs), which treat the reconstruction network parameters as random variables [9, 21, 32, 40, 60]. This allows one to quantify epistemic (i.e., model) uncertainty by measuring the variation over reconstructions generated by different draws from the parameter distribution. Posterior sampling methods [2, 6, 14, 19, 20, 30] instead draw many samples from the distribution $p(x\mid y)$ of plausible $x$ given $y$, known as the posterior, and aim to quantify the uncertainty that the measurement process imposes on $x$ (i.e., aleatoric uncertainty). It’s possible to combine BNNs with posterior sampling, as in [22].

Although the samples generated by BNNs and posterior-sampling methods could be used in many ways, they are most commonly used to compute pixel-wise uncertainty images or “maps.” To compute pixel-wise uncertainty intervals with statistical guarantees, conformal prediction can be used [4, 29, 36, 53]. Still, the value of these pixel-wise uncertainty maps is not clear. For example, when recovering images, we are usually concerned about many-pixel visual structures (e.g., lesions in MRI, hallucinations) that single-pixel statistics say little about. Secondly, uncertainty maps are not easy to interpret. They often convey little beyond the notion that there is less pixel-wise uncertainty in smooth regions as compared to near edges. Lastly, it’s not clear how pixel-wise uncertainty relates to the overall imaging goal, which is often task-oriented, such as detecting whether a tumor is present or not.

One approach to assess multi-pixel uncertainty is Bayesian Uncertainty Quantification by Optimization (BUQO) [52], which aims to test whether a particular “structure of interest” in the maximum a-posteriori (MAP) reconstruction is truly present. However, inpainting is used to hypothesize what the image would look like without the structure, the correctness of which is difficult to guarantee.

In this work, we propose a novel UQ framework for imaging inverse problems that aims to provide a more impactful measure of uncertainty. In particular, we aim to quantify to what extent a downstream task behaves differently when supplied with the reconstructed image versus the true image. Our framework supports any measurement-and-reconstruction procedure and any downstream task that outputs a real-valued scalar. Our contributions are as follows.

1. We propose to construct, using conformal prediction, an interval in the task-output space that is guaranteed to contain the true task output up to a user-specified probability. The prediction interval width provides a natural way to quantify the uncertainty that measurement-and-reconstruction contributes to the downstream task output. (See Fig. 1.)

2. For posterior-sampling-based image reconstruction, we propose to construct adaptive uncertainty intervals that shrink when the measurements offer more certainty about the true output of the downstream task. (See Fig. 2.)

3. We propose a multi-round acquisition protocol whereby measurements are accumulated until the task uncertainty is acceptably low.

4. We demonstrate our approach on accelerated MRI with the task of soft-output-classifying a meniscus tear. Several conformal predictors are evaluated and compared.

Fig. 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 $𝒞\left(\widehat{x};d_{\text{cal}}\right)\subset \mathbb{R}$ that is guaranteed to contain the true task output $z=f(x)$ in the sense that $\mathbb{P}(Z\in 𝒞(\widehat{X};D_{\text{cal}}))\ge 1-\alpha$ for some chosen error-rate $\alpha$.

Fig. 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_{\text{cal}}={\{({\{{\widehat{x}}_i^{(j)}\}}_{j=1}^p,z_i)\}}_{i=1}^n$ to construct an interval $𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{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 𝒞(\{{\widehat{X}}^{(j)}\};D_{\text{cal}}))\ge 1-\alpha$ for some chosen error-rate $\alpha$.

2. Background

Conformal prediction [3, 54] is a framework for generating uncertainty sets with prescribed statistical guarantees. Notably, it can be applied to any black-box predictor without making any distributional assumptions about the data.

We now explain the basics of conformal prediction or, more precisely, the common variant known as split conformal prediction [37, 41]. Say that we have a black-box model $f:𝒳\to 𝒵$ that predicts a target $z\in 𝒵$ from features $x\in 𝒳$. Say that we also have a calibration dataset $d_{\text{cal}}\triangleq {\left\{\left(x_i,z_i\right)\right\}}_{i=1}^n$ that was unseen when training $f$, as well as a test feature $x_{\text{test}}$ and an unknown test target $z_{\text{test}}$. In split conformal prediction, we use the calibration data to construct a prediction set $𝒞\left(x_{\text{test}};d_{\text{cal}}\right)\subset 2^{𝒵}$ that provides the “marginal coverage” [38] guarantee

$$\mathbb{P}\left(Z_{\text{test}}\in 𝒞\left(X_{\text{test}};D_{\text{cal}}\right)\right)\ge 1-\alpha ,$$ (1)

where $\alpha \in [0,1]$ is a user-specified error rate. In (1) and elsewhere, we use capital letters to denote random variables and lower-case to denote their realizations. Thus (1) can be interpreted as follows: When averaged over random calibration data $D_{\text{cal}}$ and test data $\left(X_{\text{test}},Z_{\text{test}}\right)$, the set $𝒞\left(X_{\text{test}};D_{\text{cal}}\right)$ is guaranteed to contain the correct target $Z_{\text{test}}$ with probability no less than $1-\alpha$. Although we would prefer a “conditional coverage” guarantee of the form $\mathbb{P}\left(Z_{\text{test}}\in 𝒞\left(X_{\text{test}};D_{\text{cal}}\right)\mid X_{\text{test}}=x_{\text{test}}\right)\ge 1-\alpha$, this is generally impossible to achieve [38, 55].

We now describe the standard recipe for constructing a prediction set $𝒞\left(x_{\text{test}};d_{\text{cal}}\right)$. First one chooses a nonconformity score $s(x,z;f)\in \mathbb{R}$ that assigns higher values to worse predictions. Then one computes the empirical quantile

$$\widehat{q}\triangleq \mathrm{Q}\mathrm{u}\mathrm{a}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{l}\mathrm{e}\left(\frac{\lceil (1-\alpha )(n+1)\rceil }{n};s_1,\dots ,s_n\right)$$ (2)

from the calibration scores $s_i=s\left(x_i,z_i;f\right)$. Finally one constructs

$$𝒞\left(x_{\text{test}};d_{\text{cal}}\right)=\left\{z:s\left(x_{\text{test}},z;f\right)\le \widehat{q}\right\}.$$ (3)

Under these choices, it can be proven [37,56] that the marginal coverage guarantee (1) holds when $\left(X_1,Z_1\right),\dots ,\left(X_n,Z_n\right),\left(X_{\text{test}},Z_{\text{test}}\right)$ are i.i.d., and even under the weaker condition that they are exchangeable [54].

There are many ways to construct the nonconformity score $s(x,z;f)$. For real-valued targets $z$, the simplest choice would be the absolute residual

$$s\left(x,z;f\right)=\left|z-f\left(x\right)\right|\Rightarrow 𝒞\left(x;d_{\text{cal}}\right)=\left[f\left(x\right)-\widehat{q},f\left(x\right)+\widehat{q}\right],$$ (4)

which gives an $x$-invariant interval length of $\left|𝒞\left(x;d_{\text{cal}}\right)\right|=2\widehat{q}$. We will discuss a few other choices in the sequel. For more on conformal prediction, we suggest the excellent overviews [3, 54].

3. Proposed Method

Suppose that we collect measurements $y=h\left(x\right)$ of the true image $x$, from which we compute an image recovery $\widehat{x}=g\left(y\right)$. We would ideally like that $\widehat{x}=x$, but this is impossible to guarantee with an ill-posed inverse problem. Although there are many ways to quantify the difference between $\widehat{x}$ and $x$ (e.g., PSNR, SSIM [57], LPIPS [62], DISTS [18]), we will instead assume that we are primarily interested in using $\widehat{x}$ for some downstream task $f\left(\widehat{x}\right)\in \mathbb{R}$. As a running example, we consider $x$ to be a medical image, $y$ to be accelerated MRI measurements, and $f(\cdot )\in [0,1]$ to be the soft output of a classifier that aims to detect the presence or absence of a pathology. For example, when $f\left(\widehat{x}\right)=0.7$, the classifier believes that there is a 70% chance that the pathology exists.

When image recovery is imperfect (i.e., $\widehat{x}\ne x$), we expect the task output to also be imperfect, in the sense that $\widehat{z}=f\left(\widehat{x}\right)\ne f\left(x\right)=z$. We are thus strongly motivated to understand how close $\widehat{z}$ is to the true $z$ or, even better, to construct a prediction interval $𝒞\left(\widehat{x}\right)\subset \mathbb{R}$ that contains the true $z$ with some guarantee. The interval width $\left|𝒞\right(\widehat{x}\left)\right|$ would then quantify the uncertainty that the measurement-and-reconstruction process contributes to predicting the true task output $z$.

We emphasize that our approach makes no assumptions about the task $f(\cdot )$ beyond it producing a real number. For example, if $f(\cdot )$ is a soft-output classifier, we do not assume that it is accurate or even calibrated [25]. Likewise, our approach does not aim to assess the uncertainty implicit in the task, but rather the additional uncertainty that measurement-and-reconstruction contributes to the task. For a soft-output classifier, a (true) output of $z=f\left(x\right)=0.7$ would express considerable uncertainty about the presence of a pathology in $x$. But if the true $z$ could be perfectly predicted from $\widehat{x}$, then the measurement-and-reconstruction process would bring no additional uncertainty.

To construct the interval $𝒞\left(\widehat{x}\right)$, we use conformal prediction. Adapting the methodology from Sec. 2 to the current setting, we use a calibration set $d_{\text{cal}}={\left\{\left({\widehat{x}}_i,z_i\right)\right\}}_{i=1}^n$ of (recovered-image, true-task-output) pairs, and we expect to satisfy the marginal coverage guarantee

$$\mathbb{P}\left(Z\in 𝒞(\widehat{X};D_{\text{cal}})\right)\ge 1-\alpha$$ (5)

when $({\widehat{X}}_1,Z_1),\dots ,({\widehat{X}}_n,Z_n),(\widehat{X},Z)$ are exchangeable. In 5) and in the sequel, we explicitly denote the dependence of $𝒞\left(\widehat{x}\right)$ on the calibration data. To construct the calibration set, we assume access to ground-truth examples ${\left\{x_i\right\}}_{i=1}^n$, a measurement model $h(\cdot )$, a reconstruction model $g(\cdot )$, and a task function $f(\cdot )$, from which we can construct $y_i=h\left(x_i\right)$, ${\widehat{x}}_i=g\left(y_i\right)$, and $z_i=f\left(x_i\right)$ for $i=1,\dots ,n$.

In some cases we may instead have access to a posterior-sampling-based image reconstruction model that generates $p$ recoveries ${\{{\widehat{x}}_i^{(j)}\}}_{j=1}^p$ from every measurement $y_i$ via ${\widehat{x}}_i^{(j)}=g(y_i,v_i^{(j)})$, where ${\{v_i^{(j)}\}}_{j=1}^p$ are i.i.d. code vectors and, typically, $v_i^{(j)}~𝒩(0,I)$. In this case, the prediction interval becomes $𝒞({\{{\widehat{x}}_i^{(j)}\}}_{j=1}^p;d_{\text{cal}})$. As we will see, posterior sampling facilitates locally adaptive prediction sets.

Next we describe different ways to construct the prediction intervals $𝒞\left(\widehat{x};d_{\text{cal}}\right)$ and $𝒞({\{{\widehat{x}}_i^{(j)}\}}_{j=1}^p;d_{\text{cal}})$, and later we describe a multi-round measurement protocol that exploits locally adaptive prediction intervals. See Fig. 2 for a detailed overview of our approach.

3.1. Method 1: Absolute Residuals (AR)

We first consider the case where image recovery yields a point-estimate $\widehat{x}=g(y)$ of the true $x$. As described in Sec. 2, a simple way to construct a nonconformity score is through the absolute residual (recall (4))

$$s(\widehat{x},z;f)=|z-f(\widehat{x})|.$$ (6)

Evaluating this score on the calibration set $d_{\text{cal}}$ gives ${\left\{s_i\right\}}_{i=1}^n$, whose empirical quantile $\widehat{q}$ can be computed as in (2) and used to construct the prediction interval

$$𝒞\left(\widehat{x};d_{\text{cal}}\right)=[f(\widehat{x})-\widehat{q},f(\widehat{x})+\widehat{q}],$$ (7)

which then provides the marginal coverage property (5) [37].

Note that, with this choice of score, the interval width $\left|𝒞\left(\widehat{x};d_{\text{cal}}\right)\right|=2\widehat{q}$ varies with the calibration set $d_{\text{cal}}$ but not with $\widehat{x}$. Thus, for a fixed $d_{\text{cal}}$, the score (6) provides no way to tell whether one $\widehat{x}$ will yield more task-output uncertainty than a different $\widehat{x}$.

3.2. Method 2: Locally-Weighted Residuals (LWR)

We now consider the case where we have a posterior-sampling-based recovery method that yields $p$ recoveries ${\{{\widehat{x}}^{(j)}\}}_{j=1}^p$ per measurement $y$. We make no assumption on how accurate or diverse these $p$ samples are, other than assuming that the corresponding task-outputs ${\widehat{z}}^{(j)}=f({\widehat{x}}^{(j)})$ are not all identical.

Suppose that we choose the nonconformity score

$$s(\{{\widehat{x}}^{(j)}\},z;f)=\frac{|z-\bar{z}|}{{\sigma }_z}\text{with}\left\{\begin{array}{l}\bar{z}\triangleq \frac{1}{p}\sum _{j=1}^{p}f({\widehat{x}}^{(j)})\\ {\sigma }_z\triangleq \sqrt{\frac{1}{p}\sum _{j=1}^p{\left(f({\widehat{x}}^{(j)})-\bar{z}\right)}^2}\end{array}\right.,$$ (8)

evaluate it on the calibration set $d_{\text{cal}}$ to get scores ${\left\{s_i\right\}}_{i=1}^n$, and compute their empirical quantile $\widehat{q}$ as in (2). Then the prediction interval

$$𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}})=\left[\bar{z}-{\sigma }_z\widehat{q},\bar{z}+{\sigma }_z\widehat{q}\right]$$ (9)

of this “locally weighted residual” (LWR) method provides the marginal coverage property in (5) [37].

In words, this method first computes (approximate) posterior samples ${\widehat{z}}^{(j)}~Z\mid Y=y$, which are then averaged to approximate the conditional mean ${\widehat{z}}_{\text{mmse}}\triangleq \mathrm{E}(Z\mid Y=y)\approx \bar{z}$ and square-root conditional covariance $\sqrt{\mathrm{c}\mathrm{o}\mathrm{v}(Z\mid Y=y)}\approx {\sigma }_z$. When exactly computed, the conditional covariance gives a meaningful uncertainty metric on how well the true $Z$ can be estimated from measurements $y$, because $\mathrm{c}\mathrm{o}\mathrm{v}(Z\mid Y=y)=\mathrm{E}({\left(Z-{\widehat{z}}_{\text{mmse}}\right)}^2\mid Y=y)$. However, the ${\sigma }_z$ that we compute is merely an approximation. So, with the aid of the calibration set, ${\sigma }_z$ is adjusted by the scaling $\widehat{q}$ to yield a prediction interval $\left[\bar{z}-{\sigma }_z\widehat{q},\bar{z}+{\sigma }_z\widehat{q}\right]$ that satisfies the marginal coverage criterion (5).

Importantly, the interval width $|𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}})|$ now varies with ${\{{\widehat{x}}^{(j)}\}}_{j=1}^p$ through ${\sigma }_z$. This latter property is known as “local adaptivity” [37].

3.3. Method 3: Conformalized Quantile Regression (CQR)

Another popular locally adaptive method is known as conformalized quantile regression (CQR) [45]. The idea is to construct the nonconformity score using two quantile regressors [35], one which estimates the $\frac{\alpha }{2}$ th quantile of $Z\mid Y=y$ and the other which estimates the $\left(1-\frac{\alpha }{2}\right)$ th quantile.

To compute these quantile estimates, we will once again assume access to a posterior-sampling-based recovery method that yields $p$ recoveries ${\{{\widehat{x}}^{(j)}\}}_{j=1}^p$ per measurement $y$. From the corresponding task-outputs ${\widehat{z}}^{(j)}=f({\widehat{x}}^{(j)})$, we compute the empirical quantiles $\widehat{z}\left(\frac{\alpha }{2}\right)$ and $\widehat{z}\left(1-\frac{\alpha }{2}\right)$ using

$$\widehat{z}(\omega )\triangleq \text{Quantile}\left(\omega ;{\widehat{z}}^{(1)},\dots ,{\widehat{z}}^{(p)}\right).$$ (10)

From these quantile estimates, we construct the nonconformity score

$$s(\{{\widehat{x}}^{(j)}\},z;f)=\mathrm{m}\mathrm{a}\mathrm{x}\left\{\widehat{z}\left(\frac{\alpha }{2}\right)-z,z-\widehat{z}\left(1-\frac{\alpha }{2}\right)\right\},$$ (11)

evaluate it on the calibration set $d_{\text{cal}}$ to obtain ${\left\{s_i\right\}}_{i=1}^n$, and compute their $\lceil (1-\alpha )(n+1)\rceil /n$-empirical quantile $\widehat{q}$ as in (2). Then the prediction interval

$$𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}})=\left[\widehat{z}\left(\frac{\alpha }{2}\right)-\widehat{q},\widehat{z}\left(1-\frac{\alpha }{2}\right)+\widehat{q}\right]$$ (12)

provides the marginal coverage in (5) [45]. Like (9), this interval is locally adaptive. We will compare these three conformal prediction methods in Sec. 4.

3.4. Multi-Round Measurement Protocol

In many applications, there is a significant cost to collecting a large number of measurements (i.e., acquiring a high-dimensional $y$). One example is MRI, the details of which will be discussed in Sec. 4. For these applications, we propose to collect measurements over multiple rounds, stopping as soon as the task uncertainty falls below a prescribed level $\tau$. The goal is to collect the minimal number of measurements that accomplishes the task with probability of at least $1-\alpha$.

Our approach is to use the prediction interval width $|𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}})|$ as the metric for task uncertainty. This requires the interval to be locally adaptive, as with LWR and CQR above. The details are as follows. First, a sequence of $C>1$ nested measurement configurations is chosen, so that the resulting measurement sets obey ${𝒴}^{\left[1\right]}\subset {𝒴}^{\left[2\right]}\subset \cdots \subset {𝒴}^{\left[C\right]}$. Then, for each configuration $k=1,\dots ,C$, a calibration set $d_{\mathrm{c}\mathrm{a}\mathrm{l}}^{\left[k\right]}$ is collected, from which the set-valued function $𝒞(\cdot ;d_{\text{cal}}^{\left[k\right]})$ is constructed. At test time, we begin by collecting measurements $y\in {𝒴}^{\left[1\right]}$ according to the first (i.e., minimal) configuration. From $y$ we compute the reconstructions ${\{{\widehat{x}}^{(j)}\}}_{j=1}^p$ and, from them, the task uncertainty $|𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}}^{\left[1\right]})|$. If this uncertainty falls below the desired $\tau$, we stop collecting measurements. If not, we would collect the additional measurements in ${𝒴}^{\left[2\right]}\backslash {𝒴}^{\left[1\right]}$, and repeat the procedure. Figure 3 summarizes the proposed multi-round protocol.

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

4. Numerical Experiments

We now demonstrate our task-based uncertainty quantification framework on MRI [34]. MRI offers exceptional soft-tissue contrast without ionizing radiation but suffers from very slow scan times. Accelerated MRI speeds the acquisition process by collecting a fraction $1/R$ of the measurements specified by the Nyquist sampling theorem. The integer $R$ is known as the “acceleration rate.” When $R>1$, the inverse problem is ill-posed.

In MRI, a typical task is to diagnose the presence or absence of a pathology. Although this task is typically performed by a radiologist, neural-network-based classification is expected to play a significant role in aiding radiologists [12]. Thus, in our experiments, we implement the task $f(\cdot )$ using a neural network. Details are given below.

Data:

We use the multi-coil fastMRI knee dataset [61] and in particular the non-fat-suppressed subset, which includes 484 training volumes (17286 training slices, or images) and 100 validation volumes (2188 validation images). We use pathology labels from fastMRI+ [63]. For knee-MRI, meniscus tears yield the largest fastMRI+ label set, and so we choose meniscus-tear-detection as our task. To collect measurements, we retrospectively subsample the fastMRI data in the k-space using a set of random nested masks that yield acceleration rates $R\in \{16,8,4,2\}$, the details of which are described in the Supplementary Materials.

Image Recovery:

We consider two recovery networks $g(\cdot )$. As a point estimator, we use the state-of-the-art E2E-VarNet from [49] and, as a posterior sampler, we use the conditional normalizing flow (CNF) from [58]. Both were specifically designed around the fastMRI dataset. Another option would be the MRI diffusion sampler [15], but its performance is a bit worse than the CNF and its sampling speed is 8000 × slower [58]. The E2E-VarNet and CNF were each trained to handle all four acceleration rates with a single model.

Task Network:

We used a ResNet50 [27] for the task network $f(\cdot )$. Starting from an ImageNet-based initialization, we pretrained the weights to minimize the unsupervised SimCLR loss [13] and later minimized binary-cross-entropy loss using the fastMRI+ labels. See the Supplementary Materials for details.

Empirical Validation:

Recall that the marginal coverage guarantee (5) holds on average over random test samples $(\widehat{X},Z)$ and random calibration data $D_{\text{cal}}=\{({\widehat{X}}_1,Z_1),\dots ,({\widehat{X}}_n,Z_n)\}$. To empirically validate marginal coverage and evaluate other average-performance metrics, we perform Monte-Carlo averaging over $T=10000$ trials as follows. In each trial $t$, we randomly partition the 2188-sample validation dataset into a 70% calibration fold with indices $i\in {ℐ}_{\text{cal}}[t]$ and a 30% test fold with indices $i\in {ℐ}_{\text{test}}[t]$, construct conformal predictors using the calibration data $d_{\text{cal}}[t]={\{({\{{\widehat{x}}_i^{(j)}\}}_{j=1}^p,z_i)\}}_{i\in {ℐ}_{\text{cal}}[t]}$, and evaluate performance on test fold $t$. Finally, we average performance over the $T$ trials. Further details are given below.

4.1. Effect of Acceleration Rate and Conformal Prediction Scheme

We have seen that the interval length $|𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}})|$ provides a way to quantify the uncertainty that the measurement-and-reconstruction scheme contributes to the meniscus-classification task. So a natural question is: How is the interval length affected by the MRI acceleration $R$? We study this question below.

For a fixed acceleration $R$, the interval length is also affected by the choice of conformal predictor. All else being equal, better conformal predictors yield smaller uncertainty sets [3]. So another question is: How is the interval length affected by selecting among the AR, LWR, or CQR conformal methods?

To answer these questions, we compute the “average mean interval length” $\overline{\mathrm{M}\mathrm{I}\mathrm{L}}\triangleq \frac{1}{T}\sum _{t=1}^T\mathrm{M}\mathrm{I}\mathrm{L}[t]$ using the trial-$t$ mean interval length

$$\mathrm{M}\mathrm{I}\mathrm{L}\left[t\right]\triangleq \frac{1}{\left|{ℐ}_{\text{test}}\left[t\right]\right|}\sum _{i\in {ℐ}_{\text{test}}\left[t\right]}\left|𝒞({\{{\widehat{x}}_i^{\left(j\right)}\}}_{j=1}^p;d_{\text{cal}}\left[t\right])\right|,$$ (13)

Figure 4 plots the average mean interval length versus $R$ for the AR, LWR, and CQR conformal predictors using $T=10000$ trials, $p=32$ posterior samples, and error-rate $\alpha =0.05$. The figure shows that, as expected, the average mean interval length decreases as more measurements are collected (i.e., as $R$ decreases). The figure also shows that, as expected, the (locally adaptive) LWR and CQR methods give consistently smaller average mean interval lengths than the (non-adaptive) AR method. In this sense, posterior sampling is advantageous over point sampling.

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

4.2. Effect of Number of Posterior Samples

Above, we saw that the measurement process and conformal method both affect the prediction-interval length. We conjecture that the image reconstruction process will also affect the prediction-interval length. To investigate this, we vary the number of samples $p$ produced by the posterior-sampling scheme, reasoning that smaller values of $p$ correspond to less accurate recoveries (e.g., a less accurate posterior mean approximation).

Figure 4b plots the average mean interval length versus $p$ for the LWR and CQR conformal predictors using $T=10000$ trials, acceleration $R=16$, and error-rate $\alpha =0.05$. As expected, the interval length decreases as the posterior sample size $p$ grows. But interestingly, LWR is much more sensitive to small values of $p$ than CQR. One implication is that small values of $p$ may suffice when used with an appropriate conformal prediction method.

4.3. Empirical Validation of Coverage

To verify that the marginal coverage guarantee (5) holds, we compute the empirical coverage of Monte-Carlo trial $t$ as

$$\mathrm{E}\mathrm{C}\left[t\right]\triangleq \frac{1}{\left|{ℐ}_{\text{test}}\left[t\right]\right|}\sum _{i\in {ℐ}_{\text{test}}\left[t\right]}\mathbb{1}\{z_i\in 𝒞(\{{\widehat{x}}_i^{\left(j\right)}\};d_{\text{cal}}\left[t\right])\},$$ (14)

where $\mathbb{1}\{\cdot \}$ denotes the indicator function. Existing theory (see, e.g., (3) says that when $(\{{\widehat{X}}_i^{(j)}\},Z_i)$ and $D_{\text{cal}}[t]$ in 14 are exchangeable pairs of random variables, $\mathrm{E}\mathrm{C}[t]$ is random and distributed as

$$\mathrm{E}\mathrm{C}[t]~\mathrm{B}\mathrm{e}\mathrm{t}\mathrm{a}\mathrm{B}\mathrm{i}\mathrm{n}\left(n_{\text{test}},n_{\text{cal}}+1-l_{\text{cal}},l_{\text{cal}}\right)\text{for}{l}_{\text{cal}}\triangleq ⌈\left(n_{\text{cal}}+1\right)\alpha ⌉,$$ (15)

where $n_{\text{test}}\triangleq \left|{ℐ}_{\text{test}}[t]\right|$ and $n_{\text{cal}}\triangleq \left|{ℐ}_{\text{cal}}[t]\right|$.

For each of the three conformal methods, Fig. 5 shows the histogram of ${\{\mathrm{E}\mathrm{C}[t]\}}_{t=1}^T$ from (14) for $T=10000$, error-rate $\alpha =0.05$, acceleration $R=8$, and $p=32$ posterior samples. The figure shows that this histogram is close to the histogram created from $T$ samples of the theoretical distribution in (15). Figure 5 also prints the average empirical coverage $\frac{1}{T}\sum _{t=1}^T\mathrm{E}\mathrm{C}[t]$ for each method, which is very close to the target value of $1-\alpha =0.95$. Thus we see that, in practice, conformal prediction behaves close to the theory.

Fig. 5:

For the AR, LWR, and CQR conformal methods, each subplot shows the histograms of the empirical and theoretical empirical-coverage samples ${\{\mathrm{E}\mathrm{C}[t]\}}_{t=1}^T$ across $T=10000$ Monte-Carlo trials using $=0.05$, $R=8$, and $p=32$. The subplots are also labeled with the empirical mean of ${\{\mathrm{E}\mathrm{C}[t]\}}_{t=1}^T$, which is very close to the target value of $1-\alpha =0.95$.

4.4. Multi-Round Measurements

We now investigate the application of the multi-round measurement protocol from Sec. 3.4 to accelerated MRI. For this, we simulated the collection of MRI slices over rounds $k=1,\dots ,5$, stopping as soon as the $\alpha =0.01$ interval width $|𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}}^{\left[k\right]})|$ falls below the threshold of $\tau =0.1$. The first round collects k-space measurements at acceleration rate $R^{\left[1\right]}=16$, and the remaining rounds each collect additional k-space measurements to yield $R^{\left[2\right]}=8$, $R^{\left[3\right]}=4$, $R^{\left[4\right]}=2$, and $R^{\left[5\right]}=1$ respectively. For quantitative evaluation, we randomly selected 8 multi-slice volumes from the 100-volume fastMRI validation set to act as test volumes (half of which were labeled as meniscus tears and half of which were not), and we used the remaining 92 volumes for calibration. We will refer to the corresponding index sets as ${ℐ}_{\text{test}}$ and ${ℐ}_{\text{cal}}$. Additional details about the MRI measurement procedure are given in the Supplementary Materials.

We begin by discussing the AR conformal prediction method, which uses the point-sampling E2E-VarNet [49] for image recovery. The AR method produces prediction intervals $𝒞(\widehat{x};d_{\text{cal}}^{\left[k\right]})$ that are $\widehat{x}$-invariant (i.e., not locally adaptive). Thus, immediately after calibration, it is known that $k=4$ measurement rounds (i.e., $R=2$) are necessary and sufficient to achieve the $\tau =0.1$ threshold at error-rate $\alpha =0.01$.

The LWR and CQR conformal prediction methods both use the CNF [58] with $p=32$ posterior samples and yield locally adaptive prediction intervals $𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}}^{\left[k\right]})$. This allows them to evaluate the interval length for each $\{{\widehat{x}}^{(j)}\}$ and stop the measurement process as soon as that length falls below the threshold $\tau$. For test image $i\in {ℐ}_{\text{test}}$, we denote the final measurement round as

$$k_i\triangleq \mathrm{m}\mathrm{i}\mathrm{n}\{k:|𝒞(\{{\widehat{x}}_i^{(j)}\};d_{\text{cal}}^{\left[k\right]})|<\tau \}.$$ (16)

(Note that $\{{\widehat{x}}_i^{(j)}\}$ also changes with the measurement round $k$, although the notation does not explicitly show this.) The average acceleration is then

$$\bar{R}={\left(\frac{1}{\left|{ℐ}_{\text{test}}\right|}\sum _{i\in {ℐ}_{\text{test}}}\frac{1}{R^{\left[k_i\right]}}\right)}^{-1}.$$ (17)

Table 1 shows the average acceleration $\bar{R}$ for the AR, LWR, and CQR conformal methods. We see that $\bar{R}=2$ for the AR method because it always uses four measurement rounds. The LWR and CQR methods achieve higher average accelerations $\bar{R}$ because fewer measurement rounds suffice in a large fraction of cases. Table 1 also shows that the empirical coverage is close to what we would expect given this relatively small test set.

Table 1:

Average metrics for the multi-round MRI simulation (± standard error).

Method	Average Acceleration	Empirical Coverage	Average Max Center Error
AR	2.000	0.991 ± 0.008	0.032 ± 0.017
LWR	5.157	0.992 ± 0.005	0.020 ± 0.002
CQR	6.762	0.987 ± 0.008	0.044 ± 0.009

Figure 6 plots the distribution of final-round ${\left\{k_i\right\}}_{i\in {ℐ}_{\text{test}}}$ for the AR, LWR, and CQR conformal methods. It too shows that the AR method always uses four rounds (i.e., $R=2$), while the LWR and CQR methods typically use fewer rounds. However, this plot also shows that the LWR method sometimes uses five measurement rounds. This may seem counter-intuitive but can be explained as follows. At $k=4$, the AR method is calibrated so that the true score $z$ lands in the prediction interval in all but $\alpha =1\mathrm{\%}$ of the cases, where the length of that interval is small enough to meet the $\tau =0.1$ threshold. Meanwhile, the LWR (and CQR) methods adapt the prediction interval based on the difficulty of $\{{\widehat{x}}^{(j)}\}$. In most cases, the LWR prediction interval is smaller than the AR interval, but for a few “difficult” cases the LWR prediction interval is wider, and in fact too wide to meet the $\tau =0.1$ threshold. For these difficult cases, the LWR method moves on to the fifth measurement round.

Fig. 6:

Fraction of slices accepted after a given acceleration rate.

Based on the previous discussion, one might conjecture that the prediction intervals accepted by the AR method at round $k=4$ will be somehow worse than those accepted by LWR at $k=4$, even though their lengths all meet the threshold. We can confirm this by interpreting the midpoint of the prediction interval as an estimate of $z$ and evaluating the absolute error on that estimate, which we call the “center error” (CE):

$$\mathrm{C}\mathrm{E}(\{{\widehat{x}}^{(j)}\},z)\triangleq \left|z-\frac{b_l+b_u}{2}\right|\text{where}\left[b_l,b_u\right]=𝒞(\{{\widehat{x}}^{(j)}\};d_{\text{cal}}).$$ (18)

When evaluating the center error, we take the maximum over the slices in each volume. Table 1 lists the average maximum center error and confirms that it is smaller for LWR than for AR.

Figure 7 shows examples of image reconstructions, pixel-wise standard deviation maps, and CQR prediction intervals for a test image labeled with a meniscus tear. At higher accelerations like $R=16$, relatively large variations across posterior samples $\{{\widehat{x}}^{(j)}\}$ result in relatively large variations across classifier outputs $\{{\widehat{z}}^{(j)}\}$, which result in a large prediction interval, i.e., high uncertainty about the ground-truth classifier output $z$. At lower accelerations like $R=4$, relatively small variations across posterior samples yield smaller prediction intervals, i.e., less uncertainty about $z$. While the pixel-wise standard-deviation maps also show reduced variation across posterior samples, it’s difficult to draw conclusions about uncertainty in the downstream task from them. For example, the same pixel-wise variations could result from a set of reconstructions that show clear evidence for a tear in some cases and clear evidence to the contrary in others, or from a set of reconstructions that show clear evidence for a tear in all cases but are corrupted by different noise realizations. Our uncertainty quantification methodology circumvents these issues by focusing on the task itself. Furthermore, by leveraging the framework of conformal prediction, it ensures that the uncertainty estimates are statistically meaningful.

Fig. 7:

MR Image reconstructions and CQR prediction intervals at accelerations $R=16$ and $R=4$ with error-rate $\alpha =0.01$ and a total of $p=32$ posterior samples. The fastMRI+ bounding box around the meniscus tear is magnified in red. The prediction intervals shrink as the posterior samples become more consistent in the meniscus region. The standard-deviation maps show areas of high pixel-wise uncertainty but are difficult to connect to the downstream task. Note, image brightness was increased to better highlight the tear. Best viewed when zoomed.

As far as practical implementation is concerned, for each slice in a volume, the CNF reconstructions, ResNet-50 classifier outputs, and conformal prediction intervals can be computed in 414 milliseconds for $p=32$ samples, or 7.4 milliseconds for $p=2$ samples, on a single NVIDIA A100 GPU.

5. Discussion

A number of works on uncertainty quantification for MRI have been proposed based on Bayesian neural networks and posterior sampling, e.g., [11, 15, 17, 20, 21, 30, 40, 47, 58]. They produce a set of possible reconstructions $\{{\widehat{x}}^{(j)}\}$, from which a pixel-wise uncertainty map is typically computed. Conformal prediction methods have also been proposed to generate pixel-wise uncertainty maps for MRI and other imaging inverse problems [4, 29, 36, 53], but with statistical guarantees. However, when imaging is performed with the eventual goal of performing a downstream task, pixel-wise uncertainty maps are of questionable value. In this work, we construct a conformal prediction interval that is statistically guaranteed to contain the task output from the true image. We focus on tasks that output a real-valued scalar, such as soft-output binary classification.

Other works have applied conformal prediction to MRI tasks. Lu et al. [39] consider a dataset $\left\{\left(x_i,z_i\right)\right\}$ with MRI images $x_i$ and discrete ordinal labels $z_i\in \{1,\dots ,K\}$ that rate the severity of a pathology. They design a predictor that, given test $x$, outputs a set $𝒵(x)\in 2^K$ that is guaranteed to contain the true label $z$ with probability $1-\alpha$. Different from our work, [39] involves no inverse problem and aims to quantify the uncertainty in a discrete $z$. Sankaranarayanan et al. [46] compute uncertainty intervals on the presence/absence of semantic attributes in images, and mention that one application could be pathology detection in MRI (although they do not pursue it). Although their high-level goal is similar to ours, their solution requires a trained “disentangled” generative network that, in the case of MRI, would generate MRI images from pathology probabilities. To our knowledge, no such networks exist for MRI. In contrast, our method requires only a trained pathology classifier $f(\cdot )$, which should be readily available.

Limitations:

First, our method requires a downstream task, which is not always available. Second, we demonstrated our method on only a single inverse problem and task; validation on other applications is needed. Third, our MRI application ideas are preliminary and not ready for clinical use. Since we use the conformal prediction interval width as a proxy for the diagnostic value of the reconstructed image(s), several aspects of our design (e.g., the choice of classifier $f(\cdot )$, recovery algorithm $g(\cdot )$, conformal prediction method, threshold $\tau$, and error-rate $\alpha$) would need to be tuned and validated through rigorous clinical studies. Fourth, for ease of exposition, the conformal methods that we use (AR, LWR, CQR) are somewhat simple. More advanced methods, like risk-controlling prediction sets (RCPS) [10], may perform better. Fifth, we considered only tasks that output a single real-valued scalar, such as soft-output binary classification. Extensions to more general tasks would be useful. Lastly, our posterior sampler only considers aleatoric uncertainty. In principle, epistemic uncertainty could be included by sampling the generator’s weights from a distribution, as in [22], but more work is needed in this direction.

6. Conclusion

For imaging inverse problems, we proposed a method to quantify how much uncertainty the measurement-and-reconstruction process contributes to a downstream task, such as soft-output classification. In particular, we use conformal prediction to construct an interval that is guaranteed to contain the task-output from the true image with high probability. We showed that, with posterior-sampling-based image recovery methods, the prediction intervals can be made adaptive, and we proposed a multi-round measurement protocol that stops acquiring new data when the task uncertainty is sufficiently small. We applied our method to meniscus-tear detection in accelerated knee MRI and demonstrated significant gains in acceleration rate.