# Principal Uncertainty Quantification with Spatial Correlation
for Image Restoration Problems

Omer Belhasin
Verily Research
<EMAIL_ADDRESS>Yaniv Romano
Computer Science Department - Technion
<EMAIL_ADDRESS>Daniel Freedman
Verily Research
<EMAIL_ADDRESS>Michael Elad
Verily Research
<EMAIL_ADDRESS>
###### Abstract

Uncertainty quantification for inverse problems in imaging has drawn much
attention lately. Existing approaches towards this task define uncertainty
regions based on probable values per pixel, while ignoring spatial
correlations within the image, resulting in an exaggerated volume of
uncertainty. In this paper, we propose PUQ (Principal Uncertainty
Quantification) – a novel definition and corresponding analysis of uncertainty
regions that takes into account spatial relationships within the image, thus
providing reduced volume regions. Using recent advancements in stochastic
generative models, we derive uncertainty intervals around principal components
of the empirical posterior distribution, forming an ambiguity region that
guarantees the inclusion of true unseen values with a user confidence
probability. To improve computational efficiency and interpretability, we also
guarantee the recovery of true unseen values using only a few principal
directions, resulting in ultimately more informative uncertainty regions. Our
approach is verified through experiments on image colorization, super-
resolution, and inpainting; its effectiveness is shown through comparison to
baseline methods, demonstrating significantly tighter uncertainty regions.

##  1 Introduction

Figure 1: Comparison of PUQ’s performance on the CelebA-HQ dataset in image
colorization, super-resolution, and inpainting tasks using the Exact PUQ
procedure (Section 4.2.1) applied on RGB image patches of varying size. As
demonstrated, our method provides tighter uncertainty regions with
significantly smaller uncertainty volumes ($\sim$x10 in super-resolution and
inpainting, and $\sim$x100 in colorization). The compared methods are im2im-uq
[angelopoulos2022image] and Conffusion [horwitz2022conffusion].

Restoration problems in imaging are widely encountered in various disciplines,
including cellular cameras, surveillance, experimental physics, medical
imaging, and more. Inverse problems, such as colorization, super-resolution,
and inpainting, are typically ill-posed, implying that multiple solutions can
explain the unknown target image given the measurements. In this context,
uncertainty quantification aims to characterize the range of possible
solutions, their spread, and variability. This has an especially important
role in applications such as astronomy and medical diagnosis, where it is
necessary to establish statistical boundaries for possible gray-value
deviations. The ability to characterize the range of permissible solutions
with accompanying statistical guarantees has thus become an important and
useful challenge, which we address in this paper.

Prior work on this topic [angelopoulos2022image, horwitz2022conffusion] has
chosen to address the uncertainty assessment by constructing intervals of
possible values for each pixel via quantile regression
[koenker1978regression], or other heuristics such as estimations of per-pixel
residuals. While this line of thinking is appealing due to its simplicity, it
disregards spatial correlations within the image, and thus provides an
exaggerated uncertainty range. The study presented in
[sankaranarayanan2022semantic] has improved over the above by quantifying the
uncertainty in a latent space, thus taking spatial dependencies into account.
However, by relying on a rigid, global, non-linear, non-invertible and
uncertainty-oblivious transformation, this method suffers from severe
interpretability limitations of the assessed uncertainty.

In this paper, we propose _Principal Uncertainty Quantification_ (PUQ) – a
novel approach that accounts for spatial relationships while operating in the
image domain, thus enabling a full and clear interpretation of the quantified
uncertainty region. PUQ uses the principal components of the empirical
posterior probability density function, which describe the spread of possible
solutions. This approach reduces the uncertainty volume111The definition of
this volume, which plays a critical part in this work, is further discussed in
later sections and given in Equation (3)., as demonstrated in Figure 1. This
figure presents a comparison between our proposed Exact PUQ procedure
(described in Section 4.2.1) and previous work [angelopoulos2022image,
horwitz2022conffusion], showing a much desired trend of reduced _uncertainty
volume_ that further decreases as the size of the patch under consideration
grows.

Our work aims to improve the quantification of the uncertainty volume by
leveraging recent advancements in generative models serving as stochastic
solvers for inverse problems. While our proposed approach is applicable using
any such solver (e.g., conditional GAN [mirza2014conditional]), we focus in
this work on diffusion-based techniques, which have recently emerged as the
leading image synthesis approach, surpassing GANs and other alternative
generators [dhariwal2021diffusion]. Diffusion models offer a systematic and
well-motivated algorithmic path towards the task of sampling from a prior
probability density function (PDF), $\mathbb{P}_{y}$, through the repeated
application of a trained image-denoiser [sohl2015deep, ho2020denoising]. An
important extension of these models allows the sampler to become conditional,
drawing samples from the posterior PDF, $\mathbb{P}_{y|x}$, where $x$
represents the observed measurements. This approach has recently gained
significant attention [dhariwal2021diffusion, ho2022cascaded,
saharia2022image, saharia2022palette] and yields a new and fascinating
viewpoint to inverse problems, recovering a variety of candidate solutions to
such problems while targeting high perceptual quality.

In this work, we generalize the pixelwise uncertainty assessment, as developed
in [angelopoulos2022image, horwitz2022conffusion], so as to incorporate
spatial correlations between pixels. This generalization is obtained by
considering an image-adaptive basis for a linear space that replaces the
standard basis in the pixelwise approach. To optimize the volume of the output
uncertainty region, we propose a statistical analysis of the posterior
obtained from a diffusion-based sampler (e.g., [saharia2022image,
saharia2022palette]), considering a series of candidate restorations. Our
method may be applied both globally (on the entire image) or locally (on
selected portions or patches), yielding a tighter and more accurate
encapsulation of statistically valid uncertainty regions. For the purpose of
adapting the basis, we compute and leverage the principal components of the
candidate restorations. As illustrated in Figure 2 for a simple 2-dimensional
PDFs, the pixelwise regions are less efficient and may contain vast empty
areas, and especially so in cases where pixels exhibit strong correlation.
Clearly, as the dimension increases, the gap between the standard and the
adapted uncertainty quantifications is further amplified.

Our proposed method offers two _conformal prediction_ [vovk2005algorithmic,
lei2018distribution, angelopoulos2021gentle] based calibration options
(specifically, using the Learn then Test [angelopoulos2021learn] scheme) for
users to choose from, with a trade-off between precision and complexity. These
include (i) using the entire set of principal components, (ii) using a
predetermined subset of them222We also propose a variation of this option that
enables controlling the number of necessary principal components to be used,
thus resulting in a reduced complexity at calibration and test time.. The
proposed calibration procedures ensure the validity of the uncertainty region
to contain the unknown true values with a user-specified confidence
probability, while also ensuring the recovery of the unknown true values using
the selected principal components when only a subset is used. Applying these
approaches allows for efficient navigation within the uncertainty region of
highly probable solutions.

We conduct various local and global experiments to verify our method,
considering three challenging tasks: image colorization, super-resolution, and
inpainting, all described in Section 5, and all demonstrating the advantages
of the proposed approach. For example, when applied locally on
$8$$\times$$8\times$$3$ patches, our experiments show a reduction in the
guaranteed uncertainty volume by a factor of $\sim$$10$-$100$ compared to
previous approaches, as demonstrated in Figure 1. Moreover, this local
approach can have a substantially reduced computational complexity while
retaining the statistical guarantees, by drawing far fewer posterior samples
and using a small subset of the principal components. As another example, the
global tests on the colorization task provide an unprecedented tightness in
uncertainty volumes. This is accessible via a reduced set of drawn samples,
while also allowing for efficient navigation within the solution set.

In summary, the contributions of this paper are the following:

  1. 1.

We introduce a novel generalized definition of uncertainty region that
leverages an adapted linear-space basis for better posterior coverage.

  2. 2.

We propose a new method for quantifying the uncertainty of inverse problems
that considers spatial correlation, thus providing tight uncertainty regions.

  3. 3.

We present two novel calibration procedures for the uncertainty quantification
that provide statistical guarantees for unknown data to be included in the
uncertainty region with a desired coverage ratio while being recovered with a
small error by the selected linear axes.

  4. 4.

We provide a comprehensive empirical study of three challenging image-to-image
translation tasks: colorization, super-resolution, and inpainting,
demonstrating the effectiveness of the proposed approach in all modes.

Figure 2: An illustration of uncertainty regions (in red) of 2d posterior
distributions and considering three different PDF behaviors, shown in blue,
orange, and green. The uncertainty regions are formed from intervals, as
defined in Equation (1), where $\hat{l}(x)$ and $\hat{u}(x)$ represent the
$0.05$ and $0.95$ quantiles over the dashed black axes. The top row presents
the uncertainty region in the pixel domain using standard basis vectors that
ignores the spatial correlations, while the lower row presents the regions
using the principal components as the basis. The uncertainty volume, defined
in Equation (3), is indicated in the top left corner of each plot. The $90\%$
coverage guarantee, outlined in Equation (2) with $w_{i}:=1/2$, is satisfied
by all regions. However, the lower row regions, that take spatial dependencies
into account, are significantly smaller than the pixelwise corresponding
regions in the upper row.

##  2 Related Work

Inverse problems in imaging have been extensively studied over the years; this
research domain has been deeply influenced by the AI revolution
[lathuiliere2019comprehensive, santhanam2017generalized,
yan2016attribute2image, gregor2015draw, mirza2014conditional]. A promising
recent approach towards image-to-image translation problems relies on the
massive progress made on learned generative techniques. These new tools enable
to model the conditional distribution of the output images given the input,
offering a fair sampling from this PDF. Generative-based solvers of this sort
create a new and exciting opportunity for getting high perceptual quality
solutions for the problem in hand, while also accessing a diverse set of such
candidate solutions.

Recently, _Denoising Diffusion Probabilistic Models_ (DDPM) [sohl2015deep,
ho2020denoising] have emerged as a new paradigm for image generation,
surpassing the state-of-the-art results achieved by GANs
[goodfellow2020generative, dhariwal2021diffusion]. Consequently, several
_conditional_ diffusion methods have been explored [dhariwal2021diffusion,
ho2022cascaded, saharia2022image, saharia2022palette], including SR3
[saharia2022image] – a diffusion-based method for image super-resolution,
Palette [saharia2022palette] – a diffusion-based unified framework for image-
to-image translation tasks, and more (e.g. [song2019generative,
kadkhodaie2020solving, song2020score, kawar2021stochastic, kawar2021snips,
kawar2022denoising, kadkhodaie2023learning]). Note, however, that most current
conditional algorithms for inverse problems do not offer statistical
guarantees that protect against model deviations and hallucinations.

Moving to uncertainty quantification, the field of machine learning has been
seeing rich work on the derivation of statistically rigorous confidence
intervals for predictions [romano2019conformalized,
chernozhukov2021distributional, sesia2021conformal, gupta2022nested,
kivaranovic2020adaptive]. One key paradigm in this context is _conformal
prediction_ (CP) [vovk2005algorithmic, lei2018distribution,
angelopoulos2021gentle] and _risk-controlling_ methods [bates2021distribution,
angelopoulos2021learn, angelopoulos2022conformal], which allow to rigorously
quantify the prediction uncertainty of a machine learning model with a user-
specified probability guarantee. Despite many proposed methods, only a few
have focused on mitigating uncertainty assessment in image restoration
problems, including im2im-uq [angelopoulos2022image] and Conffusion
[horwitz2022conffusion]. These works have employed a risk-controlling paradigm
[bates2021distribution] to provide statistically valid prediction intervals
over the pixel domain, ensuring the inclusion of ground-truth solutions in the
output intervals. However, both these approaches share the same limitation of
operating in the pixel domain while disregarding spatial correlations within
the image or the color layers. This leads to an unnecessarily exaggerated
volume of uncertainty.

