# How to Trust Your Diffusion Model:
A Convex Optimization Approach to Conformal Risk Control

Jacopo Teneggi111Corresponding authors<EMAIL_ADDRESS><EMAIL_ADDRESS>.
222Department of Computer Science, Johns Hopkins University, Baltimore, MD,
21218 333Mathematical Institute for Data Science (MINDS), Johns Hopkins
University, Baltimore, MD, 21218 Matthew Tivnan444Department of Biomedical
Engineering, Johns Hopkins University, Baltimore, MD, 21218 J. Webster
Stayman44footnotemark: 4 Jeremias Sulam 11footnotemark: 1 44footnotemark: 4
33footnotemark: 3

###### Abstract

Score-based generative modeling, informally referred to as _diffusion models_
, continue to grow in popularity across several important domains and tasks.
While they provide high-quality and diverse samples from empirical
distributions, important questions remain on the reliability and
trustworthiness of these sampling procedures for their responsible use in
critical scenarios. Conformal prediction is a modern tool to construct finite-
sample, distribution-free uncertainty guarantees for any black-box predictor.
In this work, we focus on image-to-image regression tasks and we present a
generalization of the Risk-Controlling Prediction Sets (RCPS) procedure, that
we term $K$-RCPS, which allows to _(i)_ provide entrywise calibrated intervals
for future samples of any diffusion model, and _(ii)_ control a certain notion
of risk with respect to a ground truth image with minimal mean interval
length. Differently from existing conformal risk control procedures, ours
relies on a novel convex optimization approach that allows for
multidimensional risk control while provably minimizing the mean interval
length. We illustrate our approach on two real-world image denoising problems:
on natural images of faces as well as on computed tomography (CT) scans of the
abdomen, demonstrating state of the art performance.

##  1 Introduction

Generative modeling is one of the longest standing tasks of classical and
modern machine learning [12]. Recently, the foundational works by Song and
Ermon [50], Song et al. [52], Pang et al. [43] on sampling via _score-
matching_ [27] and by Ho et al. [23] on _denoising diffusion models_ [49]
paved the way for a new class of _score-based_ generative models, which solve
a reverse-time stochastic differential equation (SDE) [53, 2]. These models
have proven remarkably effective on both unconditional (i.e., starting from
random noise) and conditional (e.g., inpainting, denoising, super-resolution,
or class-conditional) sample generation across a variety of fields [61, 17].
For example, score-based generative models have been applied to inverse
problems in general computer vision and medical imaging [29, 32, 31, 59, 55,
54], 3D shape generation [62, 60, 42], and even in protein design [25, 16, 58,
28].

These strong empirical results highlight the potential of score-based
generative models. However, they currently lack of precise statistical
guarantees on the distribution of the generated samples, which hinders their
safe deployment in high-stakes scenarios [26]. For example, consider a
radiologist who is shown a computed tomography (CT) scan of the abdomen of a
patient reconstructed via a score-based generative model. How confident should
they be of the fine-grained details of the presented image? Should they trust
that the model has not _hallucinated_ some of the features (e.g.,
calcifications, blood vessels, or nodules) involved in the diagnostic process?
Put differently, how different will future samples be from the presented
image, and how far can we expect them to be from the ground truth image?

In this work we focus on image-to-image regression problems, where we are
interested in recovering a high-quality ground truth image given a low-quality
observation. While our approach is general, we focus on the problem of image
denoising as a running example. We address the questions posed above on the
reliability of score-based generative models (and, more generally, of any
sampling procedure) through the lens of conformal prediction [44, 57, 38, 48,
3] and conformal risk control [10, 4, 5] which provide any black-box predictor
with _distribution-free_ , _finite-sample_ uncertainty guarantees. In
particular, the contribution of this paper is three-fold:

  1. 1.

Given a fixed score network, a low-quality observation, and any sampling
procedure, we show how to construct valid entrywise calibrated intervals that
provide coverage of future samples, i.e. future samples (on the same
observation) will fall within the intervals with high probability;

  2. 2.

We introduce a novel high-dimensional conformal risk control procedure that
minimizes the mean interval length directly, while guaranteeing the number of
pixels in the ground truth image that fall outside of these intervals is below
a user-specified level on future, unseen low-quality observations;

  3. 3.

We showcase our approach for denoising of natural face images as well as for
computed tomography of the abdomen, achieving state of the art results in mean
interval length.

Going back to our example, providing such uncertainty intervals with provable
statistical guarantees would improve the radiologist’s trust in the sense that
these intervals precisely characterize the type of tissue that could be
reconstructed by the model. Lastly, even though our contributions are
presented in the context of score-based generative modeling for regression
problems—given their recent popularity [33, 61, 18]—our results are broadly
applicable to any sampling procedure, and we will comment on potential direct
extensions where appropriate.

###  1.1 Related work

##### Image-to-Image Risk Control

Previous works have explored conformal risk control procedures for image-to-
image regression tasks. In particular, Angelopoulos et al. [6] show how to
construct set predictors from heuristic notions of uncertainty (e.g., quantile
regression [34, 45]) for any image regressor, and how to calibrate the
resulting intervals according to the original RCPS procedure of Bates et al.
[10]. Kutiel et al. [35] move beyond set predictors and propose a mask-based
conformal risk control procedure that allows for notions of distance between
the ground truth and predicted images other than interval-based ones. Finally,
and most closely to this paper, Horwitz and Hoshen [26] sketch ideas of
conformal risk control for diffusion models with the intention to integrate
quantile regression and produce heuristic sampling bounds without the need to
sample several times. Horwitz and Hoshen [26] also use the original RCPS
procedure to guarantee risk control. Although similar in spirit, the
contribution of this paper focuses on a high-dimensional generalization of the
original RCPS procedure that formally minimizes the mean interval length. Our
proposed procedure is agnostic of the notion of uncertainty chosen to
construct the necessary set predictors.

##  2 Background

First, we briefly introduce the necessary notation and general background
information. Herein, we will refer to images as vectors in ${\mathbb{R}}^{d}$,
such that ${\mathcal{X}}\subset{\mathbb{R}}^{d}$ and
${\mathcal{Y}}\subset{\mathbb{R}}^{d}$ indicate the space of high-quality
ground truth images, and low-quality observations, respectively. We assume
both ${\mathcal{X}}$ and ${\mathcal{Y}}$ to be bounded. For a general image-
to-image regression problem, given a pair $(x,y)$ drawn from an unknown
distribution ${\mathcal{D}}$ over ${\mathcal{X}}\times{\mathcal{Y}}$, the task
is to retrieve $x\in{\mathcal{X}}$ given $y\in{\mathcal{Y}}$. This is usually
carried out by means of a predictor $f:~{}{\mathcal{Y}}\to{\mathcal{X}}$ that
minimizes some notion of distance (e.g., MSE loss) between the ground truth
images and reconstructed estimates on a set
$\\{(x_{i},y_{i})\\}_{i=1}^{n}\sim{\mathcal{D}}^{n}$ of $n$ pairs of high- and
low-quality images. For example, in the classical denoising problem, one has
$y=x+v_{0}$ where $v_{0}\sim{\mathcal{N}}(0,\sigma_{0}^{2}{\mathbb{I}})$ is
random Gaussian noise with variance $\sigma_{0}^{2}$, and one wishes to learn
a denoiser $f$ such that $f(y)\approx x$.

###  2.1 Score-based Conditional Sampling

Most image-to-image regression problems are ill-posed: there exist several
ground truth images that could have generated the same low-quality
observation. This is easy to see for the classical denoising problem described
above. Instead of a point predictor $f$—which could approximate a maximum-a-
posteriori (MAP) estimate—one is often interested in devising a sampling
procedure $F:~{}{\mathcal{Y}}\to{\mathcal{X}}$ for the posterior $p(x|y)$,
which precisely describes the distribution of possible ground truth images
that generated the observation $y$. In real-world scenarios, however, the full
joint $(x,y)$ is unknown, and one must resort to approximate $p(x|y)$ from
finite data. It is known that for a general Itô process
${\text{d}}{x}=h(x,t)\,{\text{d}}{t}+g(t)\,{\text{d}}{w}$ that perturbs an
input $x$ into random noise [30], it suffices to know the _Stein score_
${\nabla}_{x}\log p_{t}(x)$ [2, 39] to sample from $p(x)$ via the reverse-time
process

${\text{d}}{x}=[h(x,t)-g(t)^{2}{\nabla}_{x}\log
p_{t}(x)]\,{\text{d}}{t}+g(t)\,{\text{d}}{\bar{w}},$ (1)

where $h(x,t)$ and $g(t)$ are a _drift_ and _diffusion_ term, respectively,
and ${\text{d}}{w}$ and ${\text{d}}{\bar{w}}$ are forward- and reverse-time
standard Brownian motion.555We will assume time $t$ continuous in $[0,1]$.
Furthermore, if the likelihood $p(y|x)$ is known—which is usually the case for
image-to-image regression problems—it is possible to condition the sampling
procedure on an observation $y$. Specifically, by Bayes’ rule, it follows that
${\nabla}_{x}\log p_{t}(x|y)={\nabla}_{x}\log p_{t}(y|x)+{\nabla}_{x}\log
p_{t}(x)$ which can be plugged-in into the reverse-time SDE in Equation 1 to
sample from $p(x|y)$.

Recent advances in generative modeling by Song and Ermon [50], Song et al.
[53] showed that one can efficiently train a _time-conditional_ score network
$s(\tilde{x},t)$ to approximate the score ${\nabla}_{x}\log p_{t}(\tilde{x})$
via _denoising score-matching_ [27]. In this way, given a forward-time SDE
that models the observation process, a score network
$s(\tilde{x},t)\approx{\nabla}_{x}\log p_{t}(\tilde{x})$, and the likelihood
term $p(y|\tilde{x})$, one can sample from $p(x|y)$ by solving the conditional
reverse-time SDE with any discretization (e.g., Euler-Maruyama) or predictor-
corrector scheme [53]. While these models perform remarkably well in practice,
limited guarantees exist on the distributions that they sample from [37].
Instead, we will provide guarantees for diffusion models by leveraging ideas
of conformal prediction and conformal risk control, which we now introduce.

###  2.2 Conformal Prediction

Conformal prediction has a rich history in mathematical statistics [57, 44,
56, 8, 9, 22].666Throughout this work, we will refer to _split_ conformal
prediction [57] simply as _conformal prediction_. It comprises various
methodologies to construct finite-sample, statistically valid uncertainty
guarantees for general predictors without making any assumption on the
distribution of the response (i.e., they are distribution-free). It
particular, these methods construct valid prediction sets that provide
coverage, which we now define.

######  Definition 2.1 (Coverage [48]).

Let $z_{1},\dots,z_{m},z_{m+1}$ be $m+1$ exchangeable random variables drawn
from the same unknown distribution ${\mathcal{Q}}$ over ${\mathcal{Z}}$. For a
desired miscoverage level $\alpha\in[0,1]$, a set ${\mathcal{C}}\subseteq
2^{{\mathcal{Z}}}$ that only depends on $z_{1},\dots,z_{m}$ provides coverage
if

${\mathbb{P}}[z_{m+1}\in{\mathcal{C}}]\geq 1-\alpha.$ (2)

We remark that the notion of coverage defined above was introduced in the
context of classification problems, where one is interested in guaranteeing
that the true, unseen label of a future sample will be in the prediction set
${\mathcal{C}}$ with high probability. It is immediate to see how conformal
prediction conveys a very precise notion of uncertainty—the larger
${\mathcal{C}}$ has to be in order to guarantee coverage, the more _uncertain_
the underlying predictor. We refer the interested reader to [48, 3] for
classical examples of conformal prediction.

In many scenarios (e.g., regression), the natural notion of uncertainty may be
different from miscoverage as described above (e.g., $\ell_{2}$ norm). We now
move onto presenting conformal risk control, which extends the coverage to any
notion of risk.

###  2.3 Conformal Risk Control

Let ${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}$ be a general
_set-valued_ predictor from ${\mathcal{Y}}$ into
${\mathcal{X}}^{\prime}\subseteq 2^{{\mathcal{X}}}$. Consider a nonnegative
loss $\ell:~{}{\mathcal{X}}\times{\mathcal{X}}^{\prime}\to{\mathbb{R}}$
measuring the discrepancy between a ground truth $x$ and the predicted
intervals ${\mathcal{I}}(y)$. We might be interested in guaranteeing that this
loss will be below a certain tolerance ${\epsilon}\geq 0$ with high
probability on future, unseen samples $y$ for which we do not know the ground
truth $x$. Conformal risk control [10, 4, 5] extends ideas of conformal
prediction in order to select a specific predictor ${\mathcal{I}}$ that
controls the risk ${\mathbb{E}}[\ell(x,{\mathcal{I}}(y))]$ in the following
sense.

######  Definition 2.2 (Risk Controlling Prediction Sets).

Let
${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}\sim{\mathcal{D}}^{n}$
be a calibration set of $n$ i.i.d. samples from an unknown distribution
${\mathcal{D}}$ over ${\mathcal{X}}\times{\mathcal{Y}}$. For a desired risk
level ${\epsilon}\geq 0$ and a failure probability $\delta\in[0,1]$, a random
set-valued predictor
${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}\subseteq
2^{{\mathcal{X}}}$ is an $({\epsilon},\delta)$-RCPS w.r.t. a loss function
$\ell:~{}{\mathcal{X}}\times{\mathcal{X}}^{\prime}\to{\mathbb{R}}$ if

${\mathbb{P}}_{{\mathcal{S}}_{{\text{cal}}}}[{\mathbb{E}}_{(x,y)\sim{\mathcal{D}}}[\ell(x,{\mathcal{I}}(y))]\leq{\epsilon}]\geq
1-\delta.$ (3)

Bates et al. [10] introduced the first conformal risk control procedure for
_monotonically nonincreasing_ loss functions, those that satisfy, for a fixed
$x$,

${\mathcal{I}}(y)\subset{\mathcal{I}}^{\prime}(y)\implies\ell(x,{\mathcal{I}}^{\prime}(y))\leq\ell(x,{\mathcal{I}}(y)).$
(4)

