Title: Uncertainty Quantification via Neural Posterior Principal Components
Venue: NeurIPS 2023
Accepted: yes
Keywords: Uncertainty Quantification

Inverse Problems

Probabilistic Modelling

Principal Components Analysis

Deep Learning
Abstract: 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, 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](https://eliasnehme.github.io/NPPC/).

Uncertainty Quantification via Neural Posterior Principal Components

Abstract

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 instanceadaptive uncertainty directions, achieving uncertainty quantification comparable with posterior samplers while being orders of magnitude faster. Code and examples are available on our webpage.

1 Introduction

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

Input

GT

Per-pixel s.t.d.

Samples

PC

-PC

Prediction

+PC

(a)

True distribution

Principal components

Per-pixel variance

NPPC

(b)

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 marginals, 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 LL T , where 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.

2 Related Work

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.

Pre-trained

Neural Posterior Principal Components

PCA Loss

h

K

∑lw (x -a)|2

k=1←

WK

w1

w1

=

ILpll

Wk

p）=-p

we)we

we)we

-20kWk

+20kWk

=

k

=

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 Λ = LL ⊤ , where L is a lower triangular matrix with predefined sparsity patterns capturing short-term correlations. The matrix 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 x ∈ R d x based on measurements y ∈ R d y . In the context of imaging, y often represents a degraded version of x ( e.g., noisy, blurry), or a measurement of the same specimen/scene acquired by a different modality. We assume that x and y are realizations of random vectors x and y with an unknown joint distribution p ( x , y ) , and that we have a training set D = { ( x i , y i ) } N d i =1 of matched input-output pairs independently sampled from that distribution.