An exception to the above is [sankaranarayanan2022semantic], which quantifies
uncertainty in the latent space of GANs. Their migration from the image domain
to the latent space is a rigid, global, non-linear, non-invertible and
uncertainty-oblivious transformation. Therefore, quantification of the
uncertainty in this domain is quite limited. More specifically, rigidity
implies that this approach cannot adapt to the complexity of the problem by
adjusting the latent space dimension; Globality suggests that it cannot be
operated locally on patches in order to better localize the uncertainty
assessments; Being non-linear implies that an evaluation of the uncertainty
volume (discussed in Section 3) in the image domain is hard and next to
impossible; Non-invertability of this transformation means that some energy is
lost from the image in the analysis and not accounted for, thus hampering the
validity of the statistical guarantees; Finally, note that the latent space
is, at best, associated with the image content, but does not represent the
prime axes of the uncertainty behavior. All these shortcomings are converted
into advantages in our proposed approach advocated in this paper. Note,
however, that due to the above, and especially the inability to provide
certified volumes of uncertainty, a comparison of our method to
[sankaranarayanan2022semantic] is impossible.

Inspired by the above contributions, we propose a novel alternative
uncertainty quantification approach that takes spatial relationships into
account. Our work provides tight uncertainty regions, compared to prior work,
with user-defined statistical guarantees through the use of a CP-based
paradigm. Specifically, we adopted the Learn then Test [angelopoulos2021learn]
procedure that provides statistical guarantees for controlling multiple risks
in a general setting.

##  3 Problem Formulation

Let $\mathbb{P}_{x,y}$ be a probability distribution over
$\mathcal{X}\times\mathcal{Y}$, where $\mathcal{X}$ and $\mathcal{Y}$
represent the input and the output space, respectively, for the inverse
problem at hand. E.g., for the task of image colorization, $\mathcal{Y}$ could
represent full-color high-quality images, while $\mathcal{X}$ represents their
colorless versions to operate on. We assume that
$\mathcal{X,Y}\subset[0,1]^{d}\subset\mathbb{R}^{d}$, where, without loss of
generality, $d$ is assumed to be the dimension of both spaces. In the context
of examining patches within output images, we define
$\mathcal{Y}_{\text{patch}}$ as the patch space of the output images. For
simplicity, we use the same notation, $d$, for $\mathcal{Y}$ and
$\mathcal{Y}_{\text{patch}}$, while it is clear that the dimension of
$\mathcal{Y}_{\text{patch}}$ is smaller and controlled by the user through the
patch size to work on.

Given an input measurement $x\in\mathbb{R}^{d}$, our objective is to quantify
the uncertainty of the possible solutions to the inverse problem, as
manifested by the estimated $d$-dimensional posterior distribution,
$\hat{\mathbb{P}}_{y|x}$. The idea is to enhance the definition of pixelwise
uncertainty intervals as introduced in [angelopoulos2022image] by integrating
the spatial correlations between pixels to yield a better structured
uncertainty region. To achieve this, we propose to construct uncertainty
intervals using a designated collection of orthonormal basis vectors for
$\mathbb{R}^{d}$ instead of intervals over individual pixels. We denote this
collection by
$\hat{B}(x)=\\{\hat{v}_{1}(x),\hat{v}_{2}(x)\dots\hat{v}_{d}(x)\\}$, where
$\hat{v}_{i}(x)\in\mathbb{R}^{d}$. Note that these vectors are instance-
dependent, thus best adapted to their task. One intuitive example of such a
basis is the standard one, $\hat{B}(x)=\\{e_{1},e_{2}\dots e_{d}\\}$, where
$e_{i}\in\mathbb{R}^{d}$ is the one-hot vector with value 1 in the $i^{th}$
entry. In our work, we use a set of principal components of
$\hat{\mathbb{P}}_{y|x}$, which will be discussed in detail in Section 4.

Similar to previous work [angelopoulos2022image, horwitz2022conffusion], we
use an interval-based method centered around the conditional mean image, i.e.,
an estimate of $\mathbb{E}[y|x]\in\mathbb{R}^{d}$, denoted by $\hat{\mu}(x)$.
Formally, we utilize the following interval-valued function that constructs
prediction intervals along each basis vector around the estimated conditional
mean:

$\displaystyle\mathcal{T}(x;\hat{B}(x))_{i}:=\Big{[}{\hat{v}_{i}(x)}^{T}\hat{\mu}(x)-\hat{l}(x)_{i}~{},~{}{\hat{v}_{i}(x)}^{T}\hat{\mu}(x)+\hat{u}(x)_{i}\Big{]}.$
(1)

In the above, $i\in\\{1,2\dots d\\}$ is a basis vector index, and
$\hat{l}(x)_{i}\in\mathbb{R}^{+}$ and $\hat{u}(x)_{i}\in\mathbb{R}^{+}$ are
the lower and upper interval boundaries for the projected values of candidate
solutions emerging from $\hat{\mathbb{P}}_{y|x}$. That is, if
$\hat{y}\sim\hat{\mathbb{P}}_{y|x}$ is such a solution,
${\hat{v}_{i}(x)}^{T}\hat{y}$ is its $i$-th projection, and this value should
fall within $\mathcal{T}(x;\hat{B}(x))_{i}$ with high probability. Returning
to the example of the standard basis, the above equation is nothing but
pixelwise prediction intervals, which is precisely the approach taken in
[angelopoulos2022image, horwitz2022conffusion]. By leveraging this
generalization, the uncertainty intervals using these basis vectors form a
$d$-dimensional hyper-rectangle, referred to as the _uncertainty region_.

Importantly, we propose that the interval-valued function, $\mathcal{T}$,
should produce valid intervals that contain a user-specified fraction of the
projected ground-truth values within a risk level of $\alpha\in(0,1)$. In
other words, more than $1-\alpha$ of the projected ground-truth values should
be contained within the intervals, similar to the approach taken in previous
work in the pixel domain. To achieve this, we propose a holistic expression
that aggregates the effect of all the intervals, $\mathcal{T}(x;\hat{B}(x))$.
This expression leads to the following condition that should be satisfied:

$\mathbb{E}\left[\sum_{i=1}^{d}\hat{w}_{i}(x)\cdot\mathbf{1}\left\\{\hat{v}_{i}(x)^{T}y\in\mathcal{T}(x;B(x))_{i}\right\\}\right]>1-\alpha,$
(2)

where $y\in\mathbb{R}^{d}$ is the unknown ground-truth and
$\hat{w}_{i}(x)\in[0,1]$ s.t. $\sum_{i=1}^{d}\hat{w}_{i}(x)=1$ are the weight
factors that set the importance of covering the projected ground-truth values
along each interval. In Section 4 we discuss the proposed holistic expression
and a specific choice of these weights. As an example, we could set
$\alpha=0.1$ and $\hat{w}_{i}(x):=1/d$, indicating that more than $90\%$ of
the projected ground-truth values onto the basis vectors are contained in the
intervals, as illustrated in a 2d example in Figure 2 for different kinds of
$\hat{\mathbb{P}}_{y|x}$.

As discussed above and demonstrated in Figure 2, if the orthonormal basis in
Equation (1) is chosen to be the standard one, we get the pixel-based
intervals that disregard spatial correlations within the image, thus leading
to an exaggerated uncertainty region. In this work, we address this limitation
by transitioning to an instance-adapted orthonormal basis of $\mathbb{R}^{d}$
that allows the description of uncertainty using axes that are not necessarily
pixel-independent, thereby providing tighter uncertainty regions. While such a
basis could have been defined analytically using, for example, orthonormal
wavelets [meyer1990orthonormal], we suggest a learned and thus a better-tuned
one. The choice to use a linear and orthonormal representation for the
uncertainty quantification comes as a natural extension of the pixelwise
approach, retaining much of the simplicity and efficiency of treating each
axis separately. Note that the orthogonality enables the decomposition of $y$
around $\hat{\mu}(x)$ via its projected values,
$y=\hat{\mu}(x)+\sum_{i=1}^{d}\left[\hat{v}_{i}(x)^{T}(y-\hat{\mu}(x))\right]\hat{v}_{i}(x)$,
which we refer to hereafter as the _exact reconstruction property_.

To evaluate the uncertainty across different uncertainty regions, we introduce
a new metric called the _uncertainty volume_ ,
$\mathcal{V}(x;\mathcal{T}(x;\hat{B}(x)))$, which represents the $d^{th}$ root
of the uncertainty volume with respect to intervals
$\mathcal{T}(x;\hat{B}(x))$, defined in the following equation:

$\displaystyle\mathcal{V}(x;\mathcal{T}(x;$
$\displaystyle\hat{B}(x))):=\sqrt[d]{\prod_{i=1}^{d}\left[\hat{u}(x)_{i}+\hat{l}(x)_{i}\right]}$
(3)
$\displaystyle\approx\exp\left(\frac{1}{d}\sum_{i=1}^{d}\log\left(\hat{u}(x)_{i}+\hat{l}(x)_{i}+\epsilon\right)\right)-\epsilon~{},$

where $\epsilon>0$ is a small hyperparameter used for numerical stability. In
Section 5 we demonstrate that our approach results in a significantly
reduction in these uncertainty volumes when compared to previous methods.

When operating in high dimensions (e.g. on the full image), providing
uncertainty intervals for all the $d$-dimensions poses severe challenges, both
in complexity and interpretability. In this case, constructing and maintaining
the basis vectors becomes infeasible. Moreover, the uncertainty quantification
using these intervals may be less intuitive compared to the conventional
pixelwise approach because of the pixel-dependency between the basis vectors,
which makes it difficult to communicate the uncertainty to the user. To
mitigate these challenges, we propose an option of using $K\ll d$ basis
vectors that capture the essence of the uncertainty. In Section 4, we discuss
how to dynamically adjust $K$ to provide fewer uncertainty axes.

While reducing the number of basis vectors benefits in interpretability and
complexity, this option does not fulfill the exact reconstruction property.
Therefore, we propose an extension to the conventional coverage validity of
Equation (2) that takes into account the reconstruction error of the
decomposed ground-truth images. Specifically, the user sets a ratio of pixels,
$q\in\mathbb{R}$, and a maximum acceptable reconstruction error over this
ratio, $\beta$ 333This should not be confused with the acceptable coverage
level parameter in Equation (2). $\in(0,1)$. This approximation allows us to
reduce the number of basis vectors used to formulate $\hat{B}(x)$, such that
the reconstruction will be valid according to the following condition:

$\displaystyle\mathbb{E}\Biggl{[}\hat{\mathcal{Q}}_{q}\Biggl{(}\left\\{\left|\sum_{j=1}^{K}\hat{v}_{j}(x)^{T}y_{c}\hat{v}_{j}(x)-y_{c}\right|_{i}\right\\}_{i=1}^{d}\Biggr{)}\Biggr{]}\leq\beta~{},$
(4)

where $y_{c}:=y-{\hat{\mu}}(x)$ is the ground-truth image centered around
${\hat{\mu}}(x)$, and $\hat{\mathcal{Q}}_{q}(\cdot)$ is the empirical quantile
function defined by the smallest $z$ satisfying
$\frac{1}{d}\sum_{i=1}^{d}\mathbf{1}\\{z_{i}\leq\hat{\mathcal{Q}}_{q}(z)\\}\geq
q$. In Section 4, we discuss this expression for assessing the validity of the
basis vectors. As an example, setting $q=0.9$ and $\beta=0.05$ would mean that
the maximal reconstruction error of $90\%$ of the ground-truth pixels is no
more than $5\%$ of the $[0,1]$ dynamic range.

$y$ $x$ $\hat{y}_{i}\sim\hat{\mathbb{P}}_{y|x}$

Figure 3: The sampling procedure for two image restoration problems using a
conditional stochastic generator. The top row corresponds to super-resolution
in local mode with patches, while the bottom row shows colorization in global
mode. The implementation details are described in Section 4.1.

Pixels Axes PCs Axes Approximation Calibration