In this way, increasing the size of the sets cannot increase the value of the
loss. Furthermore, assume that for a fixed input $y$ the family of set
predictors $\\{{\mathcal{I}}_{\lambda}(y)\\}_{\lambda\in\Lambda}$, indexed by
$\lambda\in\Lambda$,
$\Lambda\subset\overline{{\mathbb{R}}}\coloneqq{\mathbb{R}}\cup\\{\pm\infty\\}$,
satisfies the following nesting property [22]

$\lambda_{1}<\lambda_{2}\implies{\mathcal{I}}_{\lambda_{1}}(y)\subset{\mathcal{I}}_{\lambda_{2}}(y).$
(5)

Denote $R(\lambda)={\mathbb{E}}[\ell(x,{\mathcal{I}}_{\lambda}(y))]$ the risk
of ${\mathcal{I}}_{\lambda}(y)$ and $\hat{R}(\lambda)$ its empirical estimate
over a calibration set
${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}$. Finally, let
$\hat{R}^{+}(\lambda)$ be a _pointwise upper confidence bound_ (UCB) that
covers the risk, that is

${\mathbb{P}}[R(\lambda)\leq\hat{R}^{+}(\lambda)]\geq 1-\delta$ (6)

for _each_ , fixed value of $\lambda$—such that can be derived by means of
concentration inequalities (e.g., Hoeffding’s inequality [24], Bentkus’
inequality [11], or respective hybridization [10]).777We stress that Equation
6 does _not_ imply _uniform coverage_ $\forall\lambda\in\Lambda$. With these
elements, Bates et al. [10] show that choosing

$\hat{\lambda}=\inf\\{\lambda\in\Lambda:~{}\hat{R}^{+}(\lambda^{\prime})<{\epsilon},~{}\forall\lambda^{\prime}\geq\lambda\\}$
(7)

guarantees that ${\mathcal{I}}_{\hat{\lambda}}(y)$ is an
$({\epsilon},\delta)$-RCPS according to Definition 2.2. In other words,
choosing $\hat{\lambda}$ as the smallest $\lambda$ such that the UCB is below
the desired level ${\epsilon}$ for all values of $\lambda\geq\hat{\lambda}$
controls the risk al level $\epsilon$ with probability at least $1-\delta$.
For the sake of completeness, we include the original conformal risk control
procedure in Algorithm 2 in Appendix B.

Equipped with these general concepts, we now move onto presenting the
contributions of this work.

##  3 How to Trust Your Diffusion Model

We now go back to the main focus of this paper: solving image-to-image
regression problems with diffusion models. Rather than a point-predictor
$f:~{}{\mathcal{Y}}\to{\mathcal{X}}$, we assume to have access to a stochastic
sampling procedure $F:~{}{\mathcal{Y}}\to{\mathcal{X}}$ such that $F(y)$ is a
random variable with unknown distribution ${\mathcal{Q}}_{y}$—that hopefully
approximates the posterior distribution of $x$ given $y$, i.e.
${\mathcal{Q}}_{y}\approx p(x|y)$. However, we make no assumptions on the
quality of this approximation for our results to hold. As described in Section
2.1, $F$ can be obtained by means of a time-conditional score network
$s(\tilde{x},t)$ and a reverse-time SDE. While our results are applicable to
_any_ sampling procedure, we present them in the context of diffusion models
because of their remarkable empirical results and increasing use in critical
applications [61, 17].

One can identify three separate sources of randomness in a general stochastic
image-to-image regression problem: _(i)_ the unknown prior $p(x)$ over the
space of ground-truth images, as $x\sim p(x)$, _(ii)_ the randomness in the
observation process of $y$ (which can be modeled by a forward-time SDE over
$x$), and finally _(iii)_ the stochasticity in the sampling procedure $F(y)$.
We will first provide conformal prediction guarantees for a fixed observation
$y$, and then move onto conformal risk control for the ground truth image $x$.

###  3.1 Calibrated Quantiles for Future Samples

Given the same low-quality (e.g., noisy) observation $y$, where will future
unseen samples from $F(y)\sim{\mathcal{Q}}_{y}$ fall? How concentrated will
they be? Denote ${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}$ a
(random) set-valued predictor from ${\mathcal{Y}}\subset{\mathbb{R}}^{d}$ into
a space of sets ${\mathcal{X}}^{\prime}\subseteq 2^{{\mathcal{X}}}$ over
${\mathcal{X}}\subset{\mathbb{R}}^{d}$ (e.g., ${\mathcal{X}}=[0,1]^{d}$,
${\mathcal{X}}^{\prime}\subseteq 2^{[0,1]^{d}}$). We extend the notion of
coverage in Definition 2.1 to _entrywise coverage_ , which we now make
precise.

######  Definition 3.1 (Entrywise coverage).

Let $z_{1},\dots,z_{m},z_{m+1}$ be $m+1$ exchangeable random vectors drawn
from the same unknown distribution ${\mathcal{Q}}$ over
${\mathcal{X}}\subset{\mathbb{R}}^{d}$. For a desired miscoverage level
$\alpha\in[0,1]$, a set ${\mathcal{I}}\subseteq 2^{{\mathcal{X}}}$ that only
depends on $z_{1},\dots,z_{m}$ provides entrywise coverage if

${\mathbb{P}}[(z_{m+1})_{j}\in{\mathcal{I}}_{j}]\geq 1-\alpha$ (8)

for each $j\in[d]\coloneqq\\{1,\dots,d\\}$.

We stress that the definition above is different from notions of _vector
quantiles_ [13, 15] in the sense that coverage is not guaranteed over the
entire new random vector $z_{m+1}$ but rather along each dimension
independently. Ideas of vector quantile regression (VQR) are complementary to
the contribution of the current work and subject of ongoing research [21, 14,
47].

For a fixed observation $y$, we use conformal prediction to construct a set
predictor that provides entrywise coverage.

######  Lemma 3.2 (Calibrated quantiles guarantee entrywise coverage).

Let $F:~{}{\mathcal{Y}}\to{\mathcal{X}}$ be a stochastic sampling procedure
from ${\mathcal{Y}}\subset{\mathbb{R}}^{d}$ into
${\mathcal{X}}\subset{\mathbb{R}}^{d}$. Given $y\in{\mathcal{Y}}$, let
$F_{1},\dots,F_{m},F_{m+1}$ be $m+1$ i.i.d. samples from $F(y)$. For a desired
miscoverage level $\alpha\in[0,1]$ and for each $j\in[d]$, let
${\hat{l}}_{j,\alpha},{\hat{u}}_{j,\alpha}$ be the
$\lfloor(m+1)\alpha/2\rfloor/m$ and $\lceil(m+1)(1-\alpha/2)\rceil/m$
entrywise calibrated empirical quantiles of $F_{1},\dots,F_{m}$. Then,

${\mathcal{I}}^{\alpha}(y)_{j}=[{\hat{l}}_{j,\alpha},{\hat{u}}_{j,\alpha}]$
(9)

provides entrywise coverage.

The simple proof of this result is included in Section A.1. We remark that,
analogously to previous works [6, 26], the intervals in
${\mathcal{I}}^{\alpha}(y)$ are _feature-dependent_ and they capture regions
of the image where the sampling process $F(y)$ may have larger uncertainty.
The intervals in ${\mathcal{I}}^{\alpha}(y)$ are statistically valid for any
number of samples $m$ and any distribution ${\mathcal{Q}}_{y}$, i.e. they are
not a heuristic notion of uncertainty. If the sampling procedure $F$ is a
diffusion model, constructing ${\mathcal{I}}^{\alpha}(y)$ is agnostic of the
discretization scheme used to solve the reverse-time SDE [53] and it does not
require retraining the underlying score network, which can be a time-consuming
and delicate process, especially when the size of the images is considerable.
On the other hand, constructing the intervals ${\mathcal{I}}^{\alpha}(y)$
requires sampling a large enough number of times from $F(y)$, which may seen
cumbersome [26]. This is by construction and intention: diffusion models are
indeed very useful in providing good (and varied) samples from the approximate
posterior. In this way, practitioners do typically sample several realizations
to get an empirical study of this distribution. In these settings,
constructing the intervals ${\mathcal{I}}^{\alpha}(y)$ does not involve any
additional computational costs. Furthermore, note that sampling is completely
parallelizable, and so no extra complexity is incurred if a larger number of
computing nodes are available.

###  3.2 A Provable Approach to Optimal Risk Control

In this section, we will revisit the main ideas around conformal risk control
introduced in Section 2.3 and generalize them into our proposed approach,
$K$-RCPS. Naturally, one would like a good conformal risk control procedure to
yield the shortest possible interval lengths. Assume pixel intensities are
normalized between $[0,1]$ and consider the loss function

$\ell^{01}(x,{\mathcal{I}}(y))=\frac{1}{d}\sum_{j\in[d]}{\mathbbm{1}}[x_{j}\notin{\mathcal{I}}(y)_{j}],$
(10)

which counts the (average) number of ground truth pixels that fall outside of
their respective intervals in ${\mathcal{I}}(y)$. The constant set-valued
predictor ${\mathcal{U}}(y)=[0,1]^{d}$ would trivially control the risk, i.e.
$R^{01}(\lambda)={\mathbb{E}}[\ell^{01}(x,{\mathcal{U}}(y))]=0$. Alas, such a
predictor would be completely uninformative. Instead, let
$\\{{\mathcal{I}}_{\lambda}(y)\\}_{\lambda\in\Lambda}$,
$\Lambda\subset\overline{{\mathbb{R}}}$ be a family of predictors that
satisfies the nesting property in Equation 5. In particular, we propose the
following additive parametrization in $\lambda$

${\mathcal{I}}_{\lambda}(y)_{j}=[{\hat{l}}_{j}-\lambda,{\hat{u}}_{j}+\lambda]$
(11)

for some lower and upper endpoints ${\hat{l}}_{j}<{\hat{u}}_{j}$ that may
depend on $y$. For this particularly chosen family of nested predictors, it
follows that the mean interval length is

$\displaystyle\bar{I}(\lambda)=\frac{1}{d}\sum_{j\in[d]}({\hat{u}}_{j}-{\hat{l}}_{j})+2\lambda,$
(12)

a linear function of $\lambda$. Moreover, we can instantiate ${\hat{l}}_{j}$
and ${\hat{u}}_{j}$ to be the calibrated quantiles with entrywise coverage,
i.e.
${\mathcal{I}}^{\alpha}_{\lambda}(y)=[{\hat{l}}_{j,\alpha}-\lambda,{\hat{u}}_{j,\alpha}+\lambda]$.

For such a class of predictors—since the $\ell^{01}$ loss is monotonically
nonincreasing—the original RCPS procedure (see Equation 7) is equivalent to
the following constrained optimization problem

$\displaystyle\hat{\lambda}$
$\displaystyle=\operatorname*{\arg\,\min}_{\lambda\in\Lambda}~{}\bar{I}(\lambda)\quad{\text{s.t.}}\quad\hat{R}^{01+}(\lambda^{\prime})<\epsilon,~{}\forall\lambda^{\prime}\geq\lambda$
(P1)

which naturally minimizes $\lambda$. However, optimizing the mean interval
length over a single scalar parameter $\lambda$ is suboptimal in general, as
shown in Figure 1. With abuse of notation—we do not generally refer to vectors
with boldface—let
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)\\}_{{\bm{\lambda}}\in\Lambda^{d}}$ be a
family of predictors indexed by a $d$-dimensional _vector_
${\bm{\lambda}}=(\lambda_{1},\dots,\lambda_{d})$ that satisfies the nesting
property in Equation 5 in an entrywise fashion. A natural extension of
Equation 11 is then

${\mathcal{I}}_{{\bm{\lambda}}}(y)_{j}=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}],$
(13)

from which one can define an equivalent function $\bar{I}({\bm{\lambda}})$. In
particular, using the calibrated intervals as before, define

${\mathcal{I}}^{\alpha}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j,\alpha}-\lambda_{j},{\hat{u}}_{j,\alpha}+\lambda_{j}].$
(14)

Note now that $\ell^{01}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$ is entrywise
monotonically nonincreasing. For fixed vectors
$\widetilde{{\bm{\lambda}}}\in\mathbb{R}^{d}$ and
${\bm{\eta}}\in{\mathbb{R}}^{d}$ in the positive orthant (i.e.,
${\bm{\eta}}\geq 0$, entrywise), denote
${\bm{\tilde{\Lambda}}}=\widetilde{{\bm{\lambda}}}+\omega{\bm{\eta}},~{}\omega\in\Lambda$
the set of points on a line with offset $\widetilde{{\bm{\lambda}}}$ and
direction ${\bm{\eta}}$ parametrized by $\omega$. Then, the $d$-dimensional
extension of (P1) becomes

$\displaystyle\hat{{\bm{\lambda}}}$
$\displaystyle=\operatorname*{\arg\,\min}_{{\bm{\lambda}}\in{\bm{\tilde{\Lambda}}}}~{}\sum_{j\in[d]}\lambda_{j}\quad{\text{s.t.}}\quad\hat{R}^{01+}({\bm{\lambda}}+\beta{\bm{\eta}})<\epsilon,~{}\forall\beta\geq
0.$ (Pd)

We include an explicit analytical expression for
$\hat{R}^{01+}({\bm{\lambda}}+{\bm{\eta}})$ in Appendix C. Intuitively,
$\hat{{\bm{\lambda}}}$ minimizes the sum of its entries such that the UCB is
smaller than ${\epsilon}$ for all points _to its right_ along the direction of
${\bm{\eta}}$ parametrized by $\beta$. We now show a general high-dimensional
risk control result that holds for any entrywise monotonically nonincreasing
loss function $\ell$ (and not just $\ell^{01}$ as presented in (Pd)) with risk
$R({\bm{\lambda}})$, empirical estimate $\hat{R}({\bm{\lambda}})$ and
respective UCB $\hat{R}^{+}({\bm{\lambda}})$.888The version of Theorem 3.3 in
this manuscript differs from the published paper at
https://proceedings.mlr.press/v202/teneggi23a.html, which has a typo. We
detail the subtle differences and implications in Appendix F.

(a) $\mu=[-1,1]^{T}$. (b) $\mu=[-2,0.75]^{T}$.