Many image restoration methods output a single prediction ˆ x for any given input y . A common choice is to aim for the posterior mean ˆ x = E [ x | y = 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 ( x | y ) , i.e., the top K eigenvectors of the posterior covariance E [( x -ˆ x )( x -ˆ x ) ⊤ | y = y ] . The PCs capture the main directions along which x could vary, given the input 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 { x i } where x i ∈ R d x , the goal in PCA is to find a set of K orthogonal principal directions w 1 , . . . , w K along which the data varies the most. The resulting directions are ordered such that the variance of { w ⊤ 1 x i } is biggest, the variance of

{ w ⊤ 2 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 X denote the centered data matrix whose rows consist of different observations x ⊤ i after subtracting their column-wise mean, and let w 1 , . . . , w K denote the first K PCs. Then the k th PC is given by

where for the first PC, the constraint is only ∥ w ∥ = 1 . The variance along the k th PC is given by σ 2 k = 1 N ∥ Xw k ∥ 2 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 ˆ x = f ( y ; θ ) obtained through MSE minimization on the training set D . Let e i = x i -ˆ x i denote the error of the conditional mean predictor given the i th measurement, y = y i . The uncertainty directions we wish to capture per sample y i , are the PCs of the error e i around the conditional mean estimate ˆ x i . Conceptually, this task is extremely difficult as during training for every measurement y i , we have access only to a single error sample e i . If we were to estimate the PCs directly, the result would be a trivial single principal direction equaling 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 D ′ = { ( x i , y i , ˆ x i ) } N d i =1 . The key implicit assumption underlying our approach (and empirical risk minimization in general) is that the posterior mean E [ x | y = y ] and the posterior covariance E [( x -ˆ x )( x -ˆ x ) ⊤ | y = y ] vary smoothly with 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 y . This is much like models trained with MSE minimization to estimate the conditional mean E [ x | y = y ] , while being presented during training only with a single output x i for every measurement y i .

3.1 Naive solution: Iterative learning of PCs

Following the intuition from the previous section, we can parameterize the k th PC of the error using a neural network w k ( y , ˆ x ; φ k ) with parameters φ k , which has similar capacity to the pre-trained model f ( y ; θ ) outputting the conditional mean. This model accepts the measurement y and (optionally) the conditional mean estimate ˆ x , and outputs the k th PC of the error e .

Let d 1 ( y , ˆ x ; φ 1 ) be a model for predicting the first unnormalized direction, such that w 1 ( y i , ˆ x i ; φ 1 ) = d 1 ( y i , ˆ x i ; φ 1 ) / ∥ d 1 ( y i , ˆ x i ; φ 1 ) ∥ . Given a dataset of triplets D ′ = { ( x i , y i , ˆ x i ) } , we adopt the objective employed in (1), and propose to learn the parameters of the input-dependent first PC by minimizing

Next, given the model predicting the first PC, we can train a model to predict the second PC, w 2 ( y i , ˆ x i ; φ 2 ) , by manipulating the output of a model d 2 ( y i , ˆ x i ; φ 2 ) . Specifically, following the approach in (1), we optimize the same loss (2), but construct the output of the model w 2 by removing the projection of d 2 onto w 1 , and normalizing the result. This ensures that w 2 ( y i , ˆ x i ; φ 2 ) ⊥ w 1 ( y i , ˆ x i ; φ 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 { φ k } K k =1 , and do not share parameters. This is inefficient as for a given input y and corresponding mean prediction ˆ 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 w ( y i , ˆ x i ; φ ) , we introduce two key changes to the architecture inherited from the pre-trained mean estimator f ( y ; θ ) . First,

the number of filters at the output layer is multiplied by K , to accomodate the K PCs, w 1 , . . . , 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 , . . . , d K . Then, we transform them to be an orthonormal set w 1 , . . . , w K as

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

then we would only recover them up to an orthogonal matrix. To see this, let W i denote the matrix that has w k ( y i , ˆ x i ; φ ) in its k th column. Then the inner sum in (4) can be rewritten as ∥ W ⊤ i e i ∥ 2 2 . Now, it is easy to see that neither the loss nor the constraints are affected by replacing each W i with ˜ W i = W i O i , for some orthogonal matrix O i . Indeed, this solution satisfies the orthogonality constraint ˜ W ⊤ i ˜ W i = I , and attains the same loss value as the original PCs.

Note that this rotation ambiguity did not exist in the naive approach of Sec. 3.1 because there, when finding the k th PC given the preceding k -1 pre-trained PCs, the loss term | w k ( y i , ˆ x i ; φ ) ⊤ e i | 2 did not affect the learning of w 1 , . . . , w k -1 . However, when attempting to learn all directions jointly, the preceding PCs receive a gradient signal from the k th loss term as w k is a function of w 1 , . . . , w k -1 in the Gram-Schmidt procedure.

To solve this problem and decouple the learning of the different PCs while still maintaining their orthogonality, we propose a simple modification to the Gram-Schmidt procedure: Using Stopgrad for the previously derived PCs within the projection operators in (3) (see the red term in Fig. 2). This way, in each learning step the different PCs are guaranteed to be orthogonal, while solely optimizing their respective objective. This allows learning them jointly and recovering the solution of the iterative scheme in a single forward pass of a neural network with shared parameters. Please see App. D.2 for validation of the equivalence between sequential and joint PC learning.

3.3 Learning variances along PCs

Recall that the variance along the k th direction corresponds to the average squared projection of the data along that direction. We can use that to output a prediction of the variances { σ 2 k } of the PCs, by using a loss that minimizes ∑ i ( σ 2 k -| w ⊤ k e i | 2 ) 2 for every k . However, instead of adding K additional outputs to the architecture, we can encode the variances in the norms of the unnomarlized directions. To achieve this without altering the optimization objective for finding the k th PC, we again invoke the Stopgrad operator on the term | w ⊤ k e i | which is the current loss function value, and match the norm of the found PCs at the current step to this projected variance by minimizing the loss

With this, our prediction for σ 2 k at test-time is simply ∥ d k -∑ k -1 ℓ =1 ( d ⊤ k w ℓ ) w ℓ ∥ 2 . Please see App. D.3 for quantitative validation of our estimated variances.

3.4 Joint prediction of posterior mean

Thus far, we assumed we have access to a pre-trained posterior mean predictor ˆ x = f ( y ; θ ) , obtained through MSE minimization on a dataset D = { ( x i , y i ) } N d i =1 of matched pairs. Hence, our entire derivation revolved around a dataset of triplets D ′ = { ( x i , y i , ˆ x i ) } N d i =1 . However, this is not strictly necessary. In particular, the posterior mean predictor can be learned jointly alongside the PCs with an

10

5

0

5

-

5

0

5

10

p

(

x

)

(

x

i

,

y

i

)

∼

p

(

x

,

y

)

y

i

x

i

10

5

0

5

D

=

{

(

x

i

,

y

i

)

}

y

i

x

i

10

5

0

5

p

(

x

|

y

1

)

µ

x

|

y

(

y

1

)

ˆ

µ

x

|

y

(

y

1

)

GT

Ours

y

1

10

5

0

5

p

(

x

|

y

2

)

µ

x

|

y

(

y

2

)

ˆ

µ

x

|

y

(

y

2

)

GT

Ours

y

2

0

.

000

0

.

006

0

.

012

0

.

018

0

.

024

0

.

030

0

.

036

0

.

000

0

.

006

0

.

012

0

.

018

0

.

024

0

.

030

0

.

036

0

.

000

0

.

008

0

.

016

0

.

024

0

.

032

0

.

040

0

.

048

0

.

056

0

.

064

0

.

072

0

.

000

0

.

018

0

.

036

0

.

054

0

.

072

0

.

090

0

.

108

0

.

126

-

-

-

-

-

-

-

-

additional MSE loss

Note that we do not want the gradients of L w and L σ to affect the recovered mean. Therefore, we use a Stopgrad operator when inputting the error x i -ˆ x i to those loss functions.

3.5 Putting it all together

To summarize, given a dataset D = { ( x i , y i ) } N d i =1 of matched pairs, we can learn a model that predicts the posterior mean as well as the posterior PCs, using the combined loss functions of eqs. (4)-(6),

where λ 1 and λ 2 are weights balancing the contributions of the different losses.

4 Experiments

We now illustrate NPPC on several tasks and datasets. In all experiments except for the toy example, we used variants of the U-Net architecture [11, 40]. The weighting factors for the losses were chosen such that all three terms are roughly within an order of magnitude of each other. Empirically, we find that ramping up the weight factors for the terms of the directions and the variances after the mean estimate started converging, stabilizes training. Full details regarding the architectures, the scheduler, and the per-task setting of λ 1 , λ 2 are in App. A.

Toy examples Figure 3 demonstrates NPPC on a 2D denoising task, where samples x i from a two-component Gaussian mixture model are contaminated by additive white Gaussian noise to result in noisy measurements y i . To predict x from y , we trained a 5-layer MLP with 256 hidden features using MSE minimization, and around the estimated conditional mean ˆ x = ˆ µ x | y ( y ) , we trained NPPC (instantiated with a similar MLP) to output the two PCs

w 1 , w 2 , and their variances ˆ σ 2 1 , ˆ σ 2 2 . As can be seen in the right two panels of Fig. 3, NPPC accurately predicts the ground-truth (GT) PCs (those are computed from the analytical expression for the posterior covariance; see App. B). To quantify the accuracy of the recovered PCs, we computed the Wasserstein 2-distance between a Gaussian approximation of the GT posterior p ( x | y ) ≈ N ( µ x | y , Σ x | y ) , and an estimated Gaussian constructed by NPPC, ˆ p ( x | y ) = N ( ˆ µ x | y , W ⋆ W ⊤ ⋆ ) , where W ⋆ has the scaled estimated PC ˆ σ k w k in its k th column. Compared to a baseline of a point mass at the estimated

ˆ

x

x

w

1

w

2

w

3

w

4

w

5

α

=

-

3

ˆ

x

+

ασ

1

w

1

ˆ

x

+

ασ

4

w

4

α

=

-

2

α

= 0

α

= 2

α

= 3

y

(a) Denoising

ˆ

x

x

w

1

w

2

w

3

w

4

w

5

α

=

-

3

ˆ

x

+

ασ

1

w

1

ˆ

x

+

ασ

5

w

5

α

=

-

2

α

= 0

α

= 2

α

= 3

y

(b) Inpainting

h

h

4

7

7

9

9

9

conditional mean, δ ( x -ˆ x ) , NPPC reduces the Wasserstein 2-distance by 100 × from 4.05 to 0.04. Similarly, we also applied NPPC to a 100-dimensional Gaussian mixture denoising task, and computed the Wasserstein 2-distance from a Gaussian distribution whose mean is the GT posterior mean and whose covariance is the best rankK approximation of the GT posterior covariance. As clearly evident in Table 1, NPPC maintains a roughly constant distance to the analytical posterior, while the point mass baseline distance rapidly grows with K . See App. B for more details.

Handwritten digits Figure 4 demonstrates NPPC on denoising and inpainting of handwritten digits from the MNIST dataset. In the denoising task, we used noise of standard deviation σ ε = 1 , and in inpainting we used a mask that covers the top 70% of the image. As can be seen, for denoising, the learned PCs capture both inter-digit and intra-digit variations, e.g., turning a '4' into a '9'. Similarly, in inpainting, the learned PCs traverse the two likely modes around the mean estimate ˆ x , going from a '7' into a '9'. More examples are available in App. D.

Faces To test NPPC on faces, we trained on the CelebA-HQ dataset using the original split inherited from celebA [24], resulting in 24183 images for training, 2993 images for validation, and 2824 images for testing. Figure 5 presents results on 256 × 256 face images from the CelebA-HQ dataset. Here, we demonstrate NPPC on the task of inpainting, as well as on noisy 8 × super-resolution with a box downsampling filter and a noise-level of σ ε = 0 . 05 . As can be seen, the PCs generated by NPPC capture semantically meaningful uncertainty, corresponding to the eyes, mouth, and eyebrows. We also tested our approach on inpainting of the eyes area, on 4 × super-resolution (noisy and noiseless), and on image colorization. More results and examples are provided in App. D.

Comparison to posterior samplers While traversal along the predicted PCs is valuable for qualitative analysis, an important question is how our single forward pass PCs fair against state-of-the-art posterior samplers. To test this, we now compare NPPC to recent posterior samplers on CelebA-HQ 256 × 256 . Specifically, we compare NPPC to DDRM [19], DDNM [54], RePaint [26], and MAT [23] on the tasks of image super-resolution/inpainting. We use each of these methods to generate 100 samples per test image, and compute PCA on those samples. We perform comparisons over 100 test images randomly sampled from the FFHQ dataset [18]. Note that NPPC and MAT were trained on the training images of CelebA-HQ according to the original split of CelebA, whereas DDRM, DDNM, and RePaint rely on DDPM [14] as a generative prior, which was trained on the entirety of CelebA-HQ. The results are reported in Table 2. As can be seen, NPPC achieves similar root MSE (RMSE) ∥ x -ˆ x ∥ 2 , and residual error magnitude ∥ e -WW ⊤ e ∥ 2 using the first K = 5 PCs. This is while being 100 × faster than MAT, and 10 3 -10 5 × faster than diffusion-based methods. It is interesting to note that the principal angle between the subspace computed by NPPC and that computed by the baselines, is typically ≈ 90 ◦ . Namely, they are nearly orthogonal. This can be attributed to the severe ill-posedness of the tested settings, resulting in multiple different PCs with roughly the same projected variance. This is also evident in the residual error magnitude, suggesting that a very large number of PCs is required to handle such extreme degradations (see App. C).

Biological image-to-image translation image-to-image translation refers to translating images from one domain into another [16]. In biological imaging, image-to-image translation has been

ˆ

x

x

w

1

w

2

w

3

w

4

w

5

α

=

-

2

ˆ

x

+

ασ

2

w

2

ˆ

x

+

ασ

3

w

3

α

=

-

1

α

= 0

α

= 1

α

= 2

y

ˆ

x

x

w

1

w

2

w

3

w

4

w

5

α

=

-

2

ˆ

x

+

ασ

1

w

1

ˆ

x

+

ασ

5

w

5

α

=

-

1

α

= 0

α

= 1

α

= 2

y

90

B

B

(a) Inpainting

(b) Super-resolution

applied in several contexts, including for predicting fluorescent labels from bright-field images [35] and for predicting one fluorescent label ( e.g., nuclear stainings) from another ( e.g., actin stainings) [52]. However, unlike its use for artistic purposes, the use of image-to-image translation in biological imaging requires caution. Specifically, without proper uncertainty quantification, predicting cell nuclei from other fluorescent labels could lead to biased conclusions, as cell counting and tracking play central roles in microscopy-based scientific experiments ( e.g., drug testing). Here, we applied NPPC to a dataset of migrating cells imaged live for 14h (1 picture every 10min) using a spinning-disk microscope [52]. The dataset consisted of 1753 image pairs of resolution 1024 × 1024 , out of which 1748 were used for training, and 5 were used for testing following the original split by the authors. We started by training a standard U-Net [11, 40] model using MSE minimization to predict nuclear stainings from actin stainings (see Fig. 6), and then trained NPPC on the residuals. As we show in Fig. 6, the PCs learned by NPPC convey important information to experimenters. For example, the first PC highlights the fact that accurate intensity estimation is not possible in this task, and thereby a global bias is a valid uncertainty component. Furthermore, the remaining PCs reflect semantic uncertainty by adding/removing cells from the conditional mean estimate ˆ x , thereby clearly communicating reconstruction ambiguity to the user.

5 Discussion

We proposed an approach for directly predicting posterior PCs and showed its applicability across multiple tasks and datasets. Nonetheless, our method does not come without limitations. First, as evident by both the PCs of NPPC and of posterior samplers, for severely ill-posed inverse problems a linear subspace with a small number of PCs captures very little of the error. Hence, a large number of PCs is required to faithfully reconstruct the error. However, the main premise of this work was scientific imaging where scientists usually take reliable measurements, and the uncertainty is not as severe as ultimately the result should drive scientific discovery. Second, throughout this paper, we

ˆ

x

x

w

1

w

2

w

3

w

4

w

5

α

=

-

3

ˆ

x

+

ασ

2

w

2

ˆ

x

+

ασ

5

w

5

α

=

-

2

α

= 0

α

= 2

α

= 3

y

learned the PCs on the same training set of the conditional mean estimator. However, generalization error between training and validation could lead to biased PCs estimation as the learning process is usually stopped after some amount of overfitting ( e.g., until the validation loss stagnates). This can be solved by using an additional data split for 'calibration' as is typically the case with distribution-free methods [2, 3, 39, 45]. Third, different inputs y i may have posteriors of different complexities, and may thus require a different number of PCs K . In our approach, K is hard-coded within the network's architecture (it is the number of network outputs). Namely, we treat K as a hyper-parameter that needs to be set in advance prior to training, acting as an upper bound on the number of directions the user may be interested in exploring at test time. However, recall that NPPC also predicts the standard deviations along each of the K PCs. These may serve to decide whether to ignore some of the PCs. Fourth, our learned PCs cover the entire image, whereas in some cases the interesting uncertainty structure could be local ( e.g., cell data). In such circumstances, NPPC should either be applied patch-wise or the number of PCs K should be sufficiently increased. Finally, our method is tailored towards capturing uncertainty and not sampling from the posterior. While we provide the directed standard deviations, shifting along a certain direction does not guarantee samples along the way to be on the image data manifold. This can be tackled by employing NPPC in the latent space of a powerful encoder, in which case small enough steps could lead to changes on the manifold in the output image. However, this is beyond the scope of this current work.

6 Conclusion

To conclude, in this work we proposed a technique for directly estimating the principal components of the posterior distribution using a single forward pass in a neural network. We discussed key design choices, including a Gram-Schmidt procedure with a Stopgrad operator, and a principled PCA loss function. Through extensive experiments, we validated NPPC across tasks and domains, showing its wide applicability. We showed that NPPC achieves comparable uncertainty quantification to the naive approach of applying PCA on samples generated by posterior samplers while being orders of magnitude faster. Finally, we applied NPPC to the challenging task of biological image-to-image translation, demonstrating its practical benefit for safety-critical applications. In terms of broader impact, proper uncertainty quantification is crucial for trustworthy interpretable systems, particularly in healthcare applications. Thus, a method for reporting and conveniently visualizing prediction uncertainty could support users and help them avoid making flawed decisions.

References

Alexander Amini, Wilko Schwarting, Ava Soleimany, and Daniela Rus. Deep evidential regression. Advances in Neural Information Processing Systems , 33:14927-14937, 2020.

Anastasios N Angelopoulos, Amit P Kohli, Stephen Bates, Michael I Jordan, Jitendra Malik, Thayer Alshaabi, Srigokul Upadhyayula, and Yaniv Romano. Image-to-image regression with distribution-free uncertainty quantification and applications in imaging. In International Conference on Machine Learning , 2022.

Stephen Bates, Anastasios Angelopoulos, Lihua Lei, Jitendra Malik, and Michael Jordan. Distribution-free, risk-controlling prediction sets. Journal of the ACM (JACM) , 68(6):1-34, 2021.

Omer Belhasin, Yaniv Romano, Daniel Freedman, Ehud Rivlin, and Michael Elad. Principal uncertainty quantification with spatial correlation for image restoration problems. arXiv preprint arXiv:2305.10124 , 2023.

Christopher M Bishop. Mixture density networks. 1994.

Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning , pages 1613-1622. PMLR, 2015.

Eric M Christiansen, Samuel J Yang, D Michael Ando, Ashkan Javaherian, Gaia Skibinski, Scott Lipnick, Elliot Mount, Alison O'neil, Kevan Shah, Alicia K Lee, et al. In silico labeling: predicting fluorescent labels in unlabeled images. Cell , 173(3):792-803, 2018.

Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh International Conference on Learning Representations , 2022.

Garoe Dorta, Sara Vicente, Lourdes Agapito, Neill DF Campbell, and Ivor Simpson. Structured uncertainty prediction networks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 5477-5485, 2018.

Garoe Dorta, Sara Vicente, Lourdes Agapito, Neill DF Campbell, and Ivor Simpson. Training vaes under structured residuals. arXiv preprint arXiv:1804.01050 , 2018.

Thorsten Falk, Dominic Mai, Robert Bensch, Özgün Çiçek, Ahmed Abdulkadir, Yassine Marrakchi, Anton Böhm, Jan Deubner, Zoe Jäckel, Katharina Seiwald, et al. U-net: deep learning for cell counting, detection, and morphometry. Nature methods , 16(1):67-70, 2019.

Berthy T Feng, Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L Bouman, and William T Freeman. Score-based diffusion models as principled priors for inverse imaging. arXiv preprint arXiv:2304.11751 , 2023.

Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning , pages 1050-1059. PMLR, 2016.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Eliahu Horwitz and Yedid Hoshen. Conffusion: Confidence intervals for diffusion models. arXiv preprint arXiv:2211.09795 , 2022.

Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1125-1134, 2017.

Pavel Izmailov, Wesley J Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Subspace inference for bayesian deep learning. In Uncertainty in Artificial Intelligence , pages 1169-1179. PMLR, 2020.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4401-4410, 2019. URL https://github.com/NVlabs/ffhq-dataset.

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems , 2022.

Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems , 30, 2017.

Gilad Kutiel, Regev Cohen, Michael Elad, Daniel Freedman, and Ehud Rivlin. Conformal prediction masks: Visualizing uncertainty in medical imaging. In ICLR 2023 Workshop on Trustworthy Machine Learning for Healthcare .

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems , 30, 2017.

Wenbo Li, Zhe Lin, Kun Zhou, Lu Qi, Yi Wang, and Jiaya Jia. Mat: Mask-aware transformer for large hole image inpainting. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10758-10768, 2022.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE International Conference on Computer Vision , pages 3730-3738, 2015. URL http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html.

Christos Louizos and Max Welling. Multiplicative normalizing flows for variational bayesian neural networks. In International Conference on Machine Learning , pages 2218-2227. PMLR, 2017.

Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11461-11471, 2022.

Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388 , 2021.

Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks. Advances in neural information processing systems , 31, 2018.

Hila Manor and Tomer Michaeli. On the posterior distribution in denoising: Application to uncertainty quantification. arXiv preprint arXiv:2309.13598 , 2023.

Chenlin Meng, Yang Song, Wenzhe Li, and Stefano Ermon. Estimating high order gradients of the data distribution by denoising. Advances in Neural Information Processing Systems , 34: 25359-25369, 2021.

Chenlin Meng, Ruiqi Gao, Diederik P Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In NeurIPS 2022 Workshop on Score-Based Methods , 2022.

Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXiv preprint arXiv:1411.1784 , 2014.

Miguel Monteiro, Loïc Le Folgoc, Daniel Coelho de Castro, Nick Pawlowski, Bernardo Marques, Konstantinos Kamnitsas, Mark van der Wilk, and Ben Glocker. Stochastic segmentation networks: Modelling spatially correlated aleatoric uncertainty. Advances in neural information processing systems , 33:12756-12767, 2020.

Radford M Neal. Bayesian learning for neural networks , volume 118. Springer Science & Business Media, 2012.

Chawin Ounkomol, Sharmishtaa Seshamani, Mary M Maleckar, Forrest Collman, and Gregory R Johnson. Label-free prediction of three-dimensional fluorescence images from transmitted-light microscopy. Nature methods , 15(11):917-920, 2018.

Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018Conference Track Proceedings , volume 6. International Conference on Representation Learning, 2018.

Yair Rivenson, Tairan Liu, Zhensong Wei, Yibo Zhang, Kevin de Haan, and Aydogan Ozcan. Phasestain: the digital staining of label-free quantitative phase microscopy images using deep learning. Light: Science & Applications , 8(1):23, 2019.

Yair Rivenson, Hongda Wang, Zhensong Wei, Kevin de Haan, Yibo Zhang, Yichen Wu, Harun Günaydın, Jonathan E Zuckerman, Thomas Chong, Anthony E Sisk, et al. Virtual histological staining of unlabelled tissue-autofluorescence images via deep learning. Nature biomedical engineering , 3(6):466-477, 2019.

Yaniv Romano, Evan Patterson, and Emmanuel Candes. Conformalized quantile regression. Advances in neural information processing systems , 32, 2019.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACMSIGGRAPH 2022 Conference Proceedings , pages 1-10, 2022.

Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2022.

Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512 , 2022.

Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International conference on machine learning , pages 1218-1226. PMLR, 2015.

Swami Sankaranarayanan, Anastasios Nikolas Angelopoulos, Stephen Bates, Yaniv Romano, and Phillip Isola. Semantic uncertainty intervals for disentangled latent spaces. In Advances in Neural Information Processing Systems , 2022.

Nicki Skafte, Martin Jørgensen, and Søren Hauberg. Reliable training and estimation of variance networks. Advances in Neural Information Processing Systems , 32, 2019.

Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems , 28, 2015.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint arXiv:2111.08005 , 2021.

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 , 2023.

Lucas von Chamier, Romain F Laine, Johanna Jukkala, Christoph Spahn, Daniel Krentzel, Elias Nehme, Martina Lerche, Sara Hernández-Pérez, Pieta K Mattila, Eleni Karinou, et al. Democratising deep learning for microscopy with zerocostdl4mic. Nature communications , 12 (1):2276, 2021.

Guotai Wang, Wenqi Li, Michael Aertsen, Jan Deprest, Sébastien Ourselin, and Tom Vercauteren. Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks. Neurocomputing , 338:34-45, 2019.

Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. The Eleventh International Conference on Learning Representations , 2023.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes to model the full covariance matrix of aleatoric uncertainty for per-pixel image restoration tasks using a low-rank approximation. The low-rank approximation is estimated by a separate network trained end-to-end with the original predictor. The PCs, which are usually computed sequentially, are constructed in parallel using stop_gradiet operations.\n\nSTRENGTHS:\nThe approach is fast and, as far as I know, architecture agnostic.\n\nWEAKNESSES:\n- The approach requires training a separate model. \n\n- It is unclear how good parallel end-to-end training of the PCs is compared to sequential training. I could imagine that this type of approach yields pretty unstable training dynamics.\n\n- The covariance measured by this approach is over the entire image. I could image that for pixel-wise tasks there exists at least some location invariance.\n\n- Some important related work on low-rank approximations for pixel-wise aleatoric uncertainty [1] and efficient uncertainty estimation is missing (see works discussed in [2])\n\n- Minor: in eq (2) w_i is normalized to unit length whereas e_i is not. Is this intended?\n\n[1] Monteiro, M., Le Folgoc, L., Coelho de Castro, D., Pawlowski, N., Marques, B., Kamnitsas, K., van der Wilk, M. and Glocker, B., 2020. Stochastic segmentation networks: Modelling spatially correlated aleatoric uncertainty. Advances in Neural Information Processing Systems, 33, pp.12756-12767.\n\n[2] Postels, J., Segu, M., Sun, T., Sieber, L.D., Van Gool, L., Yu, F. and Tombari, F., 2022. On the practicality of deterministic epistemic uncertainty. In Proceedings of the 39th International Conference on Machine Learning (Vol. 162, pp. 17870-17909). PMLR.\n\nQUESTIONS:\n- What speaks against directly parameterizing the conditional output distribution using a multivariate gaussian with low-rank approximated covariance? As far as I know, you would also never have to explicitly evaluate the cov.\n\n- Could the authors comment on the training dynamics of the approach? How does the result compare to sequential training?\n\n- Could the authors comment on location invariance (see weaknesses)?\n\n- In l 160 mention that the error prediction model requires the right biases to work well. Can the authors elaborate on this?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFor uncertainty quantification in image recovery problems, the authors propose a way to train a neural network to produce estimates of the principle components of the posterior covariance matrix. Their approach starts with a neural network that produces the posterior mean and trains a new network (or a new prediction head on the posterior-mean network) to learn the principle components and their corresponding eigenvalues. They demonstrate their method on denoising, inpainting, super-resolution, and biological image-to-image translation problems.  They show that their method produces principle component estimates of a similar quality to those of diffusion models but thousands of times faster.\n\nSTRENGTHS:\n1. For image recovery, the idea of visualizing posterior uncertainty using principle components is a good one and, to my knowledge, has not been adequately explored in the literature.  The most common way to present image uncertainty information is to plot a pixel-wise variance map, but this does not show pixel dependencies, which are critical to understanding the structure of the uncertainty.  And trying to visually interpret uncertainty structure from dozens of posterior image samples is tedious and heuristic.\n\n2. The proposed method is fast at inference, since it uses only a single pass through a neural network.  This stands in contrast to Langevin/score/diffusion methods, which require generating many posterior samples, each of which can require thousands of passes through a neural network.\n\n3. The proposed method is widely applicable because it can be built on top of an existing conditional-mean (i.e., MMSE) estimation network, which are widely available.  This stands in contrast to, say, conditional normalizing flows, whose architectural constraints make them difficult to apply/tune on new applications.\n\n4. The numerical results are impressive.  The authors have tested their method on a wide range of applications, some of which involve consider large images (CelebA-HQ).  Also, their experiments suggest that their method gives similar RMSE and residual error magnitude to recent diffusion models. \n\n5. The paper is very clearly written.\n\nWEAKNESSES:\n1. The authors focus on a comparison to diffusion methods, which are very slow.  But modern conditional GANs and conditional normalizing flows (CNFs) can quickly generate posterior samples with performance that meets or exceeds that of recent diffusion models.  For example, a 2022 CVPR super-resolution contest (https://arxiv.org/abs/2205.05675) showed CNFs dominating other methods.  Compared to modern CNFs (e.g., https://arxiv.org/abs/2006.14200) or CGANs (e.g., https://arxiv.org/abs/2210.13389), it’s not clear that the proposed method has any speed or performance advantages.\n\n2. The proposed method is a one-trick pony, in that it generates only posterior principle components, and not high-quality image recoveries like posterior sampling methods do (e.g., diffusion, CGAN, CNF, etc.).  For example, the generated \"x+w\" face images are very blurry.  This weakness is acknowledged by the authors in Section 5.\n\n3. The proposed method seems practical for recovering only a few principle components (e.g., tens at most), whereas in some applications the uncertainty is not well described by only a few principle components.  In other words, sometimes the eigenvalues of the posterior covariance decays very slowly.  The authors acknowledge this issue in their face experiment.  The reviewer has observed, in their own work, that slow eigenvalue decay is the rule rather than the exception.\n\nQUESTIONS:\nThe paper was very clear and so I have no questions.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes using a deep neural network to predict the principle components (PCs), and the associated uncertainties, of the output directly, instead of just the most likely output. This is done by proposing a PC loss that ensures PCs are unit vectors and are orthogonal. The experimental results indicate that the proposed method performs comparably or better than other methods on super-resolution and inpainting tasks, while using much less compute when predicting for new examples.\n\nSTRENGTHS:\nThe paper proposes, what is to my knowledge, a novel PC-based method that leads to a more computationally efficient method for super-resolution and inpainting tasks that performs comparable to, or better than, competing techniques. The proposed technique also estimates the data (i.e. aleatoric) uncertainty by learning the variances associated with each PC. Knowing the uncertainty of the predictions could be helpful in understanding the outputs of the proposed method, which is necessary for many scientific applications.\n\nWEAKNESSES:\nThe main weakness of the paper is evaluations. The proposed method was only compared to two other methods on one dataset (Celeb-A-HQ). This comparison was not done on MNIST or the biological image-to-image translation dataset. The predicted PCs were not compared to the ground truth PCS. Also, the quality of the uncertainty quantification, a key claim of the paper, was not evaluated well.\n\nQUESTIONS:\n- Could the predicted principle components and variances be compared to the ground truths for a dataset such as MNIST? It was done for an extremely simple toy dataset in the supplementary material.\n\n- Is there another way to verify the improvement in uncertainty quantification that the proposed method brings?\n\n- Could the proposed method be compared with other image-to-image translation methods of the biological imaging dataset?\n\nLIMITATIONS:\nThe limitation of no guaranteed generalization performance is given. Is there a potential limitation related to choosing the number of PCs to predict? Societal impact was not addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: SUMMARY:
This work proposes to do image recovery inference using posterior principal components. The posterior principal components are often regarded as a function of observed image, but estimating it from a single image is impossible. Therefore, this work proposes to learn the principal components directly by training a neural net on the triplets of data (observed image, ground truth and posterior mean). It learns all the principal components together using shared weights, and at the same time managed to preserve orthogonality. The experimental results show its usefulness in various tasks.

STRENGTHS:
1. The paper is very well written and organized
2. The task of the paper is very important and useful
3. Designing and training a neural net that outputs all principal components at once and maintains othorgonality is nontrivial, and the paper solved the problem quite nicely.
4. The experimetal results are encouraging, as the proposed method outperforms the state-of-the-art posterior samplers

Overall, I think it is a solid paper, and its idea is simple and effective

WEAKNESSES:
There is no quality control on the predicted principal component. It is not necessarily a weakness, but it seems to be a common issue for most of deep learning based inference method.

QUESTIONS:
1. Given different observations, the posterior distribution could be quite different. In this case, do you still select the same number of principal components for each observation?
2. If you train two NNs, where the first NN output 5 components, and the second NN outputs 10 components. Are the top 5 components of the two NNs match each other?

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09