Figure 4: Illustration of our PUQ procedure in 2D ($d=2$) for a single
instance $x\in\mathcal{X}$. The top row corresponds to the case when $K=d=2$
(as in E-PUQ), while the bottom row depicts the case when $K=1<d=2$ (as in DA-
PUQ and RDA-PUQ). The procedure begins by drawing samples
$\hat{y}_{i}\sim\hat{\mathbb{P}}_{y|x}$. Next, these samples are projected
onto the PCs domain: $\hat{V}^{T}\hat{y}_{i}$, where
$\hat{V}:=[\hat{v}_{1},\dots,\hat{v}_{K}]\in\mathbb{R}^{d\times K}$. Then, we
compute bounds along the PCs to contain the samples at the correct ratio,
forming the intervals specified in Equation (1). Finally, the intervals are
scaled to statistically guarantee Equation (2) and contain the correct ratio
over solutions for unseen input instances. In the bottom row, the procedure
also statistically guarantees Equation (4) by ensuring a small recovery error
of solutions to unseen input instances, as demonstrated by the small variance
around the single PC.

##  4 PUQ: Principal Uncertainty Quantification

In this section, we present _Principal Uncertainty Quantification_ (PUQ), our
method for quantifying the uncertainty in inverse problems while taking into
account spatial dependencies between pixels. PUQ uses the _principal
components_ (PCs) of the solutions to the inverse problem for achieving its
goal. In Appendix A, we provide an intuition behind the choice of the PCs as
the basis. Our approach can be used either globally across the entire image,
referred to as the _global mode_ ; or locally within predefined patches or
segments of interest, referred to as the _local mode_. Local uncertainty
quantification can be applied to any task, where the dimensionality of the
target space is fully controlled by the user. In contrast, global
quantification is particularly advantageous for tasks that exhibit strong
spatial correlations between pixels.

Our proposed method consists of two phases. In the first, referred to as the
_approximation phase_ , a machine learning system is trained to predict the
PCs of possible solutions, denoted by
$\hat{B}(x)=\\{\hat{v}_{1}(x),\hat{v}_{2}(x),\dots,\hat{v}_{K}(x)\\}$ (where
$K\leq d$), as well as a set of importance weights,
$\hat{w}(x)\in\mathbb{R}^{K}$, referring to the vectors in $\hat{B}(x)$. In
addition, the system estimates the necessary terms in Equation (1), which
include the conditional mean, $\hat{\mu}(x)\in\mathbb{R}^{d}$, and the lower
and upper bounds, $\tilde{l}(x)\in\mathbb{R}^{K}$ and
$\tilde{u}(x)\in\mathbb{R}^{K}$ 444Note that these bounds are meant for
$\hat{\mathbb{P}}_{y|x}$ and not for $\mathbb{P}_{y|x}$, and thus marked with
tilde. The $\hat{l}(x),\hat{u}(x)$ bounds that are related to Equation (1) are
defined later in the calibration schemes (Section 4.2)., for the spread of
projected solutions over $\hat{B}(x)$. These estimations are then used to
construct heuristic uncertainty intervals in the $K$-dimensional space. In
Section 4.1, we discuss a method for deriving these terms using recent
advances in stochastic regression solvers for inverse problems, specifically
using diffusion models. Our approach is based on a sampling technique, as
illustrated in Figure 3.

The above-described approximation phase is merely an estimation, as the
corresponding heuristic intervals of Equation (1) may not contain the
projected ground-truth values with a desired ratio. Additionally, the basis
vectors may not be able to recover the ground-truth pixel values within an
acceptable threshold when $K<d$, or the basis set may contain insignificant
axes in terms of variability. Therefore, in the second, _calibration phase_ ,
we offer two calibration procedures on an held-out set of calibration data,
denoted by $\mathcal{S}_{\text{cal}}:=\\{(x_{i},y_{i})\\}_{i=1}^{n}$. These
assess the validity of our proposed uncertainty region over unseen data, which
is composed by the intervals defined in Equation (1). The choice between the
two calibration procedures depends on the user, taking into account the trade-
off between precision and complexity. The steps of our proposed method are
summarized in Algorithm 1, and the two calibration strategies are as follows:

(1) Exact PUQ (E-PUQ - Section 4.2.1): In the setting of an exact uncertainty
assessment, while assuming that $d$ PCs can be constructed and maintained in
full, the exact reconstruction property is satisfied. Consequently, the
calibration procedure is straightforward, involving only scaling of the
intervals until they contain the appropriate proportion of the projected
ground-truth values, similar to previous work over the pixel domain.

(2) Dimension-Adaptive PUQ (DA-PUQ - Section 4.2.2, RDA-PUQ - Appendix D): In
the setting of an approximate uncertainty assessment, while allowing for a
small recovery error of projected ground-truth instances to full-dimensional
instances, either due to complexity or interpretability reasons (see Section
3), the exact reconstruction property is no longer satisfied. Hence, in
addition to the scaling procedure outlined above, we must verify that the $K$
PCs can decompose the ground-truth pixel values with a small error. In this
calibration process, we also control the minimum number of the first
$\hat{k}(x)$ PCs out of the $K$ PCs, such that a small reconstruction error
can be guaranteed for unseen data. This number is dynamically determined per
input image, so that instances with greater pixel correlations are assigned
more PCs than those with weaker correlations. As manually determining $K$
might be challenging, we introduce the Reduced Dimension-Adaptive PUQ (RDA-
PUQ) procedure that also controls that value as part of the calibration - see
Appendix D.

Algorithm 1 Generating PUQ Axes and Intervals

1:Training set. Calibration set. Number of PCs $K\in\mathbb{N}$. An unseen
input instance $x\in\mathbb{R}^{d}$.

2:Statistically valid uncertainty axes and intervals for $x$.

3:$\triangleright$ Approximation phase

4:Train a machine learning system (e.g., Section 4.1) to estimate the
following: $K$ PCs of $\hat{\mathbb{P}}_{y|x}$ Importance weights of PCs The
conditional mean Lower and upper bounds on the PCs

5:$\triangleright$ Calibration phase

6:if Exact uncertainty (accurate) then

7: Apply E-PUQ using the calibration data

8:else if Approximate uncertainty (reduced complexity) then

9: Apply DA-PUQ / RDA-PUQ using the calibration data

10:end if

11:$\triangleright$ Inference

12:Provide statistically valid uncertainty axes and intervals in terms of
Equation (2) and Equation (4), applied to an unseen input instance $x$

In Section 5, we demonstrate a significant decrease in the uncertainty volume,
as defined in Equation (3) for each procedure, whether applied globally or
locally, compared to prior work. The E-PUQ procedure is the simplest and can
be applied locally to any task, while the DA-PUQ and RDA-PUQ procedures are
more complex and particularly effective in tasks in which pixels exhibit
strong correlations, such as in the image colorization task. Our method is
visually illustrated in Figure 4, showing a sampling methodology and a
calibration scheme using the full PCs or only a subset of them.

###  4.1 Diffusion Models for the Approximation Phase

The approximation phase, summarized in Algorithm 1 in RED, can be achieved in
various ways. In this section, we describe the implementation we used to
obtain the results detailed in Section 5. While we aim to construct the
uncertainty axes and intervals in the most straightforward way, further
exploration of more advanced methods to achieve the PCs is left for future
work.

In our implementation, we leverage the recent advances in stochastic
regression solvers for inverse problems based on diffusion models, which
enable to train a machine learning model to generate high-quality samples from
$\hat{\mathbb{P}}_{y|x}$. Formally, we define
$f_{\theta}:\mathcal{X}\times\mathcal{Z}\rightarrow\mathcal{Y}$ as a
stochastic regression solver for an inverse problem in global mode, where
$\mathcal{Z}$ is the noise seed space. Similarly, in local mode, we consider
$f_{\theta}:\mathcal{X}\times\mathcal{Z}\rightarrow\mathcal{Y}_{\text{patch}}$.
Given an input instance $x\in\mathbb{R}^{d}$, we propose to generate $K$
samples, denoted by $\\{f_{\theta}(x,z_{i})\\}_{i=1}^{K}$, where,
$f_{\theta}(x,z_{i})\sim\hat{\mathbb{P}}_{y|x}$. These samples are used to
estimate the PCs of possible solutions and their importance weights using the
SVD decomposition of the generated samples. The importance weights assign high
values to axes with large variance among projected samples, and low ones to
those with small variance. In Section 4.2, we elaborate on how these weights
are used in the calibration phase. Additionally, the samples are utilized to
estimate the conditional mean, $\hat{\mu}(x)$, and the lower and upper bounds,
$\tilde{l}(x)$ and $\tilde{u}(x)$, necessary for Equation (1). $\tilde{l}(x)$
and $\tilde{u}(x)$ are obtained by calculating quantiles of the projected
samples onto each PC, with a user-specified miss-coverage ratio
$\alpha\in(0,1)$.

To capture the full spread and variability of $\hat{\mathbb{P}}_{y|x}$, it is
necessary to generate at least $K=d$ samples to feed to the SVD procedure,
which is computationally challenging for high-dimensional data. As a way out,
we suggest working locally on patches, where $d$ is small and fully controlled
by the user by specifying the patch size to work on. However, for tasks with
strong pixel correlation, such as image colorization, a few PCs can describe
the variability of $\hat{\mathbb{P}}_{y|x}$ with a very small error.
Therefore, only a few samples (i.e., $K\ll d$) are required for the SVD
procedure to construct meaningful PCs for the entire image, while capturing
most of the richness in $\hat{\mathbb{P}}_{y|x}$. We formally summarize our
sampling-based methodology, in either global or local modes, in Algorithm 2.

Algorithm 2 Approximation Phase via Sampling

1:Instance $x\in\mathcal{X}$. Conditional stochastic generative model
$f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ or
$f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}_{\text{patch}}$. Maximal PCs /
samples number $K\leq d$. Misscoverage ratio $\alpha\in(0,1)$.

2:$\triangleright$ Generate samples drawn from $\hat{\mathbb{P}}_{y|x}$

3:for $i=1$ to $K$ do

4: $\hat{y}_{i}(x)\leftarrow f_{\theta}(x,z_{i})$

5:end for

6:$\triangleright$ Compute conditional mean

7:$\hat{\mu}(x)\leftarrow\frac{1}{K}\sum_{i=1}^{K}\hat{y}_{i}(x)$

8:$\triangleright$ Apply SVD decomposition and extract the PCs and weights

9:$\hat{Y}(x)\leftarrow[\hat{y}_{1}(x),\hat{y}_{2}(x)\dots\hat{y}_{K}(x)]\in\mathbb{R}^{d\times
K}$

10:$\hat{Y}(x)-\hat{\mu}(x)\cdot\mathbf{1}_{K}^{T}=\hat{V}(x)\hat{\Sigma}(x)\hat{U}(x)^{T}$

11:$\hat{B}(x)\leftarrow\\{\hat{v}_{1}(x),\hat{v}_{2}(x)\dots\hat{v}_{K}(x)\\}$,
where $\hat{v}_{i}(x)=[\hat{V}(x)]_{i}$

12:$\hat{w}(x)\leftarrow\left[\hat{\sigma}_{1}(x)^{2},\dots,\hat{\sigma}_{K}(x)^{2}\right]/c\in\mathbb{R}^{K}$,
where $\hat{\sigma}_{i}(x)=[\hat{\Sigma}(x)]_{i}$ and
$c=\sum_{j=1}^{K}{\hat{\sigma}_{j}(x)}^{2}$.

13:$\triangleright$ Compute $\alpha/2$ and $1-\alpha/2$ empirical quantiles of
projected samples onto each PC

14:for $i=1$ to $K$ do

15:
$\tilde{l}(x)_{i}\leftarrow\hat{\mathcal{Q}}_{\alpha/2}(\\{\hat{v}_{i}(x)^{T}(\hat{y}_{j}(x)-\hat{\mu}(x))\\}_{j=1}^{K})$

16:
$\tilde{u}(x)_{i}\leftarrow\hat{\mathcal{Q}}_{1-\alpha/2}(\\{\hat{v}_{i}(x)^{T}(\hat{y}_{j}(x)-\hat{\mu}(x))\\}_{j=1}^{K})$

17:end for

18:$K$ PCs $\hat{B}(x)$, importance weights $\hat{w}(x)$, conditional mean
$\hat{\mu}(x)$, lower and upper bounds $\tilde{l}(x)$ and $\tilde{u}(x)$.