Figure 1: Pictorial representation of the suboptimality of the choice of a
single scalar parameter $\lambda$ w.r.t. the mean interval length.
${\mathcal{S}}_{{\text{cal}}}\sim{\mathcal{N}}(\mu,{\mathbb{I}}_{2})^{n}$,
$n=128$, and
$({\mathcal{I}}_{{\bm{\lambda}}})_{j}=[-1-\lambda_{j},1+\lambda_{j}]$,
${\bm{\lambda}}=(\lambda_{1},\lambda_{2})$. For ${\epsilon}=\delta=0.1$,
$\hat{R}^{01+}({\bm{\lambda}})$ is obtained via Hoeffding-Bentkus
hybridization. Green areas indicate regions where
$\hat{R}^{01+}({\bm{\lambda}})\leq{\epsilon}$, and conversely for red regions.
1(a) Shows that when features are concentrated symmetrically around the
intervals, minimizing $\lambda_{1}=\lambda_{2}=\lambda$ (blue star) minimizes
the mean interval length, while 1(b) shows that in the general case, the
optimal ${\bm{\lambda}}$ (orange star) may have $\lambda_{1}\neq\lambda_{2}$.
$\Delta$ highlights the gain in mean interval length obtained by choosing the
orange star instead of the blue one.

######  Theorem 3.3 (Optimal mean interval length risk control).

Let $\ell:~{}{\mathcal{X}}\times{\mathcal{X}}^{\prime}\to{\mathbb{R}}$,
${\mathcal{X}}^{\prime}\subseteq 2^{{\mathcal{X}}}$,
${\mathcal{X}}\subset{\mathbb{R}}^{d}$ be an entrywise monotonically
nonincreasing function and let
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]\\}_{{\bm{\lambda}}\in\Lambda^{d}}$
be a family of set-valued predictors
${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}$,
${\mathcal{Y}}\subset{\mathbb{R}}^{d}$ indexed by
${\bm{\lambda}}\in\Lambda^{d}$, $\Lambda\subset\overline{{\mathbb{R}}}$, for
some lower and upper bounds ${\hat{l}}_{j}<{\hat{u}}_{j}$ that may depend on
$y$. Given
${\bm{\tilde{\Lambda}}}=\widetilde{{\bm{\lambda}}}+\omega{\bm{\eta}},~{}\omega\in\Lambda$,
for fixed $\widetilde{{\bm{\lambda}}}\in{\mathbb{R}}^{d}$ and
${\bm{\eta}}\in{\mathbb{R}}^{d},~{}{\bm{\eta}}\geq 0$, if

$\displaystyle\hat{{\bm{\lambda}}}$
$\displaystyle=\operatorname*{\arg\,\min}_{{\bm{\lambda}}\in{\bm{\tilde{\Lambda}}}}~{}\sum_{j\in[d]}\lambda_{j}\quad{\text{s.t.}}\quad\hat{R}^{+}({\bm{\lambda}}+\beta{\bm{\eta}})<{\epsilon},~{}\forall\beta\geq
0$ (15)

then ${\mathcal{I}}_{\hat{{\bm{\lambda}}}}(y)$ is an
$({\epsilon},\delta)$-RCPS and $\hat{{\bm{\lambda}}}$ minimizes the mean
interval length over ${\bm{\tilde{\Lambda}}}$.

The proof is included in Section A.2. Since $\ell^{01}$ is entrywise
monotonically nonincreasing, it follows that the solution to (Pd) controls
risk. The attentive reader will have noticed (as shown in Figure 1) that the
constraint set $\hat{R}^{01+}({\bm{\lambda}})\leq{\epsilon}$ need not be
convex. Furthermore, and as shown in Figure 2(b), $\ell^{01}$ is not convex in
${\bm{\lambda}}$. Hence, it is not possible to optimally solve (Pd) directly.
Instead, we relax it to a convex optimization problem by means of a convex
upper bound999We defer proofs to Section A.3.

$\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))=\frac{1}{d}\sum_{j\in[d]}\left[\frac{2(1+q)}{I({\bm{\lambda}})_{j}}\lvert
x_{j}-c_{j}\rvert-q\right]_{+},$ (16)

where $q=\gamma/(1-\gamma)$, $\gamma\in[0,1)$,
$I({\bm{\lambda}})_{j}={\hat{u}}_{j}-{\hat{l}}_{j}+2\lambda_{j}$,
$c_{j}=({\hat{u}}_{j}+{\hat{l}}_{j})/2$, and $[\cdot]_{+}=\max(0,\cdot)$. As
shown in Figure 2(a), the hyperparameter $\gamma$ controls the degree of
relaxation by means of changing the portion of the intervals
$[{\hat{l}}_{j},{\hat{u}}_{j}]$ where the loss is 0. This way, $\gamma=0$
retrieves the $\ell_{1}$ loss centered at $c_{j}$, and $\lim_{\gamma\to
1}\ell^{\gamma}=\infty$ if $\exists
j\in[d]:~{}x_{j}\notin[{\hat{l}}_{j},{\hat{u}}_{j}]$ and 0 otherwise.

(a) $\ell^{01},\ell^{\gamma}$ as a function of $x$. (b)
$\ell^{01},\ell^{\gamma}$ as a function of $\lambda$.

Figure 2: Visualization of $\ell^{01}(x,{\mathcal{I}}_{\lambda}(y))$ and
$\ell^{\gamma}(x,{\mathcal{I}}_{\lambda}(y))$ for
${\mathcal{I}}_{\lambda}(y)=[0.50-\lambda,1.50+\lambda]$,
$\gamma\in\\{0,0.5,0.9\\}$. In 2(a) $\lambda=0$, and in 2(b) $x=1.6$.

While one can readily propose a convex alternative to (Pd) by means of this
new loss, we instead propose a generalization of this idea in our final
problem formulation

$\displaystyle\widetilde{{\bm{\lambda}}}_{K}$
$\displaystyle=\operatorname*{\arg\,\min}_{{\bm{\lambda}}\in\Lambda^{K}}~{}\sum_{k\in[K]}n_{k}\lambda_{k}\quad{\text{s.t.}}\quad\hat{R}^{\gamma}(M{\bm{\lambda}})\leq{\epsilon},$
(PK)

for any user-defined $K$-partition of the $[d]$ features—which can be
identified by a membership matrix $M\in\\{0,1\\}^{d\times K}$ where each
feature belongs to (only) one of the $K$ groups with
$n_{k}\coloneqq\lvert\\{j\in[d]:~{}M_{jk}=1\\}\rvert$,
$\sum_{k\in[K]}n_{k}=d$. As we will shortly see, it will be useful to define
these groups as the empirical quantiles; i.e. set $M$ as that assigning each
feature to their respective $k^{\text{th}}$ quantile of the entrywise
empirical loss over the optimization set (which is a vector in
${\mathbb{R}}^{d}$). We remark that the constrain set in (PK) is defined on
the empirical estimate of the risk of ${\mathcal{I}}_{{\bm{\lambda}}}(y)$ and
it does not involve the computation of the UCB. Then, (PK) can be solved with
any standard off-the-shelf convex optimization software (e.g., CVXPY [19, 1],
MOSEK [7]).

Our novel conformal risk control procedure, $K$-RCPS, finds a vector
$\hat{{\bm{\lambda}}}_{K}$ that approximates a solution to the nonconvex
optimization problem (Pd) via a two step procedure:

  1. 1.

First obtaining the optimal solution $\widetilde{{\bm{\lambda}}}_{K}$ to a
user-defined (PK) problem, and then

  2. 2.

Choosing $\hat{\beta}\in\Lambda$ such that

$\hat{\beta}=\inf\\{\beta\in\Lambda:~{}\hat{R}^{01+}(M\widetilde{{\bm{\lambda}}}_{K}+\beta^{\prime}{\bm{1}})<{\epsilon},~{}\forall\beta^{\prime}\geq\beta\\}$

and return
$\hat{{\bm{\lambda}}}_{K}=M\widetilde{{\bm{\lambda}}}_{K}+\hat{\beta}{\bm{1}}$.

Intuitively, the $K$-RCPS algorithm is equivalent to performing the original
RCPS procedure along the line $M\widetilde{{\bm{\lambda}}}_{K}+\beta{\bm{1}}$
parametrized by $\beta$. We remark that—as noted in Theorem 3.3—any choice of
${\bm{\eta}}\geq 0$ provides a valid direction along which to perform the RCPS
procedure. Here, we choose ${\bm{1}}$ because it is precisely the gradient of
the objective function. Future work entails devising more sophisticated
algorithms to approximate the solution of (Pd).

Algorithm 1 $K$-RCPS

1: Input: risk level ${\epsilon}\geq 0$, failure probability $\delta\in[0,1]$,
calibration set ${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}$
of $n$ i.i.d. samples such that $n=n_{{\text{opt}}}+n_{{\text{RCPS}}}$,
membership function $\mathcal{M}$, family of set-valued predictors
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j}-\lambda_{j},u_{j}+\lambda_{j}]\\}_{{\bm{\lambda}}\in\Lambda^{d}}$,
initial (large) value $\beta_{\max}$, stepsize ${\text{d}}{\beta}>0$.

2: Split ${\mathcal{S}}_{{\text{cal}}}$ into
${\mathcal{S}}_{{\text{opt}}},{\mathcal{S}}_{{\text{RCPS}}}$

3: $M\leftarrow\mathcal{M}({\mathcal{S}}_{{\text{opt}}})$

4: $\widetilde{{\bm{\lambda}}}_{K}\leftarrow\texttt{SOLVE-
PK}({\mathcal{S}}_{{\text{opt}}},M)$

5: ${\bm{\lambda}}\leftarrow
M\widetilde{{\bm{\lambda}}}_{K}+\beta_{\max}{\bm{1}}$

6: $\hat{R}^{01+}({\bm{\lambda}})\leftarrow 0$

7: while $\hat{R}^{01+}({\bm{\lambda}})\leq{\epsilon}$ do

8: ${\bm{\lambda}}_{\text{prev}}\leftarrow{\bm{\lambda}}$

9: ${\bm{\lambda}}\leftarrow{\bm{\lambda}}-({\text{d}}{\beta}){\bm{1}}$

10: ${\bm{\lambda}}\leftarrow[{\bm{\lambda}}]_{+}$

11: $\hat{R}^{01}({\bm{\lambda}})\leftarrow
1/n_{{\text{RCPS}}}\cdot\sum_{(x_{i},y_{i})\in{\mathcal{S}}_{{\text{RCPS}}}}\ell^{01}(x_{i},{\mathcal{I}}_{{\bm{\lambda}}}(y_{i}))$

12:
$\hat{R}^{01+}({\bm{\lambda}})\leftarrow\text{UCB}(n_{{\text{RCPS}}},\delta,\hat{R}^{01}({\bm{\lambda}}))$

13: end while

14: $\hat{{\bm{\lambda}}}_{K}\leftarrow{\bm{\lambda}}_{\text{prev}}$

15: return $\hat{{\bm{\lambda}}}_{K}$

Algorithm 1 implements the $K$-RCPS procedure for any calibration set
${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}$, any general
family of set-valued predictors of the form
$\\{{\mathcal{I}}_{{\bm{\lambda}}}=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]\\}_{{\bm{\lambda}}\in\Lambda^{d}}$,
any membership function
$\mathcal{M}:~{}\\{{\mathcal{X}}\times{\mathcal{Y}}\\}^{n}\to\\{0,1\\}^{d\times
K}$, and a general ${\text{UCB}}(n,\delta,\hat{R}({\bm{\lambda}}))$ that
accepts a number of samples $n$, a failure probability $\delta$, and an
empirical risk $\hat{R}({\bm{\lambda}})$ and that returns a pointwise upper
confidence bound $\hat{R}^{+}({\bm{\lambda}})$ that satisfies Equation 6. We
remark that, following the _split fixed sequence testing_ framework introduced
in Angelopoulos et al. [4] and applied in previous work [36], the membership
matrix and its optimization problem (PK) are computed on a subset
${\mathcal{S}}_{{\text{opt}}}$ of the calibration set
${\mathcal{S}}_{{\text{cal}}}$, such that the direction
$M\widetilde{{\bm{\lambda}}}_{K}+\beta{\bm{1}}$ along which to perform the
RCPS procedure is chosen _before_ looking at the data
${\mathcal{S}}_{{\text{RCPS}}}={\mathcal{S}}_{{\text{cal}}}\setminus{\mathcal{S}}_{{\text{opt}}}$.
We note that $K$-RCPS allows for some of the entries in
$\hat{{\bm{\lambda}}}_{K}$ to be set to 0, which preserves the original
intervals such that—if they are obtained as described in Section 3.1—they
still provide entrywise coverage of future samples at the desired level
$\alpha$.

We now move onto showcasing the advantage of $K$-RCPS in terms on mean
interval length on two real-world high dimensional denoising problems: one on
natural images of faces as well as on CT scans of the abdomen.

##  4 Experiments

(a) CelebA dataset, $\sigma_{0}^{2}=1.0,~{}\alpha=0.10$. (b) AbdomenCT-1K
dataset, $\sigma_{0}^{2}=0.4,~{}\alpha=0.20$.

Figure 3: Calibrated quantiles ${\mathcal{I}}^{\alpha}(y)$ computed on 128
samples from $F(y)$ for noisy inputs $y$ with noise level $\sigma_{0}^{2}$.
The difference ${\hat{u}}_{\alpha}-{\hat{l}}_{\alpha}$ represents intervals
sizes (i.e., larger intervals indicate larger uncertainty).

As a reminder, the methodological contribution of this paper is two-fold:
$(i)$ we propose to use the calibrated quantiles ${\mathcal{I}}^{\alpha}(y)$
as a statistically valid notion of uncertainty for diffusion models, and
$(ii)$ we introduce the $K$-RCPS procedure to guarantee high-dimensional risk
control. Although it is natural to use the two in conjunction, we remark that
the $K$-RCPS procedure is agnostic of the notion of uncertainty and it can be
applied to any nested family of set predictors
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)\\}_{{\bm{\lambda}}\in\Lambda^{d}}$ that
satisfy the additive parametrization in Equation 13. Therefore, we compare
$K$-RCPS with the original RCPS algorithm on several baseline notions of
uncertainty: quantile regression [6], MC-Dropout [20], N-Con$ffusion$ [26],
and naive (i.e., not calibrated) quantiles. We focus on denoising problems
where $y=x+v_{0}$ with $v_{0}\sim{\mathcal{N}}(0,\sigma_{0}^{2})$, on two
imaging datasets: the CelebA dataset [40] and the AbdomenCT-1K dataset [41].
In particular—for each dataset—we train:

  * •

