Principal Uncertainty Quantification with Spatial Correlation for Image Restoration Problems

TL;DR

This work introduces Principal Uncertainty Quantification (PUQ), a spatially aware framework for uncertainty quantification in image restoration that builds uncertainty regions around principal components of the empirical posterior $\\mathbb{P}_{y|x}$. It leverages diffusion-based samplers to generate candidate restorations, extracts an instance-adaptive PCA basis, and uses conformal-style Learn-Then-Test calibrations to guarantee coverage with user-defined levels; two main configurations are Exact PUQ (E-PUQ) and Dimension-Adaptive PUQ (DA-PUQ), with a further Reduced variant for efficiency. The approach yields significantly tighter uncertainty volumes than pixelwise baselines across colorization, super-resolution, and inpainting, while preserving statistical guarantees and offering improved interpretability by using only a small number of PCs in the adaptive settings. By accounting for spatial correlations and enabling local (patch) or global quantification, PUQ provides practical, certified uncertainty regions for high-dimensional image restoration tasks and lays groundwork for broader use of adaptive linear bases in inverse problems.

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

Figures (21)

• Figure 1: Comparison of PUQ's performance on the CelebA-HQ dataset in image colorization, super-resolution, and inpainting tasks using the E-PUQ procedure (Section \ref{['sec: Exact PUQ']}) applied on RGB image patches of varying size. As seen, our method provides tighter uncertainty regions with significantly smaller uncertainty volumes ($\times10$ in super-res. and inpainting, and $\times 100$ in colorization). The compared methods are im2im-uq angelopoulos2022image and Conffusion horwitz2022conffusion.
• 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 \ref{['eq: intervals']}, 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 \ref{['eq: uncertainty volume']}, is indicated in the top left corner of each plot. The $90\%$ coverage guarantee, outlined in Equation \ref{['eq: theoretical coverage gaurentee']} with $w_i := 1/2$, is satisfied by all. As can be seen, the lower row regions take spatial dependencies into account and are significantly smaller than the pixelwise corresponding regions in the upper row.
• 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 \ref{['sec: our implementation']}.
• 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 \ref{['eq: intervals']}. Finally, the intervals are scaled to statistically guarantee Equation \ref{['eq: theoretical coverage gaurentee']} and contain the correct ratio over solutions for unseen input instances. In the bottom row, the procedure also statistically guarantees Equation \ref{['eq: theoretical reconstruction gaurentee']} by ensuring a small recovery error of solutions to unseen input instances, as demonstrated by the small variance around the single PC.
• 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.