###  4.2 Calibration Phase

In order to refine the approximation phase and obtain valid uncertainty axes
and intervals that satisfy the guarantees of Equation (2) and Equation (4), it
is necessary to apply a calibration phase, as summarized in Algorithm 1 in
BLUE. This phase includes two different options based on particular conditions
on the number of PCs to be constructed and maintained during the calibration
procedure or during inference, when applied either globally or locally. Below
we outline each of these options in more details.

####  4.2.1 Exact PUQ

The _Exact PUQ_ (E-PUQ) procedure provides the complete uncertainty of the
$d$-dimensional posterior distribution, $\mathbb{P}_{y|x}$. This approach is
applicable locally in small patches or segments for any task. In this case,
the exact reconstruction property discussed in Section 3 is satisfied, and
Equation (4) is fulfilled with $0\%$ error ($\beta=0$) across $100\%$
($q=1.0$) of the pixels. Therefore, the calibration procedure is simple,
involving only a scaling of intervals to ensure Equation (2) is satisfied with
high probability, similar to previous work [angelopoulos2022image,
horwitz2022conffusion].

Formally, for each input instance $x$ and its corresponding ground-truth value
$y\in\mathbb{R}^{d}$ in the calibration data, we use the estimators obtained
in the approximation phase to get $d$ PCs of possible solutions $\hat{B}(x)$,
their corresponding importance weights $\hat{w}(x)$, the conditional mean
$\hat{\mu}(x)$, and the lower and upper bounds, denoted by $\tilde{l}(x)$ and
$\tilde{u}(x)$. We then define the scaled intervals to be those specified in
Equation (1), with the upper and lower bounds defined as
$\hat{u}(x):=\lambda\tilde{u}(x)$ and $\hat{l}(x):=\lambda\tilde{l}(x)$, where
$\lambda\in\mathbb{R}^{+}$ is a tunable parameter that controls the scaling.
Notably, the size of the uncertainty intervals decreases as $\lambda$
decreases. We denote the scaled uncertainty intervals by
$\mathcal{T}_{\lambda}(x;\hat{B}(x))$.

The following weighted coverage loss function is used to guide our design of
$\lambda$:

$\mathcal{L}(x,y;\lambda):=\sum_{i=1}^{d}\hat{w}_{i}(x)\cdot\mathbf{1}\left\\{\hat{v}_{i}(x)^{T}y\not\in\mathcal{T}_{\lambda}(x;\hat{B}(x))_{i}\right\\}.$
(5)

This loss is closely related to the expression in Equation (2), and while it
may seem arbitrary at first, this choice is a direct extension to the one
practiced in [angelopoulos2022image, horwitz2022conffusion]. In Appendix B we
provide an additional justification for it, more tuned to the realm discussed
in this paper.

Our goal is to ensure that the expectation of $\mathcal{L}(x,y;\lambda)$ is
below a pre-specified threshold, $\alpha$, with high probability over the
calibration data. This is accomplished by a conformal prediction based
calibration scheme, and in our paper we use the LTT [angelopoulos2021learn]
procedure, which guarantees the following:

$\mathbb{P}\left(\mathbb{E}[\mathcal{L}(x,y;\hat{\lambda})]\leq\alpha\right)\geq
1-\delta~{},$ (6)

for a set of candidate values of $\lambda$, given as the set $\hat{\Lambda}$.
$\delta\in(0,1)$ is an error level on the calibration set and $\hat{\lambda}$
is the smallest value within $\hat{\Lambda}$ satisfying the above condition,
so as to provide the smallest uncertainty volume over the scaled intervals, as
defined in Equation (3), which we denote by $\mathcal{V}_{\hat{\lambda}}$.

Put in simple words, the above guarantees that more than $1-\alpha$ of the
ground-truth values projected onto the full $d$ PCs of
$\hat{\mathbb{P}}_{y|x}$ are contained in the uncertainty intervals with
probability at least $1-\delta$, where the latter probability is over the
randomness of the calibration set. The scaling factor takes into account the
weight factors to ensure that uncertainty intervals with high variability
contain a higher proportion of projected ground-truth values than those with
low variability. This is particularly important for tasks with strong pixel
correlations, where the first few PCs capture most of the variability in
possible solutions. We describe in detail the E-PUQ procedure in Algorithm 3.

Algorithm 3 Exact PUQ Procedure

1:Calibration set $\mathcal{S}_{\text{cal}}:=\\{x_{i},y_{i}\\}_{i=1}^{n}$.
Scanned calibration parameter values $\Lambda=[1\dots\lambda_{\text{max}}]$.
Approximation phase estimations
$\hat{B},\hat{w},\hat{\mu},\tilde{u},\tilde{l}$. Misscoverage ratio
$\alpha\in(0,1)$. Calibration error level $\delta\in(0,1)$.

2:for $(x,y)\in\mathcal{S}_{\text{cal}}$ do

3: $\hat{B}(x),\hat{w}(x),\hat{\mu}(x),\tilde{u}(x),\tilde{l}(x)\leftarrow$
Apply Algorithm 2 to $x$, with the choice of $K=d$ samples

4: for $\lambda\in\Lambda$ do

5:$\triangleright$ Scale uncertainty intervals

6: $\hat{u}(x)\leftarrow\lambda\tilde{u}(x)$ and
$\hat{l}(x)\leftarrow\lambda\tilde{l}(x)$

7: $\mathcal{T}_{\lambda}(x;\hat{B}(x))\leftarrow$ Equation (1) using
$\hat{\mu}(x),\hat{u}(x),\hat{l}(x)$

8:$\triangleright$ Compute weighted coverage loss, Equation (5)

9: $\mathcal{L}(x,y;\lambda)\leftarrow$

10:
$\sum_{i=1}^{d}\hat{w}_{i}(x)\cdot\mathbf{1}\left\\{{\hat{v}_{i}(x)}^{T}y\not\in\mathcal{T}_{\lambda}(x;\hat{B}(x))_{i}\right\\}$

11: end for

12:end for

13:$\hat{\Lambda}\leftarrow$ Extract valid $\lambda$s from LTT
[angelopoulos2021learn] applied on $\\{\mathcal{L}(x,y;\lambda):{(x,y)\in
S_{\text{cal}},\lambda\in\Lambda}\\}$ at risk level $\alpha$ and error level
$\delta$, referring to Equation (6).

14:$\triangleright$ Compute the minimizer for the uncert. volume, Equation (3)

15:$\hat{\lambda}\leftarrow\arg\min_{\lambda\in\hat{\Lambda}}\left\\{\frac{1}{n}\sum_{i=1}^{n}\mathcal{V}_{\lambda}(x_{i};\hat{B}(x))\right\\}$

16:Given a new instance $x\in\mathcal{X}$, obtain valid uncertainty intervals
for it, $\mathcal{T}_{\hat{\lambda}}(x;\hat{B}(x))$.

####  4.2.2 Dimension-Adaptive PUQ

The E-PUQ procedure assumes the ability to construct and maintain $d$ PCs,
which can be computationally challenging both locally and globally.
Furthermore, an uncertainty quantification over these axes may be less
intuitive, due to the many axes involved, thus harming the method’s
interpretability (see discussion in Section 3). To address these, we propose
the _Dimension-Adaptive PUQ_ (DA-PUQ) procedure, which describes the
uncertainty region with fewer axes, $K\leq d$. The use of only a few leading
dimensions, e.g., $K=3$, can lead to a more interpretable uncertainty region,
enabling an effective visual navigation within the obtained uncertainty range.

While this approach does not satisfy the exact reconstruction property (see
Section 3), the decomposed ground-truth values can still be recovered through
the $K$ PCs with a small user-defined error in addition to the coverage
guarantee. By doing so, we can achieve both the guarantees outlined in
Equation (2) and Equation (4) with high probability.

To satisfy both the coverage and reconstruction guarantees while enhancing
interpretability, we use a dynamic function,
$\hat{k}(x):\mathcal{X}\rightarrow\mathbb{N}$, and a scaling factor to control
the reconstruction and coverage risks. The function $\hat{k}(x)$ determines
the number of top PCs (out of $K$) to include in the uncertainty region, with
a focus on selecting the smallest number that can satisfy both Equation (2)
and Equation (4), so as to increase interpretability.

Formally, for each input instance $x$ and its corresponding ground-truth value
$y\in\mathbb{R}^{d}$ in the calibration data, we use the estimators obtained
in the approximation phase to estimate $K\leq d$ PCs of possible solutions,
denoted by $\hat{B}(x)$, their corresponding importance weights, denoted by
$\hat{w}(x)$, the conditional mean denoted by $\hat{\mu}(x)$, and the lower
and upper bounds denoted by $\tilde{l}(x)$ and $\tilde{u}(x)$, respectively.
We then introduce a threshold $\lambda_{1}\in(0,1)$ for the decay of the
importance weights over the PCs of solutions to $x$. The adaptive number of
PCs to be used is defined as follows:

$\hat{k}(x;\lambda_{1}):=\min_{1\leq k\leq
K}\left\\{k~{}~{}~{}\mbox{s.t.}~{}~{}~{}\sum_{i=1}^{k}\hat{w}_{i}(x)\geq\lambda_{1}\right\\}.$
(7)

Obviously, the importance weights are arranged in a descending order, starting
from the most significant axis and ending with the least significant one.

Furthermore, let $q\in(0,1)$ be a specified ratio of pixels, and
$\beta\in(0,1)$ be a maximum allowable reconstruction error over this ratio.
The reconstruction loss function to be controlled is defined as:

$\displaystyle\mathcal{L}_{1}(x,y;\lambda_{1}):=\hat{\mathcal{Q}}_{q}\left(\left\\{\left|\sum_{j=1}^{\hat{k}(x;\lambda_{1})}\hat{v}_{j}(x)^{T}y_{c}\hat{v}_{j}(x)-y_{c}\right|_{i}\right\\}_{i=1}^{d}\right)~{},$
(8)

where $\hat{\mathcal{Q}}_{q}(\cdot)$ selects the empirical $q$-quantile of the
reconstruction errors, and $y_{c}=y-{\hat{\mu}}(x)$ is the ground-truth image
centered around ${\hat{\mu}}(x)$. In Appendix C, we discuss further this
specific loss function for controlling the capability of the linear subspace
to capture the richness of the complete $d$-dimensional posterior
distribution.

At the same time, we also control the coverage risk over the $\hat{k}(x)$ PCs,
with $\alpha\in(0,1)$ representing a user-specified acceptable misscoverage
rate and $\lambda_{2}\in\mathbb{R}^{+}$ representing the calibration factor
parameter. To control this coverage risk, we define the coverage loss function
to be the same as in Equation (5), but limited to the $\hat{k}(x)$ PCs, that
is:

$\displaystyle\mathcal{L}_{2}(x$ $\displaystyle,y;\lambda_{1},\lambda_{2}):=$
(9)
$\displaystyle\sum_{i=1}^{\hat{k}(x;\lambda_{1})}\hat{w}_{i}(x)\cdot\mathbf{1}\left\\{\hat{v}_{i}(x)^{T}y\not\in\mathcal{T}_{\lambda_{2}}(x;\hat{B}(x))_{i}\right\\},$

Finally, using the reconstruction loss function of Equation (8) and the
coverage loss function of Equation (9), we seek to minimize the uncertainty
volume, defined in Equation (3), for the scaled intervals where any unused
axes (out of $d$) are fixed to zero. We denote this uncertainty volume as
$\mathcal{V}_{\lambda_{1},\lambda_{2}}$. The minimization of
$\mathcal{V}_{\lambda_{1},\lambda_{2}}$ is achieved by minimizing
$\lambda_{1}$ and $\lambda_{2}$, while ensuring that the guarantees of
Equation (2) and Equation (4) hold with high probability over the calibration
data. This can be provided, for example, through the LTT
[angelopoulos2021learn] calibration scheme, which guarantees the following:

$\displaystyle\mathbb{P}\left(\begin{array}[]{c}\mathbb{E}[\mathcal{L}_{1}(x,y;\hat{\lambda}_{1})]\leq\beta\\\
\mathbb{E}[\mathcal{L}_{2}(x,y;\hat{\lambda}_{1},\hat{\lambda}_{2})]\leq\alpha\end{array}\right)\geq
1-\delta~{},$ (12)

where $\hat{\lambda}_{1}$ and $\hat{\lambda}_{2}$ are the minimizers for the
uncertainty volume among valid calibration parameter results, $\hat{\Lambda}$,
obtained through the LTT procedure. In other words, we can reconstruct a
fraction $q$ of the ground-truth pixel values with an error no greater than
$\beta$, and a fraction of more than $1-\alpha$ of the projected ground-truth
values onto the first $\hat{k}(x;\hat{\lambda}_{1})$ PCs of $\mathbb{P}_{y|x}$
are contained in the uncertainty intervals, with a probability of at least
$1-\delta$. A detailed description of the DA-PUQ procedure is given in
Algorithm 4.

The above-described DA-PUQ procedure reduces the number of PCs to be
constructed to $K\leq d$ while using $\hat{k}(x;\hat{\lambda}_{1})\leq K$ PCs,
leading to increased efficiency in both time and space during inference.
However, determining manually the smallest $K$ value that can guarantee both
Equation (2) and Equation (4) can be challenging. To address this, we propose
an expansion of the DA-PUQ procedure; the _Reduced Dimension-Adaptive PUQ_
(RDA-PUQ) procedure that also controls the maximum number of PCs, $K$,
required for the uncertainty assessment. This approach is advantageous for
inference as it reduces the number of samples required to construct the PCs
using Algorithm 2, while ensuring both the coverage and reconstruction
guarantees of Equation (2) and Equation (4) with high probability. The RDA-PUQ
procedure is fully described in Appendix D.

Algorithm 4 Dimension-Adaptive PUQ Procedure

1: Calibration set $\mathcal{S}_{\text{cal}}:=\\{x_{i},y_{i}\\}_{i=1}^{n}$.
Scanned calibration parameter values
$\Lambda^{1}\leftarrow[1\dots\lambda_{1}{}_{\text{max}}]$ and
$\Lambda^{2}\leftarrow[1\dots\lambda_{2}{}_{\text{max}}]$. Maximal PCs number
$K\leq d$. Approximation phase estimators
$\hat{B},\hat{w},\hat{\mu},\tilde{u},\tilde{l}$. Recovered pixels ratio
$q\in(0,1)$. Reconstruction error $\beta\in(0,1)$. Misscoverage ratio
$\alpha\in(0,1)$. Calibration error level $\delta\in(0,1)$. For an effective
calibration, $\alpha,\beta,\delta$ should be close to $0$ while $q$ should be
close to $1$.

2:for $(x,y)\in\mathcal{S}_{\text{cal}}$ do

3: $\hat{B}(x),\hat{w}(x),\hat{\mu}(x),\tilde{u}(x),\tilde{l}(x)\leftarrow$
Apply Algorithm 2 to $x$, with the choice of $K$ samples

4: for $\lambda_{1}\in\Lambda_{1}$ do

5:$\triangleright$ Compute adaptive dimensionality, Equation (7)

6:
$\hat{k}(x;\lambda_{1})\leftarrow\min_{k}\left\\{k:\sum_{i=1}^{K}\hat{w}_{i}(x)\geq\lambda_{1}\right\\}$

7:$\triangleright$ Compute reconstruction loss, Equation (8)

8: $y_{c}\leftarrow y-\hat{\mu}(x)$

9:
$\mathcal{L}_{1}(x,y;\lambda_{1})\leftarrow$$\hat{\mathcal{Q}}_{q}\bigg{(}\bigg{\\{}\bigg{|}\sum_{j=1}^{\hat{k}(x;\lambda_{1})}\hat{v}_{j}(x)^{T}y_{c}\hat{v}_{j}(x)-y_{c}\bigg{|}_{i}\bigg{\\}}_{i=1}^{d}\bigg{)}$

10: for $\lambda_{2}\in\Lambda_{2}$ do

11:$\triangleright$ Scale uncertainty intervals

12: $\hat{u}(x_{i})\leftarrow\lambda_{2}\tilde{u}(x)$ and
$\hat{l}(x)\leftarrow\lambda_{2}\tilde{l}(x)$

13: $\mathcal{T}_{\lambda_{2}}(x;\hat{B}(x))\leftarrow$ Eq. (1) using
$\hat{\mu}(x),\hat{u}(x),\hat{l}(x)$

14:$\triangleright$ Compute weighted coverage loss, Equation (5)

15:
$\mathcal{L}_{2}(x,y;\lambda_{1},\lambda_{2})\leftarrow\sum_{i=1}^{\hat{k}(x;\lambda_{1})}\hat{w}_{i}(x)\cdot$
$\mathbf{1}\left\\{{\hat{v}_{i}(x)}^{T}y\not\in\mathcal{T}_{\lambda_{2}}(x;\hat{B}(x))_{i}\right\\}$

16: end for

17: end for

18:end for

19:$\hat{\Lambda}\leftarrow$ Extract valid $\lambda$s from LTT
[angelopoulos2021learn] applied on
$\\{(\mathcal{L}_{1}(x,y;\lambda_{1}),\mathcal{L}_{2}(x,y;\lambda_{1},\lambda_{2})):{(x,y)\in
S_{\text{cal}},\lambda_{1}\in\Lambda^{1},\lambda_{2}\in\Lambda^{2}}\\}$ at
risk levels ($\beta,\alpha$) and error level $\delta$, referring to Equation
(12)

20:$\triangleright$ Compute the minimizers for the uncer. volume, Equation (3)

21:$\hat{\lambda}_{1},\hat{\lambda}_{2}\leftarrow\arg\min_{\lambda_{1},\lambda_{2}\in\hat{\Lambda}}\left\\{\frac{1}{n}\sum_{i=1}^{n}\mathcal{V}_{\lambda_{1},\lambda_{2}}(x_{i};\hat{B}(x_{i}))\right\\}$

22:Given a new instance $x\in\mathcal{X}$, obtain valid uncertainty intervals
for it, $\mathcal{T}_{\hat{\lambda}_{2}}(x;\hat{B}(x))$ over
$\hat{k}(x;\hat{\lambda}_{1})\leq K$ PCs.

##  5 Empirical Study

This section presents a comprehensive empirical study of our proposed method
PUQ, applied to three challenging tasks: image colorization, super-resolution,
and inpainting, over the CelebA-HQ dataset [karras2017progressive]. Our
approximation phase starts with a sampling from the posterior, applied in our
work by the SR3 conditional diffusion model [saharia2022image]. Figure 5
presents typical sampling results for these three tasks, showing the expected
diversity in the images obtained.

The experiments we present herein verify that our method satisfies both the
reconstruction and coverage guarantees and demonstrate that PUQ provides more
confined uncertainty regions compared to prior work, including im2im-uq
[angelopoulos2022image] and Conffusion [horwitz2022conffusion]. Through the
experiments, we present superiority in uncertainty volume, as defined in
Equation (3), and in interpretability through the use of only a few PCs to
assess the uncertainty of either a patch or a complete image. All the
experiments were conducted over 100 calibration-test splits . For in-depth
additional details of our experiments, we refer the reader to Appendix E.

$x$ $y$ Samples

Figure 5: The three image recovery tasks, colorization (top), super-resolution
(middle) and inpainting (bottom). For each we present a given measurement $x$,
the ground-truth $y$, and $10$ candidate samples from the (approximated)
posterior distribution. These samples fuel the approximation phase in our
work.

###  5.1 Evaluation Metrics

Before presenting the results, we pause to discuss the metrics used to
evaluate the performance of the different methods. Although our approach is
proved to guarantee Equation (6) for E-PUQ and Equation (12) for DA-PUQ,
(through LTT [angelopoulos2021learn] procedure), we assess the validity and
tightness of these guarantees as well.

Empirical coverage risk. We measure the risk associated with the inclusion of
projected unseen ground-truth values in the uncertainty intervals. In E-PUQ,
we report the average coverage loss, defined in Equation (5). In the case of
DA-PUQ and RDA-PUQ, we report the average coverage loss, defined by Equation
(9).

Empirical reconstruction risk. We measure the risk involved in recovering
unseen ground-truth pixel values using the selected PCs. In the case of E-PUQ,
the reconstruction risk is zero by definition. However, for DA-PUQ and RDA-
PUQ, we report the average reconstruction loss, defined by Equation (8).

Interval-Size. We report the calibrated uncertainty intervals’ sizes of
Equation (1), and compare them with baseline methods. For E-PUQ, we compare
intervals over the full basis set of PCs with the intervals in the pixel
domain used in previous work. In the DA-PUQ and RDA-PUQ procedures, we apply
dimensionality reduction to $K\ll d$ dimensions. To validly compare the
intervals’ sizes of these methods to those methods over the full $d$
dimensions, we pad the remaining $d-K$ dimensions with zeros as we assume that
the error in reconstructing the ground-truth from the dimensionally reduced
samples is negligible.

Uncertainty Volume. We report the uncertainty volumes, defined in Equation
(3), for the calibrated uncertainty regions and compare them with previous
work. A smaller uncertainty volume implies a higher level of certainty in
probable solutions to $\mathbb{P}_{y|x}$. In E-PUQ, we compare volumes over
the full basis set of PCs to the pixel domain used in earlier studies. In the
DA-PUQ and RDA-PUQ procedures, we pad the remaining dimensions with zeros,
just as with the intervals’ sizes described above.

###  5.2 Local Experiments on Patches

|  | Recons. Risk | Dim. $\hat{k}(x)$ / $K$ | Uncert. Volume
---|---|---|---|---
Colorization | im2im-uq [angelopoulos2022image] | $0$ |  $192$ / $192$ | ${1.6}\mathrm{e}{-1}\pm{7.2}\mathrm{e}{-2}$
Conffusion [horwitz2022conffusion] | $0$ |  $192$ / $192$ | ${1.7}\mathrm{e}{-1}\pm{1.3}\mathrm{e}{-1}$
E-PUQ | $0$ |  $192$ / $192$ | ${2.3}\mathrm{e}{-3}\pm{9.8}\mathrm{e}{-4}$
DA-PUQ | ${2.5}\mathrm{e}{-2}\pm{5.3}\mathrm{e}{-4}$ |  $1.6\pm 0.77$ / $100$ | ${2.4}\mathrm{e}{-11}\pm{1.3}\mathrm{e}{-11}$
RDA-PUQ | ${1.7}\mathrm{e}{-2}\pm{9.7}\mathrm{e}{-4}$ |  $3.8\pm 1.8$ / $11.8\pm 6.4$ | ${6.8}\mathrm{e}{-11}\pm{3.9}\mathrm{e}{-11}$
Super- Resolution | im2im-uq [angelopoulos2022image] | $0$ |  $192$ / $192$ | ${8.8}\mathrm{e}{-2}\pm{5.7}\mathrm{e}{-2}$
Conffusion [horwitz2022conffusion] | $0$ |  $192$ / $192$ | ${8.8}\mathrm{e}{-2}\pm{6.9}\mathrm{e}{-2}$
E-PUQ | $0$ |  $192$ / $192$ | ${1.3}\mathrm{e}{-2}\pm{7.5}\mathrm{e}{-3}$
DA-PUQ | ${2.5}\mathrm{e}{-2}\pm{3.4}\mathrm{e}{-4}$ |  $11.1\pm 5.5$ / $192$ | ${3.4}\mathrm{e}{-10}\pm{3.5}\mathrm{e}{-10}$
RDA-PUQ | ${2.0}\mathrm{e}{-2}\pm{0.0}\mathrm{e}{-0}$ |  $22.8\pm 6.8$ / $70.0\pm 0.0$ | ${1.6}\mathrm{e}{-9}\pm{1.3}\mathrm{e}{-9}$
Inpainting | im2im-uq [angelopoulos2022image] | $0$ |  $192$ / $192$ | ${2.8}\mathrm{e}{-1}\pm{1.4}\mathrm{e}{-1}$
Conffusion [horwitz2022conffusion] | $0$ |  $192$ / $192$ | ${2.7}\mathrm{e}{-1}\pm{1.6}\mathrm{e}{-1}$
E-PUQ | $0$ |  $192$ / $192$ | ${1.9}\mathrm{e}{-2}\pm{1.0}\mathrm{e}{-2}$
DA-PUQ | ${1.8}\mathrm{e}{-2}\pm{1.0}\mathrm{e}{-4}$ |  $39.4\pm 11.8$ / $192$ | ${3.6}\mathrm{e}{-8}\pm{7.3}\mathrm{e}{-8}$
RDA-PUQ | ${1.8}\mathrm{e}{-2}\pm{1.6}\mathrm{e}{-3}$ |  $55.6\pm 10.2$ / $72.6\pm 30.3$ | ${1.3}\mathrm{e}{-7}\pm{8.8}\mathrm{e}{-8}$