A time-conditional score network $s(\tilde{x},t)\approx{\nabla}_{x}\log
p_{t}(\tilde{x})$ following Song et al. [53] to sample from the posterior
distribution $p(x|y)$ as described in Section 2.1, and

  * •

a time-conditional image regressor
$f:~{}{\mathcal{Y}}\times{\mathbb{R}}\to{\mathcal{X}}^{3}$ following
Angelopoulos et al. [6] such that
$f(y,t)=({\hat{q}}_{\alpha/2},\hat{x},{\hat{q}}_{1-\alpha/2})$, where
$\hat{x}\approx{\mathbb{E}}[x\mid y]$ minimizes the MSE loss between the noisy
observation $y$ and the ground truth $x$, and
${\hat{q}}_{\alpha/2},{\hat{q}}_{1-\alpha/2}$ are the $\alpha/2$ and
$1-\alpha/2$ quantile regressors of $x$, respectively [34, 45, 6].

Both models are composed of the same NCSN++ backbone [53] with dropout
$p=0.10$ for a fair comparison. We then fine-tune the original score network
$s(\tilde{x},t)$ according to the N-Con$ffusion$ algorithm proposed by Horwitz
and Hoshen [26] such that—similarly to the image regressor $f$—the resulting
time-conditional predictor
$\tilde{s}(y,t)=({\hat{q}}_{\alpha/2},{\hat{q}}_{1-\alpha/2})$ estimates the
$\alpha/2$ and $1-\alpha/2$ quantile regressors of $x$. Finally, in order to
compare with MC-Dropout, we activate the dropout layers in the image regressor
$f$ at inference time, and estimate the mean $\bar{x}$ and standard deviation
$\hat{\sigma}$ over 128 samples $\hat{x}_{1},\dots,\hat{x}_{128}$. To
summarize, we compare $K$-RCPS and RCPS on the following families of nested
set predictors:

Quantile Regression (QR)

${\mathcal{I}}_{{\bm{\lambda}},\text{QR}}(y)_{j}=[\hat{x}_{j}-\lambda_{j}({\hat{q}}_{\alpha/2})_{j},\hat{x}_{j}+\lambda_{j}({\hat{q}}_{1-\alpha/2})_{j}],$
(17)

where $f(y,t)=({\hat{q}}_{\alpha/2},\hat{x},{\hat{q}}_{1-\alpha/2})$.

MC-Dropout

${\mathcal{I}}_{{\bm{\lambda}},\text{MC-
Dropout}}(y)_{j}=[\bar{x}_{j}-\lambda_{j}\hat{\sigma}_{j},\bar{x}_{j}+\lambda_{j}\hat{\sigma}_{j}],$
(18)

where $\bar{x},\hat{\sigma}$ are the sample mean and standard deviation over
128 samples $\hat{x}_{1},\dots,\hat{x}_{128}$ obtained by activating the
dropout layers in the image regressor $f$.

N-Con$\bm{ffusion}$101010A detailed discussion of Con$ffusion$ and the
contributions of this paper is included in Appendix E. We compare two
different parametrizations—multiplicative and additive:

${\mathcal{I}}^{~{}\text{multiplicative}}_{{\bm{\lambda}},\text{Conffusion}}(y)_{j}=\left[\frac{({\hat{q}}_{\alpha/2})_{j}}{\lambda_{j}},\lambda_{j}({\hat{q}}_{1-\alpha/2})_{j}\right]\\\
$ (19)

and

${\mathcal{I}}^{~{}\text{additive}}_{{\bm{\lambda}},\text{Conffusion}}(y)_{j}=\left[({\hat{q}}_{\alpha/2})_{j}-\lambda_{j},({\hat{q}}_{1-\alpha/2})_{j}+\lambda_{j}\right],$
(20)

where $\tilde{s}(y,t)=({\hat{q}}_{\alpha/2},{\hat{q}}_{1-\alpha/2})$ is the
fine-tuned score network by means of quantile regression on 1000 additional
samples.

Naive quantiles

${\mathcal{I}}_{{\bm{\lambda}},\text{naive}}(y)_{j}=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}],$
(21)

where ${\hat{l}},{\hat{u}}$ are the naive (i.e., not calibrated) $\alpha/2$
and $1-\alpha/2$ entrywise empirical quantiles computed on 128 samples from
the diffusion model.

Calibrated quantiles

${\mathcal{I}}^{\alpha}_{{\bm{\lambda}}}(y)_{j}=[{\hat{l}}_{j,\alpha}-\lambda_{j},{\hat{u}}_{j,\alpha}+\lambda_{j}]$
(22)

where ${\hat{l}}_{\alpha},{\hat{u}}_{\alpha}$ are the entrywise calibrated
quantiles computed on 128 samples from the diffusion model as described in
Equation 9 (see Figure 3 for some examples).

(a) CelebA dataset ($\alpha=0.10$, ${\epsilon}=0.10$). (b) AbdomenCT-1K
dataset ($\alpha=0.20$, ${\epsilon}=0.05$).

Figure 4: Example optimal $\hat{{\bm{\lambda}}}_{K}$ for $K\in\\{4,8,32\\}$,
$n_{{\text{opt}}}=256$, and $d_{{\text{opt}}}=100$ with respective
conformalized uncertainty maps
${\mathcal{I}}_{\hat{{\bm{\lambda}}}}^{\alpha}(y)=[{\hat{l}}_{j,\alpha}-(\hat{\lambda}_{K})_{j},{\hat{u}}_{j,\alpha}+(\hat{\lambda}_{K})_{j}]$.
With probability at least $90\%$ no more than ${\epsilon}$ portion of the
ground truth pixels will fall outside of
${\mathcal{I}}^{\alpha}_{\hat{{\bm{\lambda}}}}$ on future, unseen samples.

We include further details on the datasets, the models, and the training and
sampling procedures in Appendix D. The implementation of $K$-RCPS with all
code and data necessary to reproduce the experiments is available at
https://github.com/Sulam-Group/k-rcps.

We compare all models and calibration procedures on 20 random draws of
calibration and validation sets
${\mathcal{S}}_{{\text{cal}}},{\mathcal{S}}_{{\text{val}}}$ of length
$n_{{\text{cal}}}$ and $n_{{\text{val}}}$, respectively. We remark that for
the $K$-RCPS procedure, $n_{{\text{opt}}}$ samples from
${\mathcal{S}}_{{\text{cal}}}$ will be used to solve the optimization problem
(PK). It follows that for a fixed $n_{{\text{cal}}}$, the concentration
inequality used in the $K$-RCPS procedure will be looser compared to the one
in the RCPS algorithm. We will show that there remains a clear benefit of
using the $K$-RCPS algorithm in terms of mean interval length given the same
amount of calibration data available (i.e., even while the concentration bound
becomes looser). In these experiments, we construct the membership matrix $M$
by assigning each feature $j\in[d]$ to the respective $k^{\text{th}}$,
$k=1,\dots,K$ quantile of the entrywise empirical estimate of the risk on
${\mathcal{S}}_{{\text{opt}}}$. Furthermore, even though (PK) is low-
dimensional (i.e., $K\ll d$), the number of constraints grows as
$dn_{{\text{opt}}}$, which quickly makes the computation of
$\widetilde{{\bm{\lambda}}}_{K}$ inefficient and time-consuming (e.g., for the
AbdomenCT-1K dataset, $dn_{{\text{opt}}}\sim 10^{8}$ when
$n_{{\text{opt}}}=128$, a mild number of samples to optimize over). In
practice, we randomly subsample a small number of features
$d_{{\text{opt}}}\ll d$ stratified by membership, which drastically speeds up
computation. In the following experiments, we set
$d_{{\text{opt}}}\in\\{50,100\\}$ and use $K\in\\{4,8,32\\}$, which reduces
the runtime of solving the optimization problem to less than a second for both
datasets. Finally, we pick $\gamma$ that minimizes the objective function over
16 values equally spaced in $[0.3,0.7]$. The choice of these heuristics makes
the runtime of $K$-RCPS comparable to that of RCPS, with a small overhead to
solve the reduced (PK) problem (potentially multiple times to optimize
$\gamma$).

Figure G.1 shows that all combinations of notion of uncertainty and
calibration procedures control the risk, as promised. In particular, we set
$\delta=0.10$ for both datasets, and ${\epsilon}=0.10,0.05$ for the CelebA and
AbdomenCT-1K dataset, respectively. We repeat all calibration procedures over
20 random samples of
${\mathcal{S}}_{{\text{cal}}},{\mathcal{S}}_{{\text{val}}}$, with
$n_{{\text{val}}}=128$, and $n_{{\text{cal}}}=640$ or $n_{{\text{cal}}}=512$
for the CelebA or AbdomenCT-1K dataset, respectively. Figure 4 showcases some
example $\hat{{\bm{\lambda}}}_{K}$’s obtained by running the $K$-RCPS
procedure with $K=4,8$, and 32 quantiles alongside their respective
conformalized uncertainty maps from ${\mathcal{S}}_{{\text{val}}}$. We can
appreciate how for both datasets, $\hat{{\bm{\lambda}}}_{K}$ captures
information about the structure of the data distribution (e.g., eyes and lips
for the CelebA dataset, and the position of lungs and the heart for the
AbdomenCT-1K dataset). Finally, we compare all baselines and calibration
procedures in terms of the guarantees each of them provide and their mean
interval length. In particular, we report whether each notion of uncertainty
provides guarantees over a diffusion model or not. Note that naive and
calibrated quantiles are the only notions of uncertainty that precisely
provide guarantees on the samples from a diffusion model. Furthermore,
calibrated quantiles are the only method that ensures entrywise coverage on
future samples on the same noisy observation. For $K$-RCPS, we perform a grid
search over $n_{{\text{opt}}}\in\\{128,256\\}$,
$d_{{\text{opt}}}\in\\{50,100\\}$, and $K\in\\{4,8,32\\}$, and we report the
optimal results in Table 1. For both datasets, $K$-RCPS provides the tightest
intervals among methods that provide both entrywise coverage and risk control
for diffusion models. When relaxing the constraint of entrywise coverage,
$K$-RCPS still provides the tightest intervals. Across the uncertainty
quantification methods that do not relate to a diffusion model, we found that
$K$-RCPS with naive sampling provides better results on the CelebA dataset.
For the AbdomenCT-1K dataset, N-Con$ffusion$ with multiplicative
parametrization and RCPS provides slightly shorter intervals compared to
Con$ffusion$ with additive parametrization and $K$-RCPS. However, we stress
that the intervals provided by N-Con$ffusion$ are computed for a fine-tuned
model that is different from the original diffusion model, and thus provide no
guarantees over the samples of the original model. We found that
N-Con$ffusion$ with multiplicative parametrization underperforms on the CelebA
dataset because the lower bounds do not decay fast enough, and the loss is
concentrated on features whose $({\hat{q}}_{\alpha/2})_{j}/\lambda_{j}>x_{j}$.

##  5 Conclusions

Table 1: Comparison of all notions of uncertainty with RCPS and $K$-RCPS in
terms of guarantees provided and mean interval length over 20 independent
draws of ${\mathcal{S}}_{{\text{cal}}}$. We refer the reader to Appendix E for
a detailed discussion of the comparison with the Con$ffusion$ framework.

Uncertainty |  Diffusion model? |  Entrywise coverage? |  Risk control? |  Calibration procedure | Mean interval length
---|---|---|---|---|---
CelebA | AbdomenCT-1K
QR | ✗ | ✗ | ✓ | RCPS | $0.4843\pm 0.0121$ | $0.2943\pm 0.0060$
MC-Dropout | ✗ | ✗ | ✓ | RCPS | $0.6314\pm 0.0109$ | $0.2810\pm 0.0013$
N-Con$ffusion$ |  |  |  |  |  |
— multiplicative | ✗ | ✗ | ✓ | RCPS | $0.6949\pm 0.0084$ | $0.1126\pm 0.0020$
— additive | ✗ | ✗ | ✓ | RCPS | $0.3314\pm 0.0040$ | $0.1164\pm 0.0024$
— additive | ✗ | ✗ | ✓ |  $K$-RCPS | $0.3131\pm 0.0056$ | $0.1136\pm 0.0019$
Naive quantiles | ✓ | ✗ | ✓ | RCPS | $0.2688\pm 0.0068$ | $0.1518\pm 0.0016$
Naive quantiles | ✓ | ✗ | ✓ |  $K$-RCPS | $0.2523\pm 0.0052$ | $0.1374\pm 0.0019$
Calibrated quantiles | ✓ | ✓ | ✓ | RCPS | $0.2762\pm 0.0059$ | $0.1506\pm 0.0014$
Calibrated quantiles | ✓ | ✓ | ✓ |  $K$-RCPS | $0.2644\pm 0.0067$ | $0.1369\pm 0.0016$

Diffusion models represent huge potential for sampling in inverse problems,
alas how to devise precise guarantees on uncertainty has remained open. We
have provided _(i)_ calibrated intervals that guarantee coverage of future
samples generated by diffusion models, (ii) shown how to extend RCPS to
$K$-RCPS, allowing for greater flexibility by conformalizing in higher
dimensions by means of a convex surrogate problem. Yet, our results are
general and hold for any data distribution and any sampling
procedure—diffusion models or otherwise. When combined, these two
contributions provide state of the art uncertainty quantification by
controlling risk with minimal mean interval length. Our contributions open the
door to a variety of new problems. While we have focused on denoising
problems, the application of these tools for other, more challenging
restoration tasks is almost direct since no distributional assumptions are
employed. The variety of diffusion models for other conditional-sampling
problem can readily be applied here too [61, 17]. Lastly—and differently from
other works that explore controlling multiple risks [36]—ours is the first
approach to provide multi-dimensional control of one risk for conformal
prediction, and likely improvements to our optimization schemes could be
possible. More generally, we envision our tools to contribute to the
responsible use of machine learning in modern settings.

