SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Elias Nehme
seliasne@campus.technion.ac.il
Omer Yair
omeryair@campus.technion.ac.il
Uncertainty quantification is crucial for the deployment of image restoration models in safety-critical domains, like autonomous driving and biological imaging. To date, methods for uncertainty visualization have mainly focused on per-pixel estimates. Yet, a heatmap of per-pixel variances is typically of little practical use, as it does not capture the strong correlations between pixels. A more natural measure of uncertainty corresponds to the variances along the principal components (PCs) of the posterior distribution. Theoretically, the PCs can be computed by applying PCA on samples generated from a conditional generative model for the input image. However, this requires generating a very large number of samples at test time, which is painfully slow with the current state-of-the-art (diffusion) models. In this work, we present a method for predicting the PCs of the posterior distribution for any input image, in a single forward pass of a neural network. Our method can either wrap around a pre-trained model that was trained to minimize the mean square error (MSE), or can be trained from scratch to output both a predicted image and the posterior PCs. We showcase our method on multiple inverse problems in imaging, including denoising, inpainting, super-resolution, colorization, and biological image-to-image translation. Our method reliably conveys instance-adaptive uncertainty directions, achieving uncertainty quantification comparable with posterior samplers while being orders of magnitude faster. Code and examples are available on our webpage.
Reliable uncertainty quantification is central to making informed decisions when using predictive models. This is especially important in domains with high stakes such as autonomous cars and biological/medical imaging, where based on visual data, the system is asked to provide predictions that could influence human life. In such domains, communicating predictive uncertainty becomes a necessity. In the particular case of image-to-image inverse problems, efficient visualization of predictive uncertainty is required. For example, in biological and medical image-to-image translation [7, 35, 37, 38], the predicted image as a whole is supposed to inform a scientific discovery or affect the diagnosis of a patient. Hence, an effective form of uncertainty to consider in this case is semantically-coordinated pixel variations that could alter the output image.
Currently, the majority of existing methods handle uncertainty in image-valued inverse problems by factorizing the output posterior into per-pixel marginals, in which case the strong correlations between pixels are completely ignored. This leads to per-pixel uncertainty estimates in the form
Figure 1: Comparison between per-pixel methods and NPPC. In structured output posteriors such as in image-to-image regression (left), uncertainty visualization in the form of per-pixel variance maps (top) fail to convey semantic variations in the prediction, leaving the user with no clue regarding the different possibilities of the solution set. Our method - NPPC (bottom) captures a more natural measure of uncertainty, by providing the user with input-adaptive PCs around the mean prediction. These can then be used to navigate the different possibilities in a meaningful manner. e.g., in 1a NPPC is used to traverse the existence of uncertain cells in an image-to-image translation task from biological imaging. Panel 1b presents a 2D motivating example illustrating the practical benefit of moving along PCs as opposed to the "standard" coordinate system (e.g., pixels in an image).
of variance heatmaps [20] or confidence intervals [2], which are of little practical use as they do not describe the joint uncertainty of different pixels. As a result, such models suffer from model misspecification, absorbing inter-pixel covariates into per-pixel margins, and presenting the user with unnecessarily inflated uncertainty estimates (see Fig. 1b). For example, consider a set of pixels along an edge in an image, in the task of image super-resolution. It is clear that for severe upsampling factors, a model would typically be uncertain around object edges. However, since these pixels are all correlated with the same object, the underlying uncertainty is whether the entire edge is shifted jointly, and not whether individual pixels can be shifted independently.
Recently, the field of image-to-image regression has seen a surge of probabilistic methods inspired by advancements in deep generative models, such as diffusion models (DMs). The recurring theme in these approaches is the use of a deep generative prior in order to sample from the posterior distribution [8, 19, 50]. In principle, using the resulting samples, it is possible to present the practitioner with the main modes of variations of the posterior uncertainty, for example using principal components analysis (PCA). However, these powerful posterior samplers come with a heavy computational price. Despite tremendous efforts over the past few years [27, 31, 43, 51], sampling from state-of-the-art methods is still unacceptably slow for many practical applications.
In this work, we bypass the need to learn a data-hungry and extremely slow conditional generative model, by directly training a neural model to predict the principal components of the posterior. Our method, which we term neural posterior principal components (NPPC), can wrap around any pretrained model that was originally trained to minimize the mean square error (MSE), or alternatively can be trained to jointly predict the conditional mean alongside the posterior principal components. In particular, when used as a post-hoc method with MSE pre-trained models, NPPC is a general technique that can be transferred across datasets and model architectures seamlessly. This is because our network architecture inherits the structure and the learning hyper-parameters of the pre-trained model with only one key change: Increasing the filter count at the last layer and employing a Gram-Schmidt process to the output to ensure our found principal components are orthogonal by construction. NPPC is then trained on the pre-trained model residuals, to endow the model's point prediction with efficient input-adaptive uncertainty visualization.
Several prior works proposed modeling the posterior as a correlated Gaussian distribution. In particular, Dorta et al. [9] approximated the output precision matrix by a matrix of the form  $\mathbf{L}\mathbf{L}^T$ , where  $\mathbf{L}$  is a lower triangular matrix with a pre-determined sparsity pattern. More recently, Monteiro et al. [33] and Meng et al. [30] approximated the covariance matrix by a sum of a diagonal matrix and a low-rank one. However, unlike NPPC, both [9] and [30] are limited to very low-resolution
images, because either they explicitly construct the (huge) covariance matrix during training or they capture only short-range correlations. Moreover, [33] suffers from training instability as multivariate Gaussians with both unknown mean and covariance are known to be notoriously unstable [46]. This limits its applicability to image segmentation, requiring proper background masking to avoid infinite covariance and overflow errors. Our method, on the other hand, does not construct the covariance matrix at any point. It directly outputs the top eigenvectors of the covariance and trains using the objective of PCA. As a result, our method is generally applicable to any inverse problem, enjoys faster and more stable training, and can handle high-resolution images as we demonstrate in Sec. 4. This is enabled by several design choices introduced to the architecture and training, including a Gram-Schmidt output layer, a PCA loss function maximizing correlation with the residuals, and a Stopgrad trick that enables learning all principal directions jointly, mitigating the need for training separate models for separate PCs.
We compare our PCs to those extracted from samples from a conditional generative model as recently proposed in [4]. As we show, we obtain comparable results, but orders of magnitude faster. Finally, we showcase NPPC on multiple inverse problems in imaging showing promising results across tasks and datasets. In particular, we apply NPPC to scientific biological datasets, showing the practical benefit of our uncertainty estimates in their ability to capture output correlations.
Model uncertainty Some methods attempt to report uncertainty that is due to model misspecification and/or out-of-distribution data. Early work on quantification of such uncertainty in deep models has focused on Bayesian modeling by imposing distributions on model weights/feature activations [6, 17, 34]. These methods employ various techniques for approximating the posterior of the weights, including using MCMC [44], variational inference [6, 25], Monte-Carlo Dropout [13], and Laplace approximation [36]. Our work rather focuses on data uncertainty as described in Section 3.
Per-pixel methods Some methods for predicting per-pixel variances assume a Gaussian distribution and learn it by maximizing the log-likelihood [20]. This approach was later combined with hierarchical priors on likelihood parameters to infer a family of distributions in a single deterministic neural network [1, 28]. Similarly, multi-modal approximations assuming a Gaussian Mixture Model per pixel have been proposed [5]. Deep Ensembles [22] and Test-time augmentations [53] have also been used to estimate uncertainty by measuring the variance of model predictions. More recently, distribution-free methods such as quantile regression and conformal prediction [39] have taken over their counterparts with firm statistical guarantees. Specifically for image-to-image regression, a work that stands out is that of Angelopoulos et al. [2]. The main shortcoming of these methods is the underlying assumption of i.i.d. pixels.
Distribution-free risk control Elaborate distribution-free methods, such as Risk Controlling Prediction Sets (RCPS) [3], take pixel correlations into account by controlling some risk factorizing all pixels in an image (such as the false discovery rate in binary image segmentation). Recently, this approach has been deployed in image-to-image regression problems using a technique called Conformal Prediction Masks [21]. The main drawback of these methods is their inability to explore the different possible options, but rather, only impose upper and lower bounds on the possible solution set. In addition, they also require an extra data split for calibration, which is not always readily available. An interesting recent work in this field [45] has proposed to control the risk of disentangled factors in the latent space of StyleGAN, and demonstrated informative uncertainty visualizations in inverse problems of facial images. However, the main drawbacks of this work are the key assumption of access to disentangled latent representations and the generalization gap from fake to real images.
Generative models and posterior samplers Generative models have been widely used to navigate prediction uncertainty, either in the form of conditional variational autoencoders [47] and conditional GANs[32], or more recently using state-of-the-art score-based and denoising diffusion models [8, 12, 14, 19, 26, 41, 42, 48, 49, 54]. While the latter have achieved astounding results in the last two years, when aiming for high-quality samples, they remain extremely slow to sample from, despite promising results reported in recent efforts [27, 31, 43, 51]. The recent Conffusion method [15] finetunes pre-trained diffusion models to output interval bounds in a single forward pass to achieve fast confidence interval prediction for posterior samplers. Nonetheless, the result is a per-pixel uncertainty map.
Figure 2: Method overview. Here, NPPC is demonstrated for image inpainting, where  $\pmb{x}$  is the ground truth image,  $\pmb{y}$  is the masked input, and  $\hat{\pmb{x}}$  is the posterior mean prediction. Our method can wrap around pre-trained conditional mean predictors  $f(\pmb{y};\pmb{\theta})$ , replicating their architecture with a slight modification at the output (see text). Using a PCA loss on the errors  $(e = x - \hat{x})$ , we learn to predict the first  $K$  PCs of the posterior  $w_{1},\ldots ,w_{K}$  directly in a single forward pass of a neural network  $w(y,\hat{x};\varphi)$ . On the right, we visualize the uncertainty captured by two PCs  $(w_{1},w_{4})$  around the mean prediction  $\hat{x}$ .
Gaussian covariance approximation Works that are particularly related to the approach presented here are [9], [33] and [30]. Specifically, Dorta et al. [9] proposed to model the output posterior with a correlated Gaussian, and approximated the precision matrix using the factorization  $\mathbf{\Lambda} = \mathbf{L}\mathbf{L}^{\top}$ , where  $\mathbf{L}$  is a lower triangular matrix with predefined sparsity patterns capturing short-term correlations. The matrix  $\mathbf{L}$  is then learned during training using a maximum likelihood objective. Similarly, [33] and [30] suggested approximating the posterior covariance matrix using a low-rank factorization with the latter involving second-order scores. However, the main drawback of [9] and [30] is their limited ability to generalize to high-resolution RGB images while efficiently capturing long-range correlations [10]. Additionally, [33] is limited by training instability, limiting its applicability to image segmentation. In contrast, our method directly estimates eigenvectors of the covariance matrix without the need for explicitly storing the covariance matrix in memory during training as in [30]. This enables us to seamlessly generalize to arbitrary inverse problems and high-resolution RGB images with more stable and less expensive training steps.
Concurrent work A related approach was recently proposed in the concurrent work of [29], where the authors demonstrate a training-free method to calculate the posterior PCs in the task of Gaussian denoising. The key advantage of our approach over this work is that we are not constrained to the task of Gaussian denoising and can handle arbitrary inverse problems, as we show in our experiments.
3 Neural Posterior Principal Components
We address the problem of predicting a signal  $\pmb{x} \in \mathbb{R}^{d_x}$  based on measurements  $\pmb{y} \in \mathbb{R}^{d_y}$ . In the context of imaging,  $\pmb{y}$  often represents a degraded version of  $\pmb{x}$  (e.g., noisy, blurry), or a measurement of the same specimen/scene acquired by a different modality. We assume that  $\pmb{x}$  and  $\pmb{y}$  are realizations of random vectors  $\mathbf{x}$  and  $\mathbf{y}$  with an unknown joint distribution  $p(\pmb{x},\pmb{y})$ , and that we have a training set  $\mathcal{D} = \{(x_i,y_i)\}_{i=1}^{N_d}$  of matched input-output pairs independently sampled from that distribution.
Many image restoration methods output a single prediction  $\hat{\pmb{x}}$  for any given input  $\pmb{y}$ . A common choice is to aim for the posterior mean  $\hat{\pmb{x}} = \mathbb{E}[\mathbf{x}|\mathbf{y} = \pmb{y}]$ , which is the predictor that minimizes the MSE. However, a single prediction does not convey to the user the uncertainty in the restoration. To achieve this goal, here we propose to also output the top  $K$  principal components of the posterior  $p(\pmb{x}|\pmb{y})$ , i.e., the top  $K$  eigenvectors of the posterior covariance  $\mathbb{E}[(\mathbf{x} - \hat{\mathbf{x}})(\mathbf{x} - \hat{\mathbf{x}})^{\top}|\mathbf{y} = \mathbf{y}]$ . The PCs capture the main directions along which  $\mathbf{x}$  could vary, given the input  $\pmb{y}$ , and thus provide the user with valuable information. However, their direct computation is computationally infeasible in high-dimensional settings. The main challenge we face is, therefore, how to obtain the posterior PCs without having to ever store or even compute the entire posterior covariance matrix. Before we describe our method, let us recall the properties of PCA in the standard (unconditional) setting.
Unconditional PCA Given a set of  $N$  data points  $\{\pmb{x}_i\}$  where  $\pmb{x}_i \in \mathbb{R}^{d_x}$ , the goal in PCA is to find a set of  $K$  orthogonal principal directions  $\pmb{w}_1, \dots, \pmb{w}_K$  along which the data varies the most. The resulting directions are ordered such that the variance of  $\{\pmb{w}_1^\top \pmb{x}_i\}$  is biggest, the variance of
$\{\pmb{w}_2^\top \pmb{x}_i\}$  is the second biggest, etc. The objective function for finding these directions has multiple equivalent forms and interpretations. The one we exploit here is the iterative maximization of variance. Specifically, let  $\pmb{X}$  denote the centered data matrix whose rows consist of different observations  $\pmb{x}_i^\top$  after subtracting their column-wise mean, and let  $\pmb{w}_1,\dots ,\pmb{w}_K$  denote the first  $K$  PCs. Then the  $k^{\mathrm{th}}$  PC is given by
$$
\boldsymbol {w} _ {k} = \underset {\boldsymbol {w}} {\arg \max } \| \boldsymbol {X} \boldsymbol {w} \| _ {2} ^ {2}, \quad \text {s . t .} \quad \| \boldsymbol {w} \| = 1, \boldsymbol {w} \perp \operatorname {s p a n} \left\{\boldsymbol {w} _ {1}, \dots , \boldsymbol {w} _ {k - 1} \right\}, \tag {1}
$$
where for the first PC, the constraint is only  $\| \pmb{w} \| = 1$ . The variance along the  $k^{\text{th}}$  PC is given by  $\sigma_k^2 = \frac{1}{N} \| \pmb{X} \pmb{w}_k \|^2$ .
The challenge in posterior PCA Going back to our goal of computing the posterior PCs, let us assume for simplicity we are given a pre-trained conditional mean predictor  $\hat{\pmb{x}} = f(\pmb{y};\pmb{\theta})$  obtained through MSE minimization on the training set  $\mathcal{D}$ . Let  $\pmb{e}_i = \pmb{x}_i - \hat{\pmb{x}}_i$  denote the error of the conditional mean predictor given the  $i^{\mathrm{th}}$  measurement,  $\mathbf{y} = \mathbf{y}_i$ . The uncertainty directions we wish to capture per sample  $\pmb{y}_i$ , are the PCs of the error  $\pmb{e}_i$  around the conditional mean estimate  $\hat{\pmb{x}}_i$ . Conceptually, this task is extremely difficult as during training for every measurement  $\pmb{y}_i$ , we have access only to a single error sample  $\pmb{e}_i$ . If we were to estimate the PCs directly, the result would be a trivial single principal direction equaling  $\pmb{e}_i$  which is unpredictable at test time. To address this challenge, here we propose to harness the implicit bias of neural models and to learn these directions from a dataset of triplets  $\mathcal{D}' = \{(x_i,y_i,\hat{x}_i)\}_{i=1}^{N_d}$ . The key implicit assumption underlying our approach (and empirical risk minimization in general) is that the posterior mean  $\mathbb{E}[\mathbf{x}|\mathbf{y} = \mathbf{y}]$  and the posterior covariance  $\mathbb{E}[(\mathbf{x} - \hat{\mathbf{x}})(\mathbf{x} - \hat{\mathbf{x}})^{\top}|\mathbf{y} = \mathbf{y}]$  vary smoothly with  $\pmb{y}$ . Hence, with the right architecture, such models can capitalize on inter-sample dependencies and properly generalize to unseen test points, by learning posterior PCs that change gracefully as a function of  $\pmb{y}$ . This is much like models trained with MSE minimization to estimate the conditional mean  $\mathbb{E}[\mathbf{x}|\mathbf{y} = \mathbf{y}]$ , while being presented during training only with a single output  $\pmb{x}_i$  for every measurement  $\pmb{y}_i$ .
3.1 Naive solution: Iterative learning of PCs
Following the intuition from the previous section, we can parameterize the  $k^{\mathrm{th}}$  PC of the error using a neural network  $\pmb{w}_k(\pmb{y},\hat{\pmb{x}};\varphi_k)$  with parameters  $\varphi_{k}$ , which has similar capacity to the pre-trained model  $f(\pmb {y};\pmb {\theta})$  outputting the conditional mean. This model accepts the measurement  $\pmb{y}$  and (optionally) the conditional mean estimate  $\hat{\pmb{x}}$ , and outputs the  $k^{\mathrm{th}}$  PC of the error  $\pmb{e}$ .
Let  $\pmb{d}_1(\pmb{y},\hat{\pmb{x}};\varphi_1)$  be a model for predicting the first unnormalized direction, such that  $\pmb{w}_1(\pmb{y}_i,\hat{\pmb{x}}_i;\varphi_1) = \pmb{d}_1(\pmb{y}_i,\hat{\pmb{x}}_i;\varphi_1) / \| \pmb{d}_1(\pmb{y}_i,\hat{\pmb{x}}_i;\varphi_1)\|$ . Given a dataset of triplets  $\mathcal{D}' = \{(\pmb{x}_i,\pmb{y}_i,\hat{\pmb{x}}_i)\}$ , we adopt the objective employed in (1), and propose to learn the parameters of the input-dependent first PC by minimizing
$$
\mathcal {L} _ {\boldsymbol {w} _ {1}} \left(\mathcal {D} ^ {\prime}, \varphi_ {1}\right) = - \sum_ {\left(\boldsymbol {x} _ {i}, \boldsymbol {y} _ {i}, \hat {\boldsymbol {x}} _ {i} \in \mathcal {D} ^ {\prime}\right)} \left| \boldsymbol {w} _ {1} \left(\boldsymbol {y} _ {i}, \hat {\boldsymbol {x}} _ {i}; \varphi_ {1}\right) ^ {\top} \boldsymbol {e} _ {i} \right| ^ {2}. \tag {2}
$$
Next, given the model predicting the first PC, we can train a model to predict the second PC,  $\pmb{w}_2(\pmb{y}_i, \hat{\pmb{x}}_i; \pmb{\varphi}_2)$ , by manipulating the output of a model  $\pmb{d}_2(\pmb{y}_i, \hat{\pmb{x}}_i; \pmb{\varphi}_2)$ . Specifically, following the approach in (1), we optimize the same loss (2), but construct the output of the model  $\pmb{w}_2$  by removing the projection of  $\pmb{d}_2$  onto  $\pmb{w}_1$ , and normalizing the result. This ensures that  $\pmb{w}_2(\pmb{y}_i, \hat{\pmb{x}}_i; \pmb{\varphi}_2) \perp \pmb{w}_1(\pmb{y}_i, \hat{\pmb{x}}_i; \pmb{\varphi}_1)$ .
While in principle this approach can be iterated  $K$  times to learn the first  $K$  PCs, it has several drawbacks that make it impractical. First, it requires a prolonged iterative training of  $K$  neural networks sequentially, preventing parallelization and leading to very long training times. Second, this approach is also inefficient at test time, as we need to compute  $K$  dependent forward passes. Finally, in this current formulation, different PCs have their own set of weights  $\{\varphi_k\}_{k=1}^K$ , and do not share parameters. This is inefficient as for a given input  $\mathbf{y}$  and corresponding mean prediction  $\hat{\mathbf{x}}$ , it is expected that the initial feature extraction stage for predicting the different PCs, would be similar. A better strategy is therefore to design an architecture that outputs all PCs at once.
3.2 Joint learning of PCs
To jointly learn the first  $K$  PCs of the error using a single neural network  $\pmb{w}(\pmb{y}_i, \hat{\pmb{x}}_i; \varphi)$ , we introduce two key changes to the architecture inherited from the pre-trained mean estimator  $f(\pmb{y}; \pmb{\theta})$ . First,
the number of filters at the output layer is multiplied by  $K$ , to accommodate the  $K$  PCs,  $\pmb{w}_1,\dots ,\pmb{w}_K$ . Second, we introduce a Gram-Schmidt procedure at the output layer, making the directions satisfy the orthogonality constraint by construction. Formally, denote the non-orthonormal predicted set of directions by  $d_{1},\ldots ,d_{K}$ . Then, we transform them to be an orthonormal set  $\pmb{w}_1,\dots ,\pmb{w}_K$  as
$$
\boldsymbol {w} _ {1} = \frac {\boldsymbol {d} _ {1}}{\| \boldsymbol {d} _ {1} \|},
$$
$$
\boldsymbol {w} _ {k} = \frac {\boldsymbol {d} _ {k} - \sum_ {\ell = 1} ^ {k - 1} \left(\boldsymbol {d} _ {k} ^ {\top} \boldsymbol {w} _ {\ell}\right) \boldsymbol {w} _ {\ell}}{\left\| \boldsymbol {d} _ {k} - \sum_ {\ell = 1} ^ {k - 1} \left(\boldsymbol {d} _ {k} ^ {\top} \boldsymbol {w} _ {\ell}\right) \boldsymbol {w} _ {\ell} \right\|}, \quad k = 2, \dots K. \tag {3}
$$
Note, however, that these changes do not yet guarantee proper learning of the PCs. Indeed, if we were to learn the directions by naively minimizing the loss
$$
\mathcal {L} _ {\boldsymbol {w}} \left(\mathcal {D} ^ {\prime}, \varphi\right) = - \sum_ {\left(\boldsymbol {x} _ {i}, \boldsymbol {y} _ {i}, \hat {\boldsymbol {x}} _ {i} \in \mathcal {D} ^ {\prime}\right)} \sum_ {k = 1} ^ {K} \left| \boldsymbol {w} _ {k} \left(\boldsymbol {y} _ {i}, \hat {\boldsymbol {x}} _ {i}; \varphi\right) ^ {\top} \boldsymbol {e} _ {i} \right| ^ {2}, \tag {4}
$$