Table 1: Local Experiments: Quantitative comparison of the means and standard
deviations of our locally applied PUQ method on RGB patch resolution of 8x8,
utilizing the two proposed procedures. Note that in this experiment $d=8\times
8\times 3=192$ and $\epsilon=1e-10$ for the volume computation. Figure 6:
Local Experiments: A comparison of E-PUQ (see Section 4.2.1) with previous
work – im2im-uq [angelopoulos2022image] and Conffusion
[horwitz2022conffusion]. These methods are applied locally on patches with
$\alpha=\delta=0.1$. Each column corresponds to a relevant metric (see Section
5.1), and each row corresponds to a specific task. The uncertainty volume was
computed with $\epsilon=1e-10$. Results indicate that our approach achieves
superior uncertainty volume. Figure 7: Local Experiments: A comparison of DA-
PUQ (see Section 4.2.2) and RDA-PUQ (see Appendix D), when applied locally on
8x8 patches with $\alpha=\delta=0.1$, $\beta=0.05$ and $q=0.9$. Each column
corresponds to a relevant metric (see Section 5.1), and each row corresponds
to a specific task. The uncertainty volume was computed with $\epsilon=1e-10$.
Here, the dimensionality is presented by two overlapping bars, where the
yellow bars represent the distribution of $K$ in DA-PUQ and $\hat{K}$ in RDA-
PUQ, and the inner bars represent the distribution of $\hat{k}(x)$ in both
cases.

$x$ $y$ im2im-uq Conffusion Our-1x1x3 Our-2x2x3 Our-4x4x3 Our-8x8x3

Figure 8: Local Experiments: Uncertainty volume maps for patches applied in
image colorization (top row), super-resolution (middle row), and inpainting
(bottom row) with E-PUQ (see Section 4.2.1) and prior work, im2im-uq
[angelopoulos2022image] and Conffusion [horwitz2022conffusion]. Each pixel in
the maps corresponds to the uncertainty volume, defined in Equation (3), of
its corresponding patch. These results indicate that our method effectively
incorporates spatial correlation between pixels, resulting in a reduction of
the uncertainty volume.

Here, we apply our proposed methods on RGB patches of increasing size — 1x1,
2x2, 4x4, and 8x8 — for image colorization, super-resolution, and inpainting
tasks. The obtained results are illustrated in Figure 6 and Figure 7, where
Figure 6 compares our exact procedure, E-PUQ, to baseline methods, and Figure
7 examines our approximation procedures, DA-PUQ and RDA-PUQ. Additionally, we
present in Table 1 a numerical comparison of uncertainty volumes across tasks
at 8x8 patch resolution. We also provide visual representations of the
uncertainty volume maps for patches at varying resolutions in Figure 8.

The results shown in Figure 6 and Figure 7 demonstrate that our method
provides smaller uncertainty volumes, and thus more confined uncertainty
regions, when compared to previous work in all tasks and patch resolutions,
and while satisfying the same statistical guarantees in all cases. More
specifically, Figure 6 compares our exact procedure, E-PUQ, to baseline
methods. Following this figure, one can see that using the E-PUQ procedure we
obtained an improvement of $\sim\times 100$ in the uncertainty volumes in
colorization and an improvement of $\sim\times 10$ in super-resolution and
inpainting, when applied to the highest resolution of 8x8. Additionally, as
the patch resolution increases, we observe a desired trend of uncertainty
volume reduction, indicating that our method takes into account spatial
correlation to reduce uncertainty. Note that even a patch size of $1\times 1$
brings a benefit in the evaluated volume, due to the exploited correlation
within the three color channels. E-PUQ reduces trivially to im2im-uq
[angelopoulos2022image] and Conffusion [horwitz2022conffusion] when applied to
scalars ($1\times 1\times 1$ patches).

In Figure 7, we examine our approximation methods, DA-PUQ and RDA-PUQ, in
which we set a relatively small reconstruction risk of $\beta=0.05$. Observe
the significantly smaller uncertainty volumes obtained; this effect is
summarized in Table 1 as well. Figure 7 also portrays the dimensionality of
the uncertainty region used with our method using two overlapping bars. The
outer bar in yellow refers to the number of PCs that need to be constructed,
denoted as $K$ in DA-PUQ and $\hat{K}$ in RDA-PUQ. The smaller this number is,
the lower the test time computational complexity. The inner bar in green
refers to the average number of the adaptively selected PCs, denoted as
$\hat{k}(x)$. A lower value of $\hat{k}(x)$ indicates better interpretability,
as fewer PCs are used at inference than those that were constructed. For
example, in the colorization task, it can be seen that the RDA-PUQ procedure
is the most computationally efficient methodology, requiring only
$\hat{K}\approx 12$ PCs to be constructed at inference, while the DA-PUQ
procedure is the most interpretable results, with uncertainty regions
consisting of only $\hat{k}(x)\in\\{1,2,3\\}$ axes.

In all experiments demonstrated in Figure 6 and Figure 7, it is noticeable
that the standard deviation of the interval-size of our approach is higher
than that of the baseline methods. This effect happens because a few intervals
along the first few PCs are wider than those along the remaining PCs. However,
the majority of the interval sizes are significantly smaller, resulting in a
much smaller uncertainty volume. Interestingly, the uncertainty intervals of
the DA-PUQ and RDA-PUQ procedures in Figure 7 exhibit larger standard
deviation compared to the E-PUQ procedure in Figure 6. We hypothesize that
this is caused when only a few intervals (e.g., 2 intervals) are used for the
calibration process while small miscoverage ratio is set by the user
($\alpha=0.1$). As an example in the case of using 2 intervals with all
samples of the calibration set, it is necessary to enlarge all the intervals
to ensure the coverage guarantee, resulting in wider intervals over the first
few PCs.

The heat maps presented in Figure 8 compare the uncertainty volumes of our
patch-based E-PUQ procedure to baseline methods. Each pixel in the presented
heat maps corresponds to the value of Equation (3) evaluated on its
corresponding patch. The results show that as the patch resolution increases,
pixels with strong correlation structure, such as pixels of the background
area, also exhibit lower uncertainty volume in their corresponding patches.
This indicates that the proposed method indeed takes into account spatial
correlation, leading to reduced uncertainty volume.

###  5.3 Global Experiments on Images

| Recons. Risk | Dim. $\hat{k}(x)$ / $K$ | Uncert. Volume
---|---|---|---
im2im-uq [angelopoulos2022image] | $0$ |  $49152$ / $49152$ | ${1.4}\mathrm{e}{-1}\pm{3.2}\mathrm{e}{-2}$
Conffusion [horwitz2022conffusion] | $0$ |  $49152$ / $49152$ | ${1.4}\mathrm{e}{-1}\pm{5.5}\mathrm{e}{-2}$
DA-PUQ | ${5.0}\mathrm{e}{-2}\pm{1.1}\mathrm{e}{-3}$ |  $2.2\pm 0.93$ / 100 | ${1.2}\mathrm{e}{-13}\pm{5.0}\mathrm{e}{-14}$
RDA-PUQ | ${4.3}\mathrm{e}{-2}\pm{2.8}\mathrm{e}{-3}$ |  $5.5\pm 4.5$ / $22.3\pm 10.9$ | ${3.1}\mathrm{e}{-13}\pm{2.4}\mathrm{e}{-13}$

Table 2: Global Experiments: Quantitative comparison of the means and standard
deviations of our globally applied PUQ method in the colorization task,
utilizing the proposed DA-PUQ (see Section 4.2.2) and RDA-PUQ (see Appendix D)
procedures. Figure 9: Global Experiments: A comparison of DA-PUQ (see Section
4.2.2) and RDA-PUQ (see Appendix D), when applied globally on the colorization
task with $\alpha=\beta=\delta=0.1$ and $q=0.95$. The uncertainty volume was
computed with $\epsilon=1e-10$.

$x$ $y$ Recons. $\hat{v}_{1}(x)$ $\hat{v}_{2}(x)$

Figure 10: Global Experiments: Visual presentation of uncertainty regions
provided by RDA-PUQ (Appendix D) when applied globally for the colorization
task. The reconstructed image is formally given by
$\hat{\mu}(x)+\sum_{i=1}^{\hat{k}(x)}\hat{v}_{i}(x)^{T}y_{c}\hat{v}_{i}(x)$,
where $y_{c}:=y-\hat{\mu}(x)$. The value of $\hat{k}(x)$ is displayed in the
top left corner of the ”Recons.” column. The importance weights,
${\hat{w}}_{1}(x)\text{ and }{\hat{w}}_{2}(x)$, are shown in the top left
corner of the $\hat{v}_{1}(x)$ and $\hat{v}_{2}(x)$ columns.

$\mid$—————— Our ——————-$\mid$ im2im-uq Conffusion

Figure 11: Global Experiments: Images sampled uniformly from the estimated
global uncertainty regions, referring to the colorization task. Using RDA-PUQ
results with high-perceptual images, while im2im-uq [angelopoulos2022image]
and Conffusion [horwitz2022conffusion] produce unlikely images. These results
indicate that our uncertainty regions are significantly more confined than
those of previous works.

We turn to examine the effectiveness and validity of DA-PUQ and RDA-PUQ when
applied to complete images at a resolution of $128\times 128$, and consider
only the colorization task. While all PUQ procedures can be applied locally
for any task, working globally is more realistic in tasks that exhibit strong
pixel correlation. Under this setting, most of the image variability could be
represented by a few PCs, i.e., we could apply DA-PUQ and RDA-PUQ while (i)
maintaining a small reconstruction risk, and (ii) using only a few PCs to
assess the uncertainty of the entire images. Note that

  * •

The tasks of super-resolution and inpainting are ill-served by a global mode
since their estimated PCs are local in nature, thus requiring huge number of
them for an effective uncertainty representation; and

  * •

The E-PUQ calibration procedure does not apply to the global approach, as it
requires computing and maintaining $d=128\times 128\times 3$ PCs.

Figure 9 visually demonstrates the performance of our approximation methods,
also summarized in Table 2. These results demonstrate that our method provides
significantly smaller uncertainty volumes compared to our local results in
Figure 7 and previous works, but this comes at the cost of introducing a small
reconstruction risk of up to $\beta=0.1$. Observe how our approximation
methods improve interpretability: the uncertainty regions consist of only 2-5
PCs in the full dimensional space of the images. The DA-PUQ procedure produces
the tightest uncertainty regions; see the uncertainty volumes in Table 2. In
addition, the mean interval-size with our procedures is very small and almost
equal to zero, indicating that the constructed uncertainty regions are tight
and narrow due to strong correlation structure of pixels. However, similar to
the previous results, the standard deviation of interval-size is spread across
a wide range. This is because few of the first PCs have wide intervals. The
RDA-PUQ procedure is the most computationally efficient as it required to
construct only $\sim$30 PCs during inference to ensure statistical validity.