## Acknowledgments

The authors sincerely thank Anastasios Angelopoulos for insightful discussions
in early stages of this work, as well as the anonymous reviewers who have
helped improve our paper.

## References

  * Agrawal et al. [2018] Akshay Agrawal, Robin Verschueren, Steven Diamond, and Stephen Boyd.  A rewriting system for convex optimization problems.  _Journal of Control and Decision_ , 5(1):42–60, 2018.
  * Anderson [1982] Brian DO Anderson.  Reverse-time diffusion equation models.  _Stochastic Processes and their Applications_ , 12(3):313–326, 1982.
  * Angelopoulos and Bates [2021] Anastasios N Angelopoulos and Stephen Bates.  A gentle introduction to conformal prediction and distribution-free uncertainty quantification.  _arXiv preprint arXiv:2107.07511_ , 2021.
  * Angelopoulos et al. [2021] Anastasios N Angelopoulos, Stephen Bates, Emmanuel J Candès, Michael I Jordan, and Lihua Lei.  Learn then test: Calibrating predictive algorithms to achieve risk control.  _arXiv preprint arXiv:2110.01052_ , 2021.
  * Angelopoulos et al. [2022a] Anastasios N Angelopoulos, Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster.  Conformal risk control.  _arXiv preprint arXiv:2208.02814_ , 2022a.
  * Angelopoulos et al. [2022b] Anastasios N Angelopoulos, Amit Pal Kohli, Stephen Bates, Michael Jordan, Jitendra Malik, Thayer Alshaabi, Srigokul Upadhyayula, and Yaniv Romano.  Image-to-image regression with distribution-free uncertainty quantification and applications in imaging.  In _International Conference on Machine Learning_ , pages 717–730. PMLR, 2022b.
  * ApS [2019] MOSEK ApS.  _The MOSEK optimization toolbox for MATLAB manual. Version 9.0._ , 2019.  URL http://docs.mosek.com/9.0/toolbox/index.html.
  * Barber et al. [2021] Rina Foygel Barber, Emmanuel J Candes, Aaditya Ramdas, and Ryan J Tibshirani.  Predictive inference with the jackknife+.  _The Annals of Statistics_ , 49(1):486–507, 2021\.
  * Barber et al. [2022] Rina Foygel Barber, Emmanuel J Candes, Aaditya Ramdas, and Ryan J Tibshirani.  Conformal prediction beyond exchangeability.  _arXiv preprint arXiv:2202.13415_ , 2022.
  * Bates et al. [2021] Stephen Bates, Anastasios Angelopoulos, Lihua Lei, Jitendra Malik, and Michael Jordan.  Distribution-free, risk-controlling prediction sets.  _Journal of the ACM (JACM)_ , 68(6):1–34, 2021\.
  * Bentkus [2004] Vidmantas Bentkus.  On hoeffding’s inequalities.  _The Annals of Probability_ , 32(2):1650–1673, 2004.
  * Bishop and Nasrabadi [2006] Christopher M Bishop and Nasser M Nasrabadi.  _Pattern recognition and machine learning_ , volume 4.  Springer, 2006.
  * Carlier et al. [2016] Guillaume Carlier, Victor Chernozhukov, and Alfred Galichon.  Vector quantile regression: an optimal transport approach.  _The Annals of Statistics_ , 44(3):1165–1192, 2016.
  * Carlier et al. [2020] Guillaume Carlier, Victor Chernozhukov, Gwendoline De Bie, and Alfred Galichon.  Vector quantile regression and optimal transport, from theory to numerics.  _Empirical Economics_ , pages 1–28, 2020.
  * Chernozhukov et al. [2017] Victor Chernozhukov, Alfred Galichon, Marc Hallin, and Marc Henry.  Monge–kantorovich depth, quantiles, ranks and signs.  _The Annals of Statistics_ , 45(1):223–256, 2017\.
  * Corso et al. [2022] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola.  Diffdock: Diffusion steps, twists, and turns for molecular docking.  _arXiv preprint arXiv:2210.01776_ , 2022.
  * Croitoru et al. [2022] Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor Ionescu, and Mubarak Shah.  Diffusion models in vision: A survey.  _arXiv preprint arXiv:2209.04747_ , 2022.
  * Croitoru et al. [2023] Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor Ionescu, and Mubarak Shah.  Diffusion models in vision: A survey.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2023.
  * Diamond and Boyd [2016] Steven Diamond and Stephen Boyd.  CVXPY: A Python-embedded modeling language for convex optimization.  _Journal of Machine Learning Research_ , 17(83):1–5, 2016.
  * Gal and Ghahramani [2016] Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In _international conference on machine learning_ , pages 1050–1059. PMLR, 2016.
  * Genevay et al. [2016] Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach.  Stochastic optimization for large-scale optimal transport.  _Advances in neural information processing systems_ , 29, 2016.
  * Gupta et al. [2022] Chirag Gupta, Arun K Kuchibhotla, and Aaditya Ramdas.  Nested conformal prediction and quantile out-of-bag ensemble methods.  _Pattern Recognition_ , 127:108496, 2022.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Hoeffding [1994] Wassily Hoeffding.  Probability inequalities for sums of bounded random variables.  In _The collected works of Wassily Hoeffding_ , pages 409–426. Springer, 1994.
  * Hoogeboom et al. [2022] Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling.  Equivariant diffusion for molecule generation in 3d.  In _International Conference on Machine Learning_ , pages 8867–8887. PMLR, 2022.
  * Horwitz and Hoshen [2022] Eliahu Horwitz and Yedid Hoshen.  Conffusion: Confidence intervals for diffusion models.  _arXiv preprint arXiv:2211.09795_ , 2022.
  * Hyvärinen and Dayan [2005] Aapo Hyvärinen and Peter Dayan.  Estimation of non-normalized statistical models by score matching.  _Journal of Machine Learning Research_ , 6(4), 2005.
  * Ingraham et al. [2022] John Ingraham, Max Baranov, Zak Costello, Vincent Frappier, Ahmed Ismail, Shan Tie, Wujie Wang, Vincent Xue, Fritz Obermeyer, Andrew Beam, et al.  Illuminating protein space with a programmable generative model.  _bioRxiv_ , 2022.
  * Kadkhodaie and Simoncelli [2021] Zahra Kadkhodaie and Eero Simoncelli.  Stochastic solutions for linear inverse problems using the prior implicit in a denoiser.  _Advances in Neural Information Processing Systems_ , 34:13242–13254, 2021.
  * Karatzas et al. [1991] Ioannis Karatzas, Ioannis Karatzas, Steven Shreve, and Steven E Shreve.  _Brownian motion and stochastic calculus_ , volume 113.  Springer Science & Business Media, 1991.
  * Kawar et al. [2021a] Bahjat Kawar, Gregory Vaksman, and Michael Elad.  Snips: Solving noisy inverse problems stochastically.  _Advances in Neural Information Processing Systems_ , 34:21757–21769, 2021a.
  * Kawar et al. [2021b] Bahjat Kawar, Gregory Vaksman, and Michael Elad.  Stochastic image denoising by sampling from the posterior distribution.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 1866–1875, 2021b.
  * Kazerouni et al. [2022] Amirhossein Kazerouni, Ehsan Khodapanah Aghdam, Moein Heidari, Reza Azad, Mohsen Fayyaz, Ilker Hacihaliloglu, and Dorit Merhof.  Diffusion models for medical image analysis: A comprehensive survey.  _arXiv preprint arXiv:2211.07804_ , 2022.
  * Koenker and Bassett Jr [1978] Roger Koenker and Gilbert Bassett Jr.  Regression quantiles.  _Econometrica: journal of the Econometric Society_ , pages 33–50, 1978.
  * Kutiel et al. [2022] Gilad Kutiel, Regev Cohen, Michael Elad, and Daniel Freedman.  What’s behind the mask: Estimating uncertainty in image-to-image problems.  _arXiv preprint arXiv:2211.15211_ , 2022.
  * Laufer-Goldshtein et al. [2022] Bracha Laufer-Goldshtein, Adam Fisch, Regina Barzilay, and Tommi Jaakkola.  Efficiently controlling multiple risks with pareto testing.  _arXiv preprint arXiv:2210.07913_ , 2022.
  * Lee et al. [2022] Holden Lee, Jianfeng Lu, and Yixin Tan.  Convergence for score-based generative modeling with polynomial complexity.  _arXiv preprint arXiv:2206.06227_ , 2022.
  * Lei and Wasserman [2014] Jing Lei and Larry Wasserman.  Distribution-free prediction bands for non-parametric regression.  _Journal of the Royal Statistical Society: Series B (Statistical Methodology)_ , 76(1):71–96, 2014.
  * Liu et al. [2016] Qiang Liu, Jason Lee, and Michael Jordan.  A kernelized stein discrepancy for goodness-of-fit tests.  In _International conference on machine learning_ , pages 276–284. PMLR, 2016.
  * Liu et al. [2018] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.  Large-scale celebfaces attributes (celeba) dataset.  _Retrieved August_ , 15(2018):11, 2018.
  * Ma et al. [2021] Jun Ma, Yao Zhang, Song Gu, Cheng Zhu, Cheng Ge, Yichi Zhang, Xingle An, Congcong Wang, Qiyuan Wang, Xin Liu, et al.  Abdomenct-1k: Is abdominal organ segmentation a solved problem.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2021.
  * Metzer et al. [2022] Gal Metzer, Elad Richardson, Or Patashnik, Raja Giryes, and Daniel Cohen-Or.  Latent-nerf for shape-guided generation of 3d shapes and textures.  _arXiv preprint arXiv:2211.07600_ , 2022.
  * Pang et al. [2020] Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu.  Efficient learning of generative models via finite-difference score matching.  _Advances in Neural Information Processing Systems_ , 33:19175–19188, 2020.
  * Papadopoulos et al. [2002] Harris Papadopoulos, Kostas Proedrou, Volodya Vovk, and Alex Gammerman.  Inductive confidence machines for regression.  In _European Conference on Machine Learning_ , pages 345–356. Springer, 2002.
  * Romano et al. [2019] Yaniv Romano, Evan Patterson, and Emmanuel Candes.  Conformalized quantile regression.  _Advances in neural information processing systems_ , 32, 2019.
  * Ronneberger et al. [2015] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _International Conference on Medical image computing and computer-assisted intervention_ , pages 234–241. Springer, 2015.
  * Rosenberg et al. [2022] Aviv A Rosenberg, Sanketh Vedula, Yaniv Romano, and Alex M Bronstein.  Fast nonlinear vector quantile regression.  _arXiv preprint arXiv:2205.14977_ , 2022.
  * Shafer and Vovk [2008] Glenn Shafer and Vladimir Vovk.  A tutorial on conformal prediction.  _Journal of Machine Learning Research_ , 9(3), 2008.
  * Sohl-Dickstein et al. [2015] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International Conference on Machine Learning_ , pages 2256–2265. PMLR, 2015.
  * Song and Ermon [2019] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Song and Ermon [2020] Yang Song and Stefano Ermon.  Improved techniques for training score-based generative models.  _Advances in neural information processing systems_ , 33:12438–12448, 2020.
  * Song et al. [2020a] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon.  Sliced score matching: A scalable approach to density and score estimation.  In _Uncertainty in Artificial Intelligence_ , pages 574–584. PMLR, 2020a.
  * Song et al. [2020b] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  _arXiv preprint arXiv:2011.13456_ , 2020b.
  * Song et al. [2021] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon.  Solving inverse problems in medical imaging with score-based generative models.  _arXiv preprint arXiv:2111.08005_ , 2021.
  * Torem et al. [2022] Nadav Torem, Roi Ronen, Yoav Y Schechner, and Michael Elad.  Towards a most probable recovery in optical imaging.  _arXiv preprint arXiv:2212.03235_ , 2022.
  * Vovk [2015] Vladimir Vovk.  Cross-conformal predictors.  _Annals of Mathematics and Artificial Intelligence_ , 74(1):9–28, 2015.
  * Vovk et al. [2005] Vladimir Vovk, Alexander Gammerman, and Glenn Shafer.  _Algorithmic learning in a random world_.  Springer Science & Business Media, 2005.
  * Watson et al. [2022] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al.  Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models.  _bioRxiv_ , 2022.
  * Xie and Li [2022] Yutong Xie and Quanzheng Li.  Measurement-conditioned denoising diffusion probabilistic model for under-sampled medical image reconstruction.  _arXiv preprint arXiv:2203.03623_ , 2022.
  * Xu et al. [2022] Jiale Xu, Xintao Wang, Weihao Cheng, Yan-Pei Cao, Ying Shan, Xiaohu Qie, and Shenghua Gao.  Dream3d: Zero-shot text-to-3d synthesis using 3d shape prior and text-to-image diffusion models.  _arXiv preprint arXiv:2212.14704_ , 2022.
  * Yang et al. [2022] Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Yingxia Shao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang.  Diffusion models: A comprehensive survey of methods and applications.  _arXiv preprint arXiv:2209.00796_ , 2022.
  * Zeng et al. [2022] Xiaohui Zeng, Arash Vahdat, Francis Williams, Zan Gojcic, Or Litany, Sanja Fidler, and Karsten Kreis.  Lion: Latent point diffusion models for 3d shape generation.  _arXiv preprint arXiv:2210.06978_ , 2022.

##  Appendix A Proofs

In this section, we include the proofs for the results presented in this
paper. Herein, denote
${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}$ a set-valued
predictor from ${\mathcal{Y}}\subset{\mathbb{R}}^{d}$ into a space of subsets
${\mathcal{X}}^{\prime}\subseteq 2^{{\mathcal{X}}}$ for
${\mathcal{X}}\subset{\mathbb{R}}^{d}$.

###  A.1 Proof of Lemma 3.2

Let $F:~{}{\mathcal{Y}}\to{\mathcal{X}}$ be a stochastic sampling procedure
from ${\mathcal{Y}}$ to ${\mathcal{X}}$ such that for a fixed
$y\in{\mathcal{Y}}$, $F(y)$ is a random vector with unknown distribution
${\mathcal{Q}}_{y}$. We show that for a desired miscoverage level
$\alpha\in[0,1]$, the entrywise calibrated empirical quantiles
${\mathcal{I}}^{\alpha}(y)_{j}=[{\hat{l}}_{j,\alpha},{\hat{u}}_{j,\alpha}]$
defined in Equation 9 provide entrywise coverage as in Definition 3.1. That
is, for each $j\in[d]\coloneqq\\{1,\dots,d\\}$

