SlowGuess/NeurIPS_2023_OCR huggingface dataset:
then we would only recover them up to an orthogonal matrix. To see this, let  $\mathbf{W}_i$  denote the matrix that has  $\mathbf{w}_k(\mathbf{y}_i, \hat{\mathbf{x}}_i; \varphi)$  in its  $k^{\text{th}}$  column. Then the inner sum in (4) can be rewritten as  $\| \mathbf{W}_i^\top \mathbf{e}_i \|_2^2$ . Now, it is easy to see that neither the loss nor the constraints are affected by replacing each  $\mathbf{W}_i$  with  $\tilde{\mathbf{W}}_i = \mathbf{W}_i \mathbf{O}_i$ , for some orthogonal matrix  $\mathbf{O}_i$ . Indeed, this solution satisfies the orthogonality constraint  $\tilde{\mathbf{W}}_i^\top \tilde{\mathbf{W}}_i = \mathbf{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^{\mathrm{th}}$  PC given the preceding  $k - 1$  pre-trained PCs, the loss term  $\left|\pmb{w}_k(\pmb{y}_i,\hat{\pmb{x}}_i;\pmb {\varphi})^\top \pmb {e}_i\right|^2$  did not affect the learning of  $\pmb{w}_{1},\dots ,\pmb{w}_{k - 1}$ . However, when attempting to learn all directions jointly, the preceding PCs receive a gradient signal from the  $k^{\mathrm{th}}$  loss term as  $\pmb{w}_{k}$  is a function of  $\pmb{w}_{1},\dots ,\pmb{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^{\mathrm{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  $\{\sigma_k^2\}$  of the PCs, by using a loss that minimizes  $\sum_{i}(\sigma_{k}^{2} - |\pmb{w}_{k}^{\top}\pmb{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 unnomarized directions. To achieve this without altering the optimization objective for finding the  $k^{\mathrm{th}}$  PC, we again invoke the Stopgrad operator on the term  $|\pmb{w}_k^\top \pmb {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
$$
\mathcal {L} _ {\sigma} \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(\left\| \boldsymbol {d} _ {k} - \sum_ {\ell = 1} ^ {k - 1} \left(\boldsymbol {d} _ {k} ^ {\top} \boldsymbol {w} _ {\ell}\right) \boldsymbol {w} _ {\ell} \right\| _ {2} ^ {2} - \left| \boldsymbol {w} _ {k} ^ {\top} \boldsymbol {e} _ {i} \right| ^ {2}\right) ^ {2}. \tag {5}
$$
With this, our prediction for  $\sigma_k^2$  at test-time is simply  $\| \pmb{d}_k - \sum_{\ell = 1}^{k - 1}(\pmb{d}_k^\top \pmb{w}_\ell)\pmb{w}_\ell \| ^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  $\hat{\pmb{x}} = f(\pmb{y};\pmb{\theta})$ , obtained through MSE minimization on a dataset  $\mathcal{D} = \{(x_i,y_i)\}_{i=1}^{N_d}$  of matched pairs. Hence, our entire derivation revolved around a dataset of triplets  $\mathcal{D}' = \{(x_i,y_i,\hat{x}_i)\}_{i=1}^{N_d}$ . However, this is not strictly necessary. In particular, the posterior mean predictor can be learned jointly alongside the PCs with an
Figure 3: Denoising samples from a 2D Gaussian mixture. The left two panels depict the underlying (unknown) signal prior  $p(\pmb{x})$  (blue heatmap), exemplar matched samples from the joint distribution  $(\pmb{x}_i, \pmb{y}_i) \sim p(\pmb{x}, \pmb{y})$ , and the resulting training set  $\mathcal{D}$ . The right two panels show the analytical posterior  $p(\pmb{x}|\pmb{y})$  (red heatmap) for two different test points  $\pmb{y}_1$  and  $\pmb{y}_2$  (marked as red dots) on top of the signal prior. Our estimated conditional mean  $\hat{\mu}_{\mathbf{x}|\mathbf{y}}(\pmb{y})$  (black star) and PCs scaled by the estimated std (dashed black ellipse) coincide with the analytical posterior mean  $\mu_{\mathbf{x}|\mathbf{y}}(\pmb{y})$  (red star), and scaled posterior PCs derived from the analytical covariance (solid orange ellipse).
additional MSE loss
$$
\mathcal {L} _ {\boldsymbol {\mu}} (\mathcal {D}, \varphi) = \sum_ {(\boldsymbol {x} _ {i}, \boldsymbol {y} _ {i} \in \mathcal {D})} \| \boldsymbol {x} _ {i} - \hat {\boldsymbol {x}} _ {i} \| _ {2} ^ {2}. \tag {6}
$$
Note that we do not want the gradients of  $\mathcal{L}_w$  and  $\mathcal{L}_{\sigma}$  to affect the recovered mean. Therefore, we use a Stopgrad operator when inputting the error  $\boldsymbol{x}_i - \hat{\boldsymbol{x}}_i$  to those loss functions.
3.5 Putting it all together
To summarize, given a dataset  $\mathcal{D} = \{(x_i,y_i)\}_{i = 1}^{N_d}$  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),
$$
\mathcal {L} _ {\text {a l l}} = \mathcal {L} _ {\boldsymbol {\mu}} + \lambda_ {1} \mathcal {L} _ {\boldsymbol {w}} + \lambda_ {2} \mathcal {L} _ {\sigma}, \tag {7}
$$
where  $\lambda_{1}$  and  $\lambda_{2}$  are weights balancing the contributions of the different losses.
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  $\lambda_{1},\lambda_{2}$  are in App. A.
Toy examples Figure 3 demonstrates
NPPC on a 2D denoising task, where samples  $\pmb{x}_i$  from a two-component Gaussian mixture model are contaminated by additive white Gaussian noise to result in noisy measurements  $\pmb{y}_i$ . To predict  $\mathbf{x}$  from  $\mathbf{y}$ , we trained a 5-layer MLP with 256 hidden features using MSE minimization, and around the estimated conditional mean  $\hat{\pmb{x}} = \hat{\mu}_{\mathbf{x}|\mathbf{y}}(\pmb{y})$ , we trained NPPC (instantiated with a similar MLP) to output the two PCs
Table 1: Comparison of the Wasserstein 2-distance from the rank  $K$  Gaussian approximation of the GT posterior  $p(\boldsymbol{x}|\boldsymbol{y}) \approx \mathcal{N}(\boldsymbol{x};\boldsymbol{\mu}_{\boldsymbol{x}|\boldsymbol{y}},\boldsymbol{\Sigma}_{\boldsymbol{x}|\boldsymbol{y}})$  (i.e.,  $K = \mathrm{rank}(\Sigma_{\mathbf{x}|\mathbf{y}})$ ), on 5000 test samples (see text).
K=0	K=3	K=6	K=9	K=12
Baseline	0.4	180.6	272.5	317.5	325.7
NPPC	0.4	30.3	31.2	26.6	28.8
$\pmb{w}_{1}, \pmb{w}_{2}$ , and their variances  $\hat{\sigma}_{1}^{2}, \hat{\sigma}_{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(\pmb{x}|\pmb{y}) \approx \mathcal{N}(\pmb{\mu}_{\mathbf{x}|\mathbf{y}}, \pmb{\Sigma}_{\mathbf{x}|\mathbf{y}})$ , and an estimated Gaussian constructed by NPPC,  $\hat{p}(\pmb{x}|\pmb{y}) = \mathcal{N}(\hat{\mu}_{\mathbf{x}|\mathbf{y}}, \pmb{W}_{\star} \pmb{W}_{\star}^{\top})$ , where  $\pmb{W}_{\star}$  has the scaled estimated PC  $\hat{\sigma}_{k} \pmb{w}_{k}$  in its  $k$ th column. Compared to a baseline of a point mass at the estimated
(a) Denoising
(b) Inpainting
Figure 4: MNIST denoising and inpainting. (a) Here we show the application of NPPC to extreme image denoising. On the left is the noisy measurement  $\pmb{y}$ , the estimated conditional mean  $\hat{\pmb{x}}$ , and the (unknown) ground truth test image  $\pmb{x}$ . On the right, the top row shows the first  $K = 5$  predicted PCs  $\pmb{w}_k$ , and the bottom rows show a traversal of 3 standard deviations around  $\hat{\pmb{x}}$  for  $\pmb{w}_1$ ,  $\pmb{w}_4$ . At an extreme noise level of  $\sigma_{\varepsilon} = 1$ , the digit is either a "4" or a "9". (b) Here we show the result of NPPC on image inpainting from only the 8 bottom rows. The PCs reveal the digit is either a "7" or a "9".
conditional mean,  $\delta (\mathbf{x} - \hat{\mathbf{x}})$ , NPPC reduces the Wasserstein 2-distance by  $100\times$  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 rank-  $K$  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  $\sigma_{\varepsilon} = 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  $\hat{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 \times 256$  face images from the CelebA-HQ dataset. Here, we demonstrate NPPC on the task of inpainting, as well as on noisy  $8 \times$  super-resolution with a box downsampling filter and a noise-level of  $\sigma_{\varepsilon} = 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 \times$  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 \times 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)  $\| \pmb{x} - \hat{\pmb{x}} \|_2$ , and residual error magnitude  $\| e - W W^\top e \|_2$  using the first  $K = 5$  PCs. This is while being  $100 \times$  faster than MAT, and  $10^3 - 10^5 \times$  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  $\approx 90^\circ$ . 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
(a) Inpainting
(b) Super-resolution
Figure 5: CelebA-HQ inpainting and  $8 \times$  noisy super-resolution. (a) Application of NPPC to image inpainting with a mask around the mouth. NPPC received the masked image  $\mathbf{y}$  and the mean prediction  $\hat{\mathbf{x}}$ , and predicted the first 5 posterior PCs  $\mathbf{w}_1, \ldots, \mathbf{w}_5$ . On the bottom, we show a traversal of 2 standard deviations around the mean estimate  $\hat{\mathbf{x}}$  for  $\mathbf{w}_2$  and  $\mathbf{w}_3$ , capturing uncertainty regarding open/closed mouth, and background vs shadow next to the jawline. (b) Application of NPPC to  $8 \times$  noisy super-resolution with a noise standard deviation  $\sigma_{\varepsilon} = 0.05$ . Similarly, the PC traversals on the bottom capture uncertainty of eye size and cheek color/position.
Table 2: Quantitative comparison of  $\| \pmb{x} - \hat{\pmb{x}}\|_2 \downarrow / \| \pmb{e} - \pmb{W}\pmb{W}^\top \pmb{e}\|_2 \downarrow$  with posterior samplers on 100 test images from FFHQ. Mean prediction and PCs were computed using 100 samples per test image, and compute is reported in neural function evaluations (NFEs).
Super-resolution	Inpainting	NFEs↓
4× noiseless	4× noisy	8× noiseless	8× noisy	Eyes	Mouth
DDRM [19]	8.94/8.89	11.29/11.25	13.68/13.47	16.20/16.00	-/-	-/-	2·103
DDNM [54]	8.43/8.38	10.74/10.70	13.12/12.95	15.73/15.55	13.27/11.24	13.30/10.37	10·103
RePaint [26]	-/-	-/-	-/-	-/-	12.24/10.3	12.55/9.72	457·103
MAT [23]	-/-	-/-	-/-	-/-	14.12/12.94	13.08/11.74	100
NPPC (Ours)	8.4/8.24	10.41/10.35	13.06/12.87	15.29/15.11	13.55/11.43	12.23/10.27	1
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  $10\mathrm{min}$ ) using a spinning-disk microscope [52]. The dataset consisted of 1753 image pairs of resolution  $1024 \times 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  $\hat{x}$ , thereby clearly communicating reconstruction ambiguity to the user.
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
Figure 6: Biological image-to-image translation. NPPC applied to the task of translating one fluorescent label  $\mathbf{y}$  (actin staining) to another fluorescent label  $\mathbf{x}$  (nuclear staining) in migrating cells (left). On the right, we show the predicted PCs at the top, and traversals along the rows. Yellow arrows point to captured uncertain cells in  $\hat{\mathbf{x}}$  (either deforming or being completely erased).
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.
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.
This research was partially supported by the Israel Science Foundation (grant no. 2318/22), by the Ollendorff Minerva Center, ECE faculty, Technion, and by a gift from KLA. The authors gratefully acknowledge fruitful discussions with Daniel Freedman, Regev Cohen, and Rotem Mulayoff throughout this work. The authors also thank Matan Kleiner, Hila Manor, and Noa Cohen for their help with the figures.
[1] Alexander Amini, Wilko Schwarting, Ava Soleimany, and Daniela Rus. Deep evidential regression. Advances in Neural Information Processing Systems, 33:14927-14937, 2020.
[2] 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.
[3] 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.
[4] 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.
[5] Christopher M Bishop. Mixture density networks. 1994.
[6] 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.
[7] Eric M Christiansen, Samuel J Yang, D Michael Ando, Ashkan Javaherian, Gaia Skibinski, Scott Lipnick, Elliot Mount, Alison O'neil, Kevin Shah, Alicia K Lee, et al. In silico labeling: predicting fluorescent labels in unlabeled images. Cell, 173(3):792-803, 2018.
[8] 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.
[9] 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.
[10] Garoe Dorta, Sara Vicente, Lourdes Agapito, Neill DF Campbell, and Ivor Simpson. Training vaes under structured residuals. arXiv preprint arXiv:1804.01050, 2018.
[11] 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.
[12] 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.
[13] 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.
[14] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.