Figure 10 presents selected uncertainty regions that were provided by our
proposed RDA-PUQ procedure when applied globally. As can be seen, the
projected ground-truth images using only a few555We use ${\hat{k}}(x)$ PCs and
this value is given in Figure 10. PCs results in images that are very close to
the originals. This indicates that the uncertainty region can describe the
spread and variability among solutions with small reconstruction errors. The
first two axes of our uncertainty regions exhibit semantic content, which is
consistent with a method that accounts for spatial pixel correlation. The fact
that these PCs capture foreground/background or full-object content highlights
a unique strength of our approach. We provide the importance weights of the
first two PCs, indicating impressive proportions of variability among
projected samples onto these components (see Section 4.1). For example, in the
third row, we observe that $77\%$ of the variability in
$\hat{\mathbb{P}}_{y|x}$ is captured by $\hat{v}_{1}(x)$, which mostly
controls a linear color range of the pixels associated with the hat in the
image.

In Figure 11 we visually compare samples that were generated from the
corresponding estimated uncertainty regions, by sampling uniformly a high
dimensional point (i.e., an image) from the corresponding hyper-rectangle. As
can be seen, the samples extracted from our uncertainty region are of high
perceptual quality, whereas im2im-uq [angelopoulos2022image] and Conffusion
[horwitz2022conffusion] produce highly improbable images. This testifies to
the fact that our method provides much tighter uncertainty regions, whereas
previous work results in exaggerated uncertainty regions that contain unlikely
images.

##  6 Concluding Remarks

This paper presents “Principal Uncertainty Quantification” (PUQ), a novel and
effective approach for quantifying uncertainty in any image-to-image task. PUQ
takes into account the spatial dependencies between pixels in order to achieve
significantly tighter uncertainty regions. We introduced two different
procedures for uncertainty assessment, which are selected based on the user’s
preference for balancing precision and complexity. The experimental results
demonstrate that PUQ outperforms existing methods in image colorization,
super-resolution and inpainting, by improving the uncertainty volume.
Additionally, by allowing for a small reconstruction error when recovering
ground-truth images using a selected linear subspace, PUQ produces tight
uncertainty regions with a few axes and thus improves computational complexity
and interpretability at inference. As a result, PUQ achieves state-of-the-art
performance in uncertainty quantification for image-to-image problems.

Referring to future research, more sophisticated choices that rely on recent
advancements in stochastic image regression models could be explored, so as to
improve the complexity of our proposed approximation phase. Additionally,
investigating alternative geometries for uncertainty regions could be
interesting in order to further reduce the gap between the provided region of
uncertainty and the high-density areas of the true posterior distribution.
This includes an option to divide the spatial domain into meaningful segments,
while minimizing the uncertainty volume, or consider a mixture of Gaussians
modeling of the samples of the estimated posterior distribution.

##  Appendix A Motivation for employing the PCs

In this section, we provide an intuition behind employing principal components
(PCs) as the basis, $\hat{B}(x)$, in Equation (1), as depicted in Figure 12.
The figure showcases the estimation of the first three PCs using our globally
applied PUQ and visualizes the uncertainty region formed by these axes. This
approach facilitates efficient exploration within the uncertainty region,
thanks to linear axes that incorporate spatial correlation, as illustrated by
the visualization of $\hat{v}_{1}(x)$, $\hat{v}_{2}(x)$, and $\hat{v}_{3}(x)$.

Figure 12: A visual representation demonstrating the intuition behind
utilizing principal components (PCs) as the basis, $\hat{B}(x)$, in Equation
(1) for the colorization task. The left part illustrates that the PCs
incorporate spatial correlation, with $\hat{v}_{1}(x)$ primarily controlling
the hat color, $\hat{v}_{2}(x)$ governing the background color, and
$\hat{v}_{3}(x)$ influencing the clothing color. On the right side, an
illustration of the uncertainty region is presented, composed of these axes,
where the origin is $\hat{\mu}(x)$, and each image is defined by
$\hat{\mu}(x)+\hat{v}_{i}(x)^{T}y_{c}+a$, where $y_{c}:=y-\hat{\mu}(x)$, and
$a\in\mathbb{R}$ is a controllable parameter that moves along the axis.

##  Appendix B Coverage Loss Justification

This section aims to justify our choice for the loss-function for tuning
$\lambda$ in Equation (5), and the weights used in it, $\hat{w}_{i}(x)$.
Recall, this expression is given as:

$\mathcal{L}(x,y;\lambda):=\sum_{i=1}^{d}\hat{w}_{i}(x)\cdot\mathbf{1}\left\\{\hat{v}_{i}(x)^{T}y\not\in\mathcal{T}_{\lambda}(x;\hat{B}(x))_{i}\right\\}.$

Our starting point is the given $d$-dimensional hyper-rectangle obtained from
the approximation phase, oriented along the $d$ PC directions. This shape
serves as our initially estimated uncertainty region. Given the calibration
data, $\mathcal{S}_{\text{cal}}:=\\{(x_{i},y_{i})\\}_{i=1}^{n}$, our goal is
to inflate (or deflate, if this body proves to be exaggerated) this shape
uniformly across all axes so that it contains the majority of the ground truth
examples.

Focusing on a single pair from this dataset, $(x,y)$, the degraded image $x$
is used to ignite the whole approximation phase, while the ground truth $y$
serves for assessing the obtained hyper-rectangle, by considering the
projected coordinates $\\{\hat{v}_{i}(x)^{T}y_{c}\\}_{i=1}^{d}$, where
$y_{c}:=y-\hat{\mu}(x)$. The following function measures a potential deviation
in the $i$-th axis,

$\displaystyle h_{i}(x,y):=$
$\displaystyle\max\bigl{\\{}\hat{v}_{i}(x)^{T}y_{c}-\hat{u}(x)_{i},0\bigr{\\}}$
(13)
$\displaystyle+\max\bigl{\\{}-\hat{v}_{i}(x)^{T}y_{c}+\hat{l}(x)_{i},0\bigr{\\}}.$

Written differently, this expression is also given by

$h_{i}(x,y):=\begin{cases}\hat{v}_{i}(x)^{T}y_{c}-\hat{u}(x)_{i}&\text{if
}\hat{v}_{i}(x)^{T}y_{c}>\hat{u}(x)_{i}>0\\\
\hat{l}(x)_{i}-\hat{v}_{i}(x)^{T}y_{c}&\text{if
}\hat{v}_{i}(x)^{T}y_{c}<-\hat{l}(x)_{i}<0\\\ 0&\text{otherwise}.\end{cases}$

If positive, this implies that in this axis the example spills outside the
range of the rectangle, and the value itself is the distance from it’s border.

The following expression quantifies the weighted amount of energy that should
be invested in projecting back the $\\{\hat{v}_{i}(x)^{T}y\\}_{i=1}^{d}$
coordinates to the closest border point:

$\displaystyle\text{Energy}(x,y)=\sum_{i=1}^{d}\hat{\sigma}_{i}(x)^{2}h_{i}(x,y)^{2}.$
(14)

Note that in our weighting we prioritize high-variance axes, in which
deviation from the boundaries is of greater impact. Naturally, we should tune
$\lambda$, which scales $\hat{u}(x)_{i}$ and $\hat{l}(x)_{i}$, so as to reduce
this energy below a pre-chosen threshold, thus guaranteeing that the majority
of ground truth images fall within the hyper-rectangle. While this expression
is workable, it suffers from two shortcomings: (1) It is somewhat involved to
compute; and (2) The threshold to use with it is hard to interpret and thus to
choose. Therefore, similar to previous approaches [angelopoulos2022image,
horwitz2022conffusion], we opted in this work for a binary version of Equation
(13) of the form

$b_{i}(x,y):=\begin{cases}1&\text{if }h_{i}(x,y)>0\\\
0&\text{otherwise}.\end{cases}$ (15)

In addition, we divide the energy expression, defined in Equation (14), by the
sum of squares of all the singular values, and this way obtain exactly
$\mathcal{L}(x,y;\lambda)$ as in Equation (5). Observe that, by definition, we
get that $0\leq\mathcal{L}(x,y;\lambda)\leq 1$, where the bottom bound
corresponds to a point fully within the rectangle, and the upper bound for the
case where the point is fully outside in all axes. Therefore, thresholding the
expectation of this value with $\alpha\ll 1$ is intuitive and meaningful.

##  Appendix C Reconstruction Loss Justification

This section aims to discuss our choice for the loss function for tuning
$\lambda_{1}$ in Equation (8), given by

$\displaystyle\mathcal{L}_{1}(x,y;\lambda_{1}):=\hat{\mathcal{Q}}_{q}\left(\left\\{\left|\sum_{j=1}^{\hat{k}(x;\lambda_{1})}\hat{v}_{j}(x)^{T}y_{c}\hat{v}_{j}(x)-y_{c}\right|_{i}\right\\}_{i=1}^{d}\right)~{}.$

Recall the process: We begin with $K\leq d$ PCs obtained from the
approximation phase, and then choose $\hat{k}(x;\lambda_{1})\leq K$ of them as
instance-specific number of PCs for the evaluation of the uncertainty.

Given a calibration pair $(x,y)$, $x$ is used to derive
$\hat{k}(x;\lambda_{1})$, defining a low-dimensional subspace
$\hat{V}(x):=[\hat{v}_{1}(x),\dots,\hat{v}_{\hat{k}(x;\lambda_{1})}(x)]\in\mathbb{R}^{\hat{k}(x;\lambda_{1})\times
d}$. This, along with the conditional-mean, $\hat{\mu}(x)$, represent
$\mathbb{P}_{y|x}$ as an affine subspace. The ground-truth image $y$ is then
projected onto this slab via:

$\displaystyle\text{Projection}(y)$
$\displaystyle:=\hat{\mu}(x)+\hat{V}(x)\hat{V}(x)^{T}y_{c}$ (16)
$\displaystyle=\hat{\mu}(x)+\sum_{j=1}^{\hat{k}(x;\lambda_{1})}\hat{v}_{j}(x)^{T}y_{c}\hat{v}_{j}(x)~{},$

where $y_{c}:=y-\hat{\mu}(x)$.

The parameter $\lambda_{1}$ should be tuned so as to guarantee that this
projection entails a bounded error, $dist(y,\text{Projection}(y))$ in
expectation. A natural distance measure to use here is the $L_{2}$-norm of the
difference, which aligns well with our choice to use SVD in the approximation
phase. However, $L_{2}$ accumulates the error over the whole support, thus
losing local interpretability. An alternative is using $L_{\infty}$ which
quantifies the worst possible pixelwise error induced by the low-dimensional
projection,

$\displaystyle dist(y,\text{Projection}(y))$
$\displaystyle:=\left\|\sum_{j=1}^{\hat{k}(x;\lambda_{1})}\hat{v}_{j}(x)^{T}y_{c}\hat{v}_{j}(x)-y_{c}\right\|_{\infty}~{}.$

While this measure is applicable in many tasks, there are cases (e.g.,
inpainting) in which controlling a small maximum error requires the use of a
large number of PCs, $\hat{k}(x;\lambda_{1})$. To address this, we propose a
modification by considering the maximum error over a user-defined ratio of
pixels, $q\in(0,1)$, a value close to $1$. This is equivalent to determining
the $q$-th empirical quantile, $\hat{\mathcal{Q}}_{q}$, of the error values
among the pixels, providing a more flexible and adaptive approach, which also
aligns well with the rationale of uncertainty quantification, in which the
statistical guarantees are given with probabilistic restrictions.

##  Appendix D Reduced Dimension-Adaptive PUQ

The DA-PUQ procedure (see Section 4.2.2) reduces the number of PCs to be
constructed to $K\leq d$ while using $\hat{k}(x;\hat{\lambda}_{1})\leq K$ PCs,
leading to increased efficiency in both time and space during inference.
However, determining manually the smallest $K$ value that can guarantee both
Equation (2) and Equation (4) with high probability can be challenging. To
address this, we propose an expansion of the DA-PUQ procedure; the _Reduced
Dimension-Adaptive PUQ_ (RDA-PUQ) procedure that also controls the maximum
number of PCs required for the uncertainty assessment. While this approach is
computationally intensive during calibration, it is advantageous for inference
as it reduces the number of samples required to construct the PCs using
Algorithm 2.