${\mathbb{P}}[F(y)_{j}\in{\mathcal{I}}^{\alpha}(y)_{j}]\geq 1-\alpha.$ (23)

###### Proof.

The proof is a variation of the classical split conformal prediction coverage
guarantee (see Angelopoulos and Bates [3], Theorem D.1). Let
$F_{1},\dots,F_{m},F_{m+1}$ be $m+1$ i.i.d. samples from $F(y)$. For a desired
miscoverage level $\alpha\in[0,1]$ and for each $j\in[d]$ denote

${\hat{l}}_{j,\alpha}=\inf\left\\{l:~{}\frac{\lvert\\{k:~{}(F_{k})_{j}\leq
l\\}\rvert}{m}\geq\frac{\lfloor(m+1)\alpha/2\rfloor}{m}\right\\}$ (24)

and

${\hat{u}}_{j,\alpha}=\inf\left\\{u:~{}\frac{\lvert\\{k:~{}(F_{k})_{j}\leq
u\\}\rvert}{m}\geq\frac{\lceil(m+1)(1-\alpha/2)\rceil}{m}\right\\}$ (25)

the $\lfloor(m+1)\alpha/2\rfloor/m$ and $\lceil(m+1)(1-\alpha/2)\rceil/m$
entrywise calibrated empirical quantiles of $F_{1},\dots,F_{m}$. Assume that
for each $j\in[d]$, the first $m$ samples are ordered in ascending order, i.e.
$(F_{1})_{j}<\dots<(F_{m})_{j}$ such that

${\hat{l}}_{j,\alpha}=(F_{\lfloor(m+1)\alpha/2\rfloor})_{j}\quad~{}\text{and}~{}\quad{\hat{u}}_{j,\alpha}=(F_{\lceil(m+1)(1-\alpha/2)\rceil})_{j}.$
(26)

Note that by symmetry of $(F_{1})_{j},\dots,(F_{m})_{j}$ it follows that
$(F_{m+1})_{j}$ is equally likely to fall between any of the first $m$
samples. That is, for any two indices $m_{1}<m_{2}$

${\mathbb{P}}[(F_{m+1})_{j}\in[(F_{m_{1}})_{j},(F_{m_{2}})_{j}]]=\frac{m_{2}-m_{1}}{m+1}.$
(27)

Instantiating the above equality with
${\mathcal{I}}^{\alpha}(y)_{j}=[{\hat{l}}_{j,\alpha},{\hat{u}}_{j,\alpha}]$
yields

$\displaystyle{\mathbb{P}}[(F_{m+1})_{j}\in{\mathcal{I}}^{\alpha}(y)_{j}]$
$\displaystyle={\mathbb{P}}\left[(F_{m+1})_{j}\in[{\hat{l}}_{j,\alpha},{\hat{u}}_{j,\alpha}]\right]$
(28)
$\displaystyle={\mathbb{P}}\left[(F_{m+1})_{j}\in[(F_{\lfloor(m+1)\alpha/2\rfloor})_{j},(F_{\lceil(m+1)(1-\alpha/2)\rceil})_{j}]\right]$
(29)
$\displaystyle=\frac{\lceil(m+1)(1-\alpha/2)\rceil-\lfloor(m+1)\alpha/2\rfloor}{m+1}$
(30) $\displaystyle\geq\frac{(m+1)(1-\alpha)}{m+1}=1-\alpha$ (31)

which concludes the proof. ∎

###  A.2 Proof of Theorem 3.3

Recall that for a calibration set
${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}$ of $n$ i.i.d.
samples from an unknown distribution ${\mathcal{D}}$ over
${\mathcal{X}}\times{\mathcal{Y}}$, a loss function
$\ell:~{}{\mathcal{X}}\times{\mathcal{X}}^{\prime}\to{\mathbb{R}}$, and a
family
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)\\}_{{\bm{\lambda}}\in\Lambda^{d}}$ of
set-valued predictors indexed by a $d$-dimensional vector
${\bm{\lambda}}=(\lambda_{1},\dots,\lambda_{d})\in\Lambda^{d}$,
$\Lambda\subset\overline{{\mathbb{R}}}={\mathbb{R}}\cup\\{\pm\infty\\}$

$R({\bm{\lambda}})={\mathbb{E}}_{(x,y)\sim{\mathcal{D}}}\left[\ell(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))\right]\quad~{}\text{and}~{}\quad\hat{R}({\bm{\lambda}})=\frac{1}{n}\sum_{(x_{i},y_{i})\in{\mathcal{S}}_{{\text{cal}}}}^{n}\ell(x_{i},{\mathcal{I}}_{{\bm{\lambda}}}(y_{i}))$
(32)

denote the risk of ${\mathcal{I}}_{{\bm{\lambda}}}(y)$ and its empirical
estimate on the calibration set, respectively. Furthermore, let
$\hat{R}^{+}({\bm{\lambda}})$ be a pointwise upper confidence bound (UCB) such
that for each fixed ${\bm{\lambda}}\in\Lambda^{d}$ and $\forall\delta\in[0,1]$

${\mathbb{P}}[R({\bm{\lambda}})\leq\hat{R}^{+}({\bm{\lambda}})]\geq 1-\delta$
(33)

as presented in Equation 6. Equivalently to Definition 2.2, for a risk level
${\epsilon}\geq 0$, we say that ${\mathcal{I}}_{{\bm{\lambda}}}(y)$ is an
$({\epsilon},\delta)$-RCPS if

${\mathbb{P}}_{{\mathcal{S}}_{{\text{cal}}}}[R({\bm{\lambda}})\leq{\epsilon}]\geq
1-\delta.$ (34)

Given fixed vectors $\widetilde{{\bm{\lambda}}}\in{\mathbb{R}}^{d}$ and
${\bm{\eta}}\in{\mathbb{R}}^{d},~{}{\bm{\eta}}\geq 0$, denote
${\bm{\tilde{\Lambda}}}=\widetilde{{\bm{\lambda}}}+\omega{\bm{\eta}},~{}\omega\in\Lambda$
the line with offset $\widetilde{{\bm{\lambda}}}$ and direction ${\bm{\eta}}$
parametrized by $\omega$. We show that for entrywise monotonically
nonincreasing loss functions and for the family of set-valued predictors of
the form
${\mathcal{I}}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]$,
for some lower and upper bounds ${\hat{l}}_{j}<{\hat{u}}_{j}$ that may depend
on $y$, if

$\hat{{\bm{\lambda}}}=\operatorname*{\arg\,\min}_{{\bm{\lambda}}\in{\bm{\tilde{\Lambda}}}}~{}\sum_{j\in[d]}\lambda_{j}\quad{\text{s.t.}}\quad\hat{R}^{+}({\bm{\lambda}}+\beta{\bm{\eta}})<{\epsilon},\forall\beta\geq
0$ (35)

${\mathcal{I}}_{\hat{{\bm{\lambda}}}}(y)$ is an $({\epsilon},\delta)$-RCPS. We
start by reminding the following definitions.

######  Definition A.1 (Entrywise monotonically nonincreasing function).

A loss function $\ell$ is entrywise monotonically nonincreasing if, for a
fixed $x$, $\forall j\in[d]$

${\mathcal{I}}(y)_{j}\subset{\mathcal{I}}^{\prime}(y)_{j}\implies\ell(x,{\mathcal{I}}^{\prime}(y))\leq\ell(x,{\mathcal{I}}(y)).$
(36)

######  Definition A.2 (Entrywise nesting property).

A family of set predictors
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)\\}_{{\bm{\lambda}}\in\Lambda^{d}}$ is
entrywise nested if, for a fixed $y$, $\forall j\in[d]$

$\lambda_{j,1}<\lambda_{j,2}\implies{\mathcal{I}}_{[\lambda_{j,1},{\bm{\lambda}}_{-j}]}(y)_{j}\subset{\mathcal{I}}_{[\lambda_{j,2},{\bm{\lambda}}_{-j}]}(y)_{j},$
(37)

where $[\lambda_{j},{\bm{\lambda}}_{-j}]$ is the vector that takes value
$\lambda_{j}$ in its $j^{\text{th}}$ entry and $\lambda_{-j}$ in its
complement $-j\coloneqq[d]\setminus\\{j\\}$.

###### Proof.

The proof is a high-dimensional extension of the validity of the original RCPS
calibration procedure. Note that the family of set-valued predictors
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]\\}_{{\bm{\lambda}}\in\Lambda^{d}}$
satisfies the entrywise nesting property in Definition A.2. For
$\hat{{\bm{\lambda}}}$ chosen as in Equation 35, denote
$L:~{}\Lambda\to{\mathbb{R}}$ the one-dimensional function such that

$L(\beta)=R(\hat{{\bm{\lambda}}}+\beta{\bm{\eta}})\quad~{}\text{and}~{}\quad\hat{L}^{+}(\beta)=\hat{R}^{+}(\hat{{\bm{\lambda}}}+\beta{\bm{\eta}}).$
(38)

It follows that $L$ is monotonically nonincreasing because $\ell$ is entrywise
monotonically nonincreasing by assumption, and ${\bm{\eta}}$ belongs to the
positive orthant. Furthermore,
${\mathbb{P}}[L(\beta)\leq\hat{L}^{+}(\beta)]\geq 1-\delta$ by definition of
$\hat{R}^{+}({\bm{\lambda}})$. Denote

$\beta^{*}=\inf\\{\beta\in\Lambda:L(\beta)\leq{\epsilon}\\}$ (39)

and suppose $R(\hat{{\bm{\lambda}}})=L(0)>{\epsilon}$. By monotonicity of $L$
it follows that $\beta^{*}>0$, and by definition of $\hat{{\bm{\lambda}}}$,
$\hat{L}^{+}(\beta^{*})=\hat{R}^{+}(\hat{{\bm{\lambda}}}+\beta^{*}{\bm{\eta}})<{\epsilon}$.
However, since $L(\beta^{*})={\epsilon}$ and
${\mathbb{P}}[L(\beta^{*})\leq\hat{L}^{+}(\beta^{*})]\geq 1-\delta$, we
conclude that this event happens with probability at most $\delta$. Hence,
${\mathbb{P}}[R(\hat{{\bm{\lambda}}})\leq{\epsilon}]\geq 1-\delta$ and
${\mathcal{I}}_{\hat{{\bm{\lambda}}}}(y)$ is an $({\epsilon},\delta)$-RCPS.

Lastly, it is easy to see that $\hat{{\bm{\lambda}}}$ minimizes the mean
interval length $\bar{I}({\bm{\lambda}})$ over ${\bm{\tilde{\Lambda}}}$. Note
that

$\bar{I}({\bm{\lambda}})=\frac{1}{d}\sum_{j\in[d]}({\hat{u}}_{j}-{\hat{l}}_{j})+2\sum_{j\in[d]}\lambda_{j},$
(40)

and the statement follows by definition. ∎

###  A.3 Proof that $\ell^{\gamma}$ is a convex upper bound to $\ell^{01}$
(see Equation 16)

Recall that for a family of set-valued predictors
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]\\}_{{\bm{\lambda}}\in\Lambda^{d}}$
indexed by a $d$-dimensional vector
${\bm{\lambda}}=(\lambda_{1},\dots,\lambda_{d})\in\Lambda^{d}$,
$\Lambda\subset\overline{{\mathbb{R}}}\coloneqq{\mathbb{R}}\cup\\{\pm\infty\\}$
for some general lower and upper bounds ${\hat{l}}_{j}<{\hat{u}}_{j}$ that may
depend on $y$, we define $\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$
to be

$\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))=\frac{1}{d}\sum_{j\in[d]}\left[\frac{2(1+q)}{I({\bm{\lambda}})_{j}}\lvert
x_{j}-c_{j}\rvert-q\right]_{+},$ (41)

where $q=\gamma/(1-\gamma)$, $\gamma\in[0,1)$,
$I({\bm{\lambda}})_{j}={\hat{u}}_{j}-{\hat{l}}_{j}+2\lambda_{j}$,
$c_{j}=({\hat{u}}_{j}+{\hat{l}}_{j})/2$, and $[u]_{+}=\max(0,u)$. First, we
show that $\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$ is convex in
${\bm{\lambda}}$ for ${\bm{\lambda}}\geq 0$.

###### Proof.

Note that $\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$ is separable in
${\bm{\lambda}}$. Hence, it suffices to show that
$\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$ is convex in each entry
$\lambda_{j}$. That is, we want to show that

$\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}=\left[\frac{2(1+q)}{{\hat{u}}_{j}-{\hat{l}}_{j}+2\lambda_{j}}\lvert
x_{j}-c_{j}\rvert-q\right]_{+}$ (42)

is convex in $\lambda_{j}$. Note that:

  * •

The term $1/({\hat{u}}_{j}-{\hat{l}}_{j}+2\lambda_{j})$ behaves like
$1/\lambda_{j}$, hence it is convex for $\lambda_{j}\geq
0>-({\hat{u}}_{j}-{\hat{l}}_{j})/2$,

  * •

$C=2(1+q)\lvert x_{j}-c_{j}\rvert$ is nonnegative, hence $C\cdot
1/({\hat{u}}_{j}-{\hat{l}}_{j}+2\lambda_{j})$ is convex,

  * •

$q$ does not depend on $\lambda_{j}$, hence $C\cdot
1/({\hat{u}}_{j}-{\hat{l}}_{j}+2\lambda_{j})-q$ is still convex, and finally

  * •

the positive part $[u]_{+}=\max(0,u)$ is a convex function of its argument,
hence $[C\cdot 1/({\hat{u}}_{j}-{\hat{l}}_{j}+2\lambda_{j})-q]_{+}$ is convex.

We conclude that $\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}$ is
convex in each entry $j\in[d]$ for $\lambda_{j}\geq 0$, hence
$\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$ is convex for
${\bm{\lambda}}\geq 0$. ∎

Note that $\ell^{\gamma}$ is an upper bound to $\ell^{01}$ by construction.
One can see that
$\forall{\bm{\lambda}},~{}\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))\geq\ell^{01}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$.

###### Proof.

We will now show that
$\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}\geq\ell^{01}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}$
entrywise. To this end, we will show that both functions attain the same value
(i.e., 1) at the extremes of the intervals (i.e., ${\hat{l}}-{\bm{\lambda}}$
and ${\hat{u}}+{\bm{\lambda}}$), and use the fact that
$\ell^{\gamma}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}$ is convex
nonnegative.

First, note that $\forall{\bm{\lambda}}$, it holds that

$\displaystyle\ell^{\gamma}({\hat{l}}-{\bm{\lambda}},{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}$
$\displaystyle=\ell^{\gamma}({\hat{u}}+{\bm{\lambda}},{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}=\left[\frac{2(1+q)}{I({\bm{\lambda}})_{j}}\cdot\bigg{\lvert}\frac{I({\bm{\lambda}})_{j}}{2}\bigg{\rvert}-q\right]_{+}=1,$
(43)
$\displaystyle\ell^{01}({\hat{l}}-{\bm{\lambda}},{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}$
$\displaystyle=\ell^{01}({\hat{u}}+{\bm{\lambda}},{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}=1.$
(44) Furthermore,
$\displaystyle\ell^{01}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))_{j}$
$\displaystyle=0,~{}\forall
x_{j}\in[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]$ (45)

by definition. We conclude that $\ell^{\gamma}$ upper bounds $\ell^{01}$
because $\ell^{\gamma}$ is convex, nonnegative, it intersects $\ell^{01}$ at
${\hat{l}}-{\bm{\lambda}}$ and ${\hat{u}}+{\bm{\lambda}}$ (for which both
losses are equal to 1), and $\ell^{01}$ is exactly 0 between
${\hat{l}}_{j}-\lambda_{j}$ and ${\hat{u}}_{j}+\lambda_{j}$. ∎

##  Appendix B Risk Controlling Prediction Sets (RCPS) [10]

In this section, we present the original RCPS procedure presented in Bates et
al. [10]. Let
$\ell:~{}{\mathcal{Y}}\times{\mathcal{Y}}^{\prime}\to{\mathbb{R}}$,
${\mathcal{Y}}^{\prime}\subseteq 2^{{\mathcal{Y}}}$ be a monotonically
nonincreasing loss function (see Equation 4) over
${\mathcal{X}}\subset{\mathbb{R}}^{d}$ and
${\mathcal{Y}}\subset{\mathbb{R}}^{d}$, and let
$\\{{\mathcal{I}}_{\lambda}(y)\\}_{\lambda\in\Lambda}$,
$\Lambda\subset\overline{{\mathbb{R}}}$ be a family of set-valued predictors
${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}$ that satisfies the
nesting property in Equation 5. Here,

Algorithm 2 summarizes the original conformal risk control procedure
introduced in Bates et al. [10] for a general bounding function
${\text{UCB}}(n,\delta,\hat{R}(\lambda))$ that accepts the number of samples
in a calibration set
${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}$ of $n$ i.i.d.
samples from an unknown distribution ${\mathcal{D}}$ over
${\mathcal{X}}\times{\mathcal{Y}}$, a failure probability $\delta\in[0,1]$,
the empirical estimate
$\hat{R}(\lambda)=1/n\cdot\sum_{i=1}^{n}\ell(x_{i},{\mathcal{I}}_{\lambda}(y_{i}))$
evaluated on ${\mathcal{S}}_{{\text{cal}}}$, and that returns a pointwise
upper confidence bound $\hat{R}^{+}(\lambda)$ that satisfies

${\mathbb{P}}[R(\lambda)\leq\hat{R}^{+}(\lambda)]\geq 1-\delta$ (46)

as presented in Equation 6. For example, for losses bounded above by 1,
Hoeffding’s inequality [24] yields

$\hat{R}^{+}(\lambda)=\text{UCB}(n,\delta,\hat{R}(\lambda))=\hat{R}(\lambda)+\sqrt{\frac{1}{2n}\log\left(\frac{1}{\delta}\right)}.$
(47)

In practice, tighter alternatives exist (see Bates et al. [10], Section 3.1
for a thorough discussion).

Algorithm 2 $({\epsilon},\delta)$-RCPS (see [6], Algorithm 2)

1: Input: risk level ${\epsilon}\geq 0$, failure probability $\delta\in[0,1]$,
calibration set ${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}$
of $n$ i.i.d. samples, monotonically nonincreasing loss function $\ell$,
family of nested set-valued predictors
$\\{{\mathcal{I}}_{\lambda}\\}_{\lambda\in\Lambda}$,
$\Lambda\subset\overline{{\mathbb{R}}}$, initial (large) value
$\lambda_{\max}$, stepsize ${\text{d}}{\lambda}>0$.

2: $\lambda\leftarrow\lambda_{\max}$

3: $\hat{R}^{+}(\lambda)\leftarrow 0$

4: while $\hat{R}^{+}(\lambda)\leq{\epsilon}$ do

5: $\lambda\leftarrow\lambda-{\text{d}}{\lambda}$

6: $\hat{R}(\lambda)\leftarrow
1/n\cdot\sum_{(x_{i},y_{i})\in{\mathcal{S}}_{{\text{cal}}}}\ell(x_{i},{\mathcal{I}}_{\lambda}(y_{i}))$

7: $\hat{R}^{+}(\lambda)\leftarrow\text{UCB}(n,\delta,\hat{R}(\lambda))$

8: end while

9: $\hat{\lambda}\leftarrow\lambda+{\text{d}}{\lambda}$

10: return $\hat{\lambda}$

##  Appendix C Computing the UCB

In this section, we include a detailed analytical example on how one can
compute the UCB $\hat{R}^{01+}({\bm{\lambda}}+{\bm{\eta}})$ by means of
concentration inequalities. Recall that, for a family of nested set predictors
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)\\}_{{\bm{\lambda}}\in\Lambda^{d}}$
indexed by a _vector_ ${\bm{\lambda}}=(\lambda_{1},\dots,\lambda_{d})$, the
$01$ loss $\ell^{01}(x,{\mathcal{I}}_{{\bm{\lambda}}}(y))$ counts the number
of pixels in $x$ that fall outside of their respective intervals in
${\mathcal{I}}_{{\bm{\lambda}}}(y)$. For example, choosing
${\mathcal{I}}_{{\bm{\lambda}}}(y)_{j}=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]$
as proposed in Equation 11, given a vector
${\bm{\eta}}=(\eta_{1},\dots,\eta_{d})$ yields

$\ell^{01}(x,{\mathcal{I}}_{{\bm{\lambda}}+{\bm{\eta}}}(y))=\frac{1}{d}\sum_{j\in[d]}{\mathbbm{1}}\left[x_{j}\notin[{\hat{l}}_{j}-(\lambda_{j}+\eta_{j}),{\hat{u}}_{j}+(\lambda_{j}+\eta_{j})]\right].$
(48)

Then,
$R^{01}({\bm{\lambda}}+{\bm{\eta}})={\mathbb{E}}[\ell^{01}(x,{\mathcal{I}}_{{\bm{\lambda}}+{\bm{\eta}}}(y))]$
denotes the risk of ${\mathcal{I}}_{{\bm{\lambda}}+{\bm{\eta}}}$ and its
empirical estimate over a calibration set
${\mathcal{S}}_{{\text{cal}}}=\\{(x_{i},y_{i})\\}_{i=1}^{n}$ of $n$ i.i.d.
samples from an unknown distribution ${\mathcal{D}}$ over
${\mathcal{X}}\times{\mathcal{Y}}$ is

$\hat{R}^{01}({\bm{\lambda}}+{\bm{\eta}})=\frac{1}{n}\sum_{i\in[n]}\ell^{01}(x_{i},{\mathcal{I}}_{{\bm{\lambda}}+{\bm{\eta}}}(y_{i})).$
(49)

In turn, for a user-specified failure probability $\delta\in[0,1]$, the UCB
$R^{01+}({\bm{\lambda}}+{\bm{\eta}})$ can be obtained by means of common
concentration inequalities. For example, Hoeffding’s inequality implies

$\hat{R}^{01+}({\bm{\lambda}}+{\bm{\eta}})=\hat{R}^{01}({\bm{\lambda}}+{\bm{\eta}})+\sqrt{\frac{1}{2n}\log\left(\frac{1}{\delta}\right)}.$
(50)

##  Appendix D Experimental Details

In this section, we include further experimental details. All experiments were
performed on a private cluster with 8 NVIDIA RTX A5000 with 24 GB of memory.

###  D.1 Datasets

(a) CelebA dataset. (b) AbdomenCT-1K dataset.

Figure D.1: Example images.

####  D.1.1 CelebA dataset

The CelebA dataset [40] (available at
http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) contains more than
$200\times 10^{3}$, $178\times 218$ pixel images of celebrity faces with
several landmark locations and binary attributes (e.g., eyeglasses, bangs,
smiling). In this work, we center-crop all images to $128\times 128$ pixels
and normalize them between $[-1,1]$.

####  D.1.2 AbdomenCT-1K dataset