Specifically, for each input instance $x$ and its corresponding ground-truth
value $y$ in the calibration data, we use the estimators obtained in the
approximation phase, to estimate $\hat{K}_{\lambda_{3}}$ PCs of possible
solutions, denoted by $\hat{B}(x)$, their corresponding importance weights,
denoted by $\hat{w}(x)$, the conditional mean denoted by $\hat{\mu}(x)$, and
the lower and upper bounds denoted by $\tilde{l}(x)$ and $\tilde{u}(x)$,
respectively. Note that these estimates are now depend on $\lambda_{3}$, we
omit the additional notation for simplicity. Then, for each choice of
$\lambda_{3}$, we use these $\hat{K}_{\lambda_{3}}$-dimensional estimates
exactly as in the DA-PUQ procedure to achieve both the coverage and
reconstruction guarantees of Equation (2) and Equation (4) with high
probability.

Similar to previous approaches, we aim to minimize the uncertainty volume,
defined in Equation (3), for the scaled $\hat{K}_{\lambda_{3}}$-dimensional
intervals where any additional axis ($d-\hat{K}_{\lambda_{3}}$ axes) is fixed
to zero. We denote the uncertainty volume in this setting as
$\mathcal{V}_{\lambda_{1},\lambda_{2},\lambda_{3}}$. The minimization of
$\mathcal{V}_{\lambda_{1},\lambda_{2},\lambda_{3}}$ is achieved by minimizing
$\lambda_{1}$, $\lambda_{2}$ and $\lambda_{3}$, while ensuring that the
guarantees of Equation (2) and Equation (4) are satisfied with high
probability. This can be provided using a conformal prediction scheme, for
example, through the LTT [angelopoulos2021learn] calibration scheme, which
ensures that the following holds:

$\displaystyle\mathbb{P}\left(\begin{array}[]{c}\mathbb{E}[\mathcal{L}_{1}(x,y;\hat{\lambda}_{1},\hat{\lambda}_{3})]\leq\beta\\\
\mathbb{E}[\mathcal{L}_{2}(x,y;\hat{\lambda}_{1},\hat{\lambda}_{2},\hat{\lambda}_{3})]\leq\alpha\end{array}\right)\geq
1-\delta~{},$ (19)

where $\hat{\lambda}_{1}$, $\hat{\lambda}_{2}$ and $\hat{\lambda}_{3}$ are the
minimizers for the uncertainty volume among valid calibration parameter
results, $\hat{\Lambda}$, obtained through the LTT procedure. Note that the
loss functions, $\mathcal{L}_{1}$ and $\mathcal{L}_{2}$, in the above are
exactly those of the DA-PUQ procedure, defined in Equation (8) and Equation
(9), while replacing $K$ with $\hat{K}_{\lambda_{3}}$.

Intuitively, Equation (19) guarantees that a fraction $q$ of the ground-truth
pixel values is recovered with an error no greater than $\beta$ using no more
than $\hat{K}_{\hat{\lambda}_{3}}$ principal components, and a fraction of
more than $1-\alpha$ of the projected ground-truth values onto the first
$\hat{k}(x;\hat{\lambda}_{1})$ principal components (out of
$\hat{K}_{\hat{\lambda}_{3}})$ are contained in the uncertainty intervals,
with a probability of at least $1-\delta$. The RDA-PUQ procedure is formally
described in Algorithm 5.

Algorithm 5 Reduced Dimensional Adaptive PUQ Proc.

1:Calibration set $\mathcal{S}_{\text{cal}}:=\\{x_{i},y_{i}\\}_{i=1}^{n}$.
Scanned calibration parameter values
$\Lambda^{1}\leftarrow[1\dots\lambda_{1}{}_{\text{max}}]$,
$\Lambda^{2}\leftarrow[1\dots\lambda_{2}{}_{\text{max}}]$ and
$\Lambda^{3}\leftarrow[1\dots\lambda_{3}{}_{\text{max}}]$. Maximal PCs number
$K\leq d$. Approximation phase estimators
$\hat{B},\hat{w},\hat{\mu},\tilde{u},\tilde{l}$. Recovered pixels ratio
$q\in(0,1)$. Reconstruction error $\beta\in(0,1)$. Misscoverage ratio
$\alpha\in(0,1)$. Calibration error level $\delta\in(0,1)$.

2:for $(x,y)\in\mathcal{S}_{\text{cal}}$ do

3: for $\lambda_{3}\in\Lambda_{3}$ do

4:$\triangleright$ Reduce dimensionality

5: $\hat{K}_{\lambda_{3}}\leftarrow\lfloor K\cdot\lambda_{3}\rfloor$

6: $\hat{B}(x),\hat{w}(x),\hat{\mu}(x),\tilde{u}(x),\tilde{l}(x)\leftarrow$
Apply Algorithm 2 to $x$, with the choice of $\hat{K}_{\lambda_{3}}$ samples

7: for $\lambda_{1}\in\Lambda_{1}$ do

8:$\triangleright$ Compute adaptive dimensionality, Equation (7)

9: $\hat{k}(x;\lambda_{1},\lambda_{3})\leftarrow$
$\min_{k}\left\\{k:\sum_{i=1}^{\hat{K}_{\lambda_{3}}}\hat{w}_{i}(x)\geq\lambda_{1}\right\\}$

10:$\triangleright$ Compute reconstruction loss, Equation (8)

11: $y_{c}\leftarrow y-\hat{\mu}(x)$

12:
$\mathcal{L}_{1}(x,y;\lambda_{1},\lambda_{3})\leftarrow$$\hat{\mathcal{Q}}_{q}\bigg{(}\bigg{\\{}\bigg{|}\sum_{j=1}^{\hat{k}(x;\lambda_{1},\lambda_{3})}\hat{v}_{j}(x)^{T}y_{c}\hat{v}_{j}(x)-y_{c}\bigg{|}_{i}\bigg{\\}}_{i=1}^{d}\bigg{)}$

13: for $\lambda_{2}\in\Lambda_{2}$ do

14:$\triangleright$ Scale uncertainty intervals

15: $\hat{u}(x)\leftarrow\lambda_{2}\tilde{u}(x)$ and
$\hat{l}(x)\leftarrow\lambda_{2}\tilde{l}(x)$

16: $\mathcal{T}_{\lambda_{2}}(x;\hat{B}(x))\leftarrow$ Equation (1) using
$\hat{\mu}(x),\hat{u}(x),\hat{l}(x)$

17:$\triangleright$ Compute weighted coverage loss, Equation (5)

18:
$\mathcal{L}_{2}(x,y;\lambda_{1},\lambda_{2},\lambda_{3})\leftarrow\sum_{i=1}^{\hat{k}(x;\lambda_{1},\lambda_{3})}\hat{w}_{i}(x)\cdot$
$\mathbf{1}\left\\{{\hat{v}_{i}(x)}^{T}y\not\in\mathcal{T}_{\lambda_{2}}(x;\hat{B}(x))_{i}\right\\}$

19: end for

20: end for

21: end for

22:end for

23:$\hat{\Lambda}\leftarrow$ Extract valid $\lambda$s from LTT
[angelopoulos2021learn] applied on
$\\{(\mathcal{L}_{1}(x,y;\lambda_{1},\lambda_{3}),\mathcal{L}_{2}(x,y;\lambda_{1},\lambda_{2},\lambda_{3})):{(x,y)\in
S_{\text{cal}},\lambda_{1}\in\Lambda^{1},\lambda_{2}\in\Lambda^{2},\lambda_{3}\in\Lambda^{3}}\\}$
at risk levels ($\beta,\alpha$).

24:$\triangleright$ Compute the minimizers for the uncer. volume, Equation (3)

25:$\hat{\lambda}_{1},\hat{\lambda}_{2},\hat{\lambda}_{3}\leftarrow\arg\min_{\lambda_{1},\lambda_{2},\lambda_{3}\in\hat{\Lambda}}\left\\{\frac{1}{n}\sum_{i=1}^{n}\mathcal{V}_{\lambda_{1},\lambda_{2}}(x_{i};\hat{B}(x_{i}))\right\\}$

26:Given a new instance $x\in\mathcal{X}$, obtain valid uncertainty intervals
for it, $\mathcal{T}_{\hat{\lambda}_{2}}(x;\hat{B}(x))$ over
$\hat{k}(x;\hat{\lambda}_{1})\leq\hat{K}_{\hat{\lambda}_{3}}$ PCs.

##  Appendix E Experimental Details

This section provides details of the experimental methodology employed in this
study, including the datasets used, architectures implemented, and the
procedural details and hyperparameters of our method.

###  E.1 Datasets and Preprocessing

Our machine learning system was trained using the Flickr-Faces-HQ (FFHQ)
dataset [karras2019style], which includes 70,000 face images at a resolution
of 128x128. We conducted calibration and testing on the CelebA-HQ (CelebA)
dataset [karras2017progressive], which also consists of face images and was
resized to match the resolution of our training data. To this end, we randomly
selected 2,000 instances from CelebA, of which 1,000 were used for calibration
and 1,000 for testing. For the colorization experiments, a grey-scale
transformation was applied to the input images. For the super-resolution
experiments, patches at a resolution of 32x32 were averaged to reduce the
input image resolution by a factor of 4 in each dimension. For the inpainting
experiments, we randomly cropped pixels from the input images during the
training phase, either in squares or irregular shapes; while for the
calibration and testing data, we cropped patches at a resolution of 64x64 at
the center of the image.

###  E.2 Architecture and Training

In all our experiments, we applied the approximation phase using recent
advancements in conditional image generation through diffusion-based models,
while our proposed general scheme in Algorithm 1 can accommodate any
stochastic regression solvers for inverse problems, such as conditional GANs
[mirza2014conditional]. In all tasks, we utilized the framework for
conditional diffusion-based models proposed in the SR3 work
[saharia2022image], using a U-Net architecture. For each of the three tasks,
we trained a diffusion model separately and followed the training regimen
outlined in the code of [saharia2022image]. To ensure a valid comparison with
the baseline methods, we implemented them using the same architecture and
applied the same training regimen. All experiments, including the baseline
methods, were trained for 10,000 epochs with a batch size of 1,024 input
images.

###  E.3 PUQ Procedures and Hyperparameters

Our experimental approach follows the general scheme presented in Algorithm 1
and consists of 2 sets of experiments: local experiments on patches and global
experiments on entire images. For the local experiments, we conducted 4
experiments of the E-PUQ procedure (detailed in Section 4.2.1) on RGB patch
resolutions of 1x1, 2x2, 4x4, and 8x8. We used $K=3$, $K=12$, $K=48$, and
$K=192$ PCs for each resolution, respectively. We set $\alpha=\delta=0.1$ to
be the user-specified parameters of the guarantee, defined in Equation (6). In
addition, we conducted another 2 experiments of the DA-PUQ (detailed in
Section 4.2.2), and RDA-PUQ (detailed in Appendix D) procedures on RGB patch
resolution of 8x8. We set $q=0.9$, $\beta=0.05$ and $\alpha=\delta=0.1$, to be
the user-specified parameters of the guarantees of both Equation (12) and
Equation (19). In total, we conducted $18$ local experiments across three
tasks. For the global experiments, we used entire images at a resolution of
128x128, in which we applied the DA-PUQ and the RDA-PUQ procedures. As global
working is suitable for tasks that exhibit strong pixel correlation, we
applied these experiment only on the task of image colorization. We set
$q=0.95$, $\beta=\alpha=\delta=0.1$, to be the user-specified parameters of
the guarantees of both Equation (12) and Equation (19). Both locally and
globally, for the DA-PUQ and RDA-PUQ experiments, we used $K=100$ PCs in the
colorization task and $K=200$ PCs in super-resolution and inpainting. We note
that in the RDA-PUQ experiments, we used $\hat{K}$ PCs during inference, as
discussed in Appendix D. In all experiments we used $\epsilon=1e-10$ for the
computation of the uncertainty volume, defined in Equation (3).