The AbdomenCT-1K dataset [41] (available at
https://github.com/JunMa11/AbdomenCT-1K) comprises more than 1000 abdominal CT
scans (for a total of more than $200\times 10^{3}$, $512\times 512$ pixel
individual images) aggregated from 6 existing datasets. Scans are provided in
NIfTI format, so we first convert them to their Hounsfield unit (HU) values
and subsequently window them with the standard abdomen setting (WW = 400, WL =
40) such that pixels intensities are normalized in $[0,1]$ and they represent
the same tissue across images.

Figure D.1 shows some example images for both datasets.

###  D.2 Models

Recall that in this work, we train both:

  * •

A score network $s(\tilde{x},t)\approx{\nabla}_{x}\log p_{t}(\tilde{x})$ to
sample from $p(x|y)$ as introduced in Section 2.1, and

  * •

a modified time-conditional image regressor
$f:~{}{\mathcal{Y}}\times{\mathbb{R}}\to{\mathcal{X}}^{3}$ such that
$f(y,t)=({\hat{q}}_{\alpha/2},\hat{x},{\hat{q}}_{\alpha/2})$, where
$\hat{x}\approx{\mathbb{E}}[x\mid y]$ and
${\hat{q}}_{\alpha/2},{\hat{q}}_{1-\alpha/2}$ are the entrywise $\alpha/2$ and
$1-\alpha/2$ quantile regressors of $x$ [34, 45, 6], respectively.

Both models are implemented wit the same U-net-like [46] backbone: NCSN++,
which was introduced by Song et al. [53] (code is available at
https://github.com/yang-song/score_sde). We use the original NCSN++
configurations presented in Song et al. [53] for the CelebA dataset and, for
the AbdomenCT-1K dataset, for the FFHQ dataset (available at
https://github.com/NVlabs/ffhq-dataset) given the larger image size of
$512\times 512$. For the image regressor $f$, we use the original
implementation of the quantile regression head used in Angelopoulos et al. [6]
(available at https://github.com/aangelopoulos/im2im-uq) on top of the NCSN++
backbone. This allows us to maintain a time-conditional backbone and extend
the original image regressor presented in Angelopoulos et al. [6] to all noise
levels as the score network.

###  D.3 Training

####  D.3.1 Data Augmentation

For both datasets, we augment the training data by means of random horizontal
and vertical flips, and random rotations between $[-\pi/2,\pi/2]$ degrees.

####  D.3.2 Forward SDE

(a) CelebA dataset. (b) AbdomenCT-1K dataset.

Figure D.2: Example of perturbed images via the forward SDE. The final level
of noise is $\sigma^{2}=1$.

Recall that in this work, we are interested in solving the classical denoising
problem where $y=x+v_{0}$ is a noisy observation of a ground truth image $x$
perturbed with random Gaussian noise with known variance $\sigma_{0}^{2}$. As
done by previous works [50, 51, 54, 31], we model the observation process with
a _variance exploding_ (VE) forward-time SDE

${\text{d}}{x}=\sqrt{\frac{{\text{d}}{[\sigma^{2}(t)]}}{{\text{d}}{t}}}\,{\text{d}}{w}\quad~{}\text{where}~{}\quad\sigma(t)=\sigma_{\min}\left(\frac{\sigma_{\max}}{\sigma_{\min}}\right)^{t},~{}t\in[0,1]$
(51)

such that $\sigma(0)=\sigma_{\min}$ and $\sigma(1)=\sigma_{\max}$. In
particular, we set $\sigma_{\min}=0$ and $\sigma_{\max}=90$ for the CelebA
dataset and $\sigma_{\max}=1$ for the AbdomenCT-1K dataset. It has been shown
[53] that the above forward-time SDE is equivalent to the following discrete
Markov chain

$x_{t}=x_{t-1}+\sqrt{\sigma^{2}_{t}-\sigma^{2}_{t-1}}z,~{}z\sim{\mathcal{N}}(0,{\mathbb{I}}),\quad~{}\text{for}~{}\quad
t=1,\dots,N$ (52)

and $\\{\sigma_{i}\\}_{i=0}^{N}$ noise levels. That is, $x_{t}=x+z_{t}$, where
$z_{t}\sim(0,\sigma_{t}^{2})$. Figure D.2 shows some example images from both
datasets perturbed via the forward SDE described in this section.

####  D.3.3 Denoising Score-matching

Here, we briefly describe the loss function used to train the time-conditional
score network $s(\tilde{x},t)$. Denote $\theta\in\Theta$ the parametrization
of $s$, then, following Song and Ermon [50], Song et al. [53], we have
$s(\tilde{x},t)=s_{\theta^{*}}(\tilde{x},t)$, where

$\displaystyle\theta^{*}$
$\displaystyle=\operatorname*{\arg\,\min}_{\theta\in\Theta}{\mathbb{E}}_{t\sim
U(0,1)}\left[\xi(t){\mathbb{E}}_{x\sim
p(x),~{}x(t)|x}\left[\|s_{\theta}(x(t),t)-{\nabla}_{x}\log
p_{t}(x(t)|x)\|^{2}\right]\right]$ (53)
$\displaystyle=\operatorname*{\arg\,\min}_{\theta\in\Theta}{\mathbb{E}}_{t\sim
U(0,1)}\left[\xi(t){\mathbb{E}}_{x\sim
p(x),~{}x(t)|x}\left[\bigg{\|}s_{\theta}(x(t),t)+\frac{x(t)-x}{\sigma(t)}\bigg{\|}^{2}\right]\right],$
(54)

where $\xi(t)\propto\sigma^{2}(t)$ and $U(0,1)$ is the uniform distribution
over $[0,1]$.

####  D.3.4 Quantile Regression

(a) CelebA dataset. (b) AbdomenCT-1K dataset.

Figure D.3: Example of results with the modified image regressor.

Similarly to above, we briefly describe the loss function used to train the
modified time-conditional image regressor $f$. Denote $\theta\in\Theta$ the
parametrization of $f$, and recall that for a desired $\alpha$ quantile and
its respective quantile regressor ${\hat{q}}_{\alpha}$, the _quantile loss
function_ [34, 45, 6] $\ell^{\alpha}(x,{\hat{q}}_{\alpha})$ is

$\ell^{\alpha}(x,{\hat{q}}_{\alpha})=\alpha(x-{\hat{q}}_{\alpha})\cdot{\mathbbm{1}}[x>{\hat{q}}_{\alpha}]+(1-\alpha)({\hat{q}}_{\alpha}-x)\cdot{\mathbbm{1}}[x\leq{\hat{q}}_{\alpha}].$
(55)

For the image regressor $f$, we set $\alpha=0.10$ and use the original
implementation of the quantile regression head from Angelopoulos et al. [6]
(available at https://github.com/aangelopoulos/im2im-uq) which minimizes the
multi-loss objective

$\mathcal{L}(x,f(y,t))=\ell^{\alpha/2}(x,{\hat{q}}_{\alpha/2})+(x-\hat{x})^{2}+\ell^{1-\alpha/2}(x,{\hat{q}}_{1-\alpha/2})$
(56)

on top of the NCSN++ backbone. This allows us to maintain a time-conditional
backbone and extend the original image regressor presented in Angelopoulos et
al. [6] to all noise levels as the score network. Then,
$f(y,t)=f_{\theta^{*}}(y,t)$ where

$\theta^{*}=\operatorname*{\arg\,\min}_{\theta\in\Theta}{\mathbb{E}}_{t\sim
U(0,1)}\left[\mathcal{L}(x,f(x(t),t)\right].$ (57)

###  D.4 Sampling

(a) CelebA dataset. (b) AbdomenCT-1K dataset.

Figure D.4: Example of images sampled via Algorithm 3.

Here, we briefly describe the sampling procedure used in this work to solve
the conditional reserve-time SDE

${\text{d}}{x}=[h(x,t)-g(t)^{2}{\nabla}_{x}\log
p_{t}(x|y)]\,{\text{d}}{t}+g(t)\,{\text{d}}{\bar{w}}$ (58)

where $y$ is the initial noisy observation with known noise level
$\sigma_{0}^{2}$, and
$g(t)=\sqrt{{\text{d}}{[\sigma^{2}(t)]}/{\text{d}}{t}},~{}h(x,t)=0$ as in
Equation 51. Although several different discretization schemes exist [53], we
use the classical Euler-Maruyama discretization [30] since the contribution of
this work is not focused on improving existing diffusion models. Recall that
for a general Itô process

${\text{d}}{x}=h(x,t)\,{\text{d}}{t}+g(t)\,{\text{d}}{w},$ (59)

its Euler-Maruyama discretization with step $\Delta{t}$ is

$x_{t+1}=x_{t}+h(x,t)\,\Delta{t}+g(t)\sqrt{\Delta{t}}\,z,~{}z\sim{\mathcal{N}}(0,{\mathbb{I}}).$
(60)

Furthermore, under the forward SDE described in Equation 51, and as shown by
previous works [32, 31, 29], it holds that

${\nabla}_{x}\log p_{t}(x(t)|y)={\nabla}_{x}\log
p_{t}(x(t))+\frac{y-x(t)}{\sigma^{2}_{0}-\sigma^{2}(t)}.$ (61)

We conclude that the reverse-time Euler-Maruyama discretization of the
conditional SDE above is

$x_{t-1}=x_{t}+g^{2}_{t}\left[{\nabla}_{x}\log
p_{t}(x_{t})+\frac{y-x_{t}}{\sigma^{2}_{0}-\sigma^{2}_{t}}\right]\,\Delta{t}+g_{t}\sqrt{\Delta{t}}\,z,~{}z\sim{\mathcal{N}}(0,{\mathbb{I}}).$
(62)

Algorithm 3 implements the sampling procedure for a general score network
$s(\tilde{x},t)$.

Algorithm 3 Denoising Reverse-time SDE

1: Input: observation $y$, initial noise level $\sigma_{0}$,
$\sigma_{\max},\sigma_{\min}$, number of steps $N$

2:
$t_{0}\leftarrow(\log\sigma_{0}-\log\sigma_{\min})/(\log\sigma_{\max}-\log\sigma_{\min})$

3: $\Delta{t}\leftarrow 1/N$

4: $n\leftarrow\lfloor t_{0}/\Delta{t}\rfloor$

5: $x_{n}\leftarrow y$

6: for $i=n,\dots,1$ do

7: Draw $z_{i}\sim{\mathcal{N}}(0,{\mathbb{I}})$

8: $t_{i}\leftarrow i\Delta t$

9:
$\sigma_{i}\leftarrow\sigma_{\min}\cdot(\sigma_{\max}/\sigma_{\min})^{t_{i}}$

10: $g_{i}\leftarrow\sigma_{i}\cdot\sqrt{2\log(\sigma_{\max}/\sigma_{\min})}$

11: $x_{i-1}\leftarrow
x_{i}+g^{2}_{i}[s(x_{i},t_{i})+(y-x_{i})/(\sigma_{0}^{2}-\sigma^{2}_{i})]\,\Delta{t}+g_{i}\sqrt{\Delta{t}}\,z_{i}$

12: end for

13: return $x_{0}$

##  Appendix E Comparison with Conffusion

In this section, we discuss the differences between the Con$ffusion$ framework
proposed by Horwitz and Hoshen [26] and the contributions of this paper.
Although similar in spirit and motivation, the two are fundamentally different
and address separate and distinct questions. Con$ffusion$ deploys ideas of
quantile regression [34] to fine-tune an existing score network and obtain an
image regressor equipped with heuristic uncertainty intervals. Such intervals
can then be conformalized to provide risk control. On the other hand, our
$K$-RCPS is a novel high-dimensional calibration procedure for any stochastic
sampler and any notion of uncertainty (i.e., it is agnostic of the notion of
uncertainty), including diffusion models and quantile regression. In this
paper, we present $K$-RCPS for diffusion models given their remarkable
performance in solving inverse problems via conditional sampling.

In order to contextualize the choices of parametrization for the families of
set predictors in Equations 19 and 20, we highlight a gap between the
presentation of Con$ffusion$ in the arXiv paper available at
https://arxiv.org/abs/2211.09795v1 and the official code release in
https://github.com/eliahuhorwitz/Conffusion. In particular Section 3.2 of
Horwitz and Hoshen [26] introduces the following family of set predictors:

${\mathcal{I}}^{~{}\text{arXiv}}_{{\bm{\lambda}},\text{Conffusion}}(y)_{j}=[\lambda_{j}({\hat{q}}_{\alpha/2})_{j},\lambda_{j}({\hat{q}}_{1-\alpha/2})_{j}],$
(63)

which does not satisfy the nesting property in Equation 5. It is easy to see
that—without any further assumptions—as $\lambda_{j}$ increases,
${\mathcal{I}}^{~{}\text{arXiv}}_{{\bm{\lambda}},\text{Conffusion}}(y)_{j}$
does not cover the interval $[0,1]$ and the $01$ loss in Equation 10 is not
entrywise monotonically non-increasing. This renders the RCPS procedure (and
most calibration procedures) not applicable as presented in the arXiv paper.
While this family of set predictors is thus not applicable, we resorted to the
authors’ GitHub release where such parametrization is different (see
https://github.com/eliahuhorwitz/Conffusion/blob/fffe5c946219cf9dead1a1c921a131111e31214e/inpainting_n_conffusion/core/calibration_masked.py#L28).
More precisely, the GitHub implementation reflects the multiplicative
parametrization presented in Equation 19, which does indeed satisfy the
entrywise nesting property. Finally, we additionally propose the additive
parametrization in Equation 20 to showcase how $K$-RCPS is agnostic of the
notion of uncertainty and it can be used in conjunction with Con$ffusion$.

##  Appendix F Erratum of Theorem 3.3

In this section, we address the differences between the version of Theorem 3.3
presented in this manuscript and the published paper available at
https://proceedings.mlr.press/v202/teneggi23a.html, which contains a small
typo. We include the two theorems side-by-side with differences highlighted.

######  Theorem (This manuscript).

Let $\ell:~{}{\mathcal{X}}\times{\mathcal{X}}^{\prime}\to{\mathbb{R}}$,
${\mathcal{X}}^{\prime}\subseteq 2^{{\mathcal{X}}}$,
${\mathcal{X}}\subset{\mathbb{R}}^{d}$ be an entrywise monotonically
nonincreasing function and let
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]\\}_{{\bm{\lambda}}\in\Lambda^{d}}$
be a family of set-valued predictors
${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}$,
${\mathcal{Y}}\subset{\mathbb{R}}^{d}$ indexed by
${\bm{\lambda}}\in\Lambda^{d}$, $\Lambda\subset\overline{{\mathbb{R}}}$, for
some lower and upper bounds ${\hat{l}}_{j}<{\hat{u}}_{j}$ that may depend on
$y$. Given
${\bm{\tilde{\Lambda}}}=\widetilde{{\bm{\lambda}}}+\omega{\bm{\eta}},~{}\omega\in\Lambda$,
for fixed $\widetilde{{\bm{\lambda}}}\in{\mathbb{R}}^{d}$ and
${\bm{\eta}}\in{\mathbb{R}}^{d},~{}{\bm{\eta}}\geq 0$, if

$\displaystyle\hat{{\bm{\lambda}}}$
$\displaystyle=\operatorname*{\arg\,\min}_{{\color[rgb]{0,0.88,0}{\bm{\lambda}}\in{\bm{\tilde{\Lambda}}}}}~{}\sum_{j\in[d]}\lambda_{j}\quad{\text{s.t.}}\quad\hat{R}^{+}({\bm{\lambda}}+\beta{\bm{\eta}})<{\epsilon}$

$\forall\beta\geq 0$, then ${\mathcal{I}}_{\hat{{\bm{\lambda}}}}(y)$ is an
$({\epsilon},\delta)$-RCPS and $\hat{{\bm{\lambda}}}$ minimizes the mean
interval length over ${\bm{\tilde{\Lambda}}}$.

######  Theorem (Published paper).

Let $\ell:~{}{\mathcal{X}}\times{\mathcal{X}}^{\prime}\to{\mathbb{R}}$,
${\mathcal{X}}^{\prime}\subseteq 2^{{\mathcal{X}}}$,
${\mathcal{X}}\subset{\mathbb{R}}^{d}$ be an entrywise monotonically
nonincreasing function and let
$\\{{\mathcal{I}}_{{\bm{\lambda}}}(y)=[{\hat{l}}_{j}-\lambda_{j},{\hat{u}}_{j}+\lambda_{j}]\\}_{{\bm{\lambda}}\in\Lambda^{d}}$
be a family of set-valued predictors
${\mathcal{I}}:~{}{\mathcal{Y}}\to{\mathcal{X}}^{\prime}$,
${\mathcal{Y}}\subset{\mathbb{R}}^{d}$ indexed by
${\bm{\lambda}}\in\Lambda^{d}$, $\Lambda\subset\overline{{\mathbb{R}}}$, for
some lower and upper bounds ${\hat{l}}_{j}<{\hat{u}}_{j}$ that may depend on
$y$. For a fixed vector ${\bm{\eta}}\in{\mathbb{R}}^{d},~{}{\bm{\eta}}\geq 0$,
if $\displaystyle\hat{{\bm{\lambda}}}$
$\displaystyle=\operatorname*{\arg\,\min}_{{\bm{\lambda}}\in\Lambda^{d}}~{}\sum_{j\in[d]}\lambda_{j}\quad{\text{s.t.}}\quad\hat{R}^{+}({\bm{\lambda}}+\beta{\bm{\eta}})<{\epsilon}$
$\forall\beta\geq 0$, then ${\mathcal{I}}_{\hat{{\bm{\lambda}}}}(y)$ is an
$({\epsilon},\delta)$-RCPS, and $\hat{{\bm{\lambda}}}$ minimizes the mean
interval length.

##### Reason for the differences

The domain of the optimization problem in the published paper, i.e.
${\bm{\lambda}}\in\Lambda^{d}$, does not reflect the _pointwise_ coverage
guarantee provided by $\hat{R}^{+}$. That is, for each _fixed_
${\bm{\lambda}}\in\Lambda^{d}$,
${\mathbb{P}}[R({\bm{\lambda}})\leq\hat{R}^{+}({\bm{\lambda}})]\geq 1-\delta$.
To solve the optimization problem over ${\bm{\lambda}}\in\Lambda^{d}$,
$\hat{R}^{+}$ should provide _uniform_ coverage, i.e.
$\forall{\bm{\lambda}}\in\Lambda^{d},~{}{\mathbb{P}}[R({\bm{\lambda}})\leq\hat{R}^{+}({\bm{\lambda}})]\geq
1-\delta$.

##### Effect on results

We remark that the changes in Theorem 3.3 do not affect the $K$-RCPS procedure
nor the results presented in the published version of the paper.

##  Appendix G Figures

This section contains supplementary figures.

(a) CelebA dataset. (b) AbdomenCT-1K dataset.

Figure G.1: Empirical estimates of risk over 20 random draws of
${\mathcal{S}}_{{\text{cal}}}$. All combinations of notion of uncertainty and
calibration procedure successfully control risk at level
${\epsilon}=0.10,0.05$ for the CelebA and AbdomenCT-1K dataset, respectively,
with probability at least $1-\delta,~{}\delta=0.10$.