Uncertainty Quantification via Neural Posterior Principal Components

Summary

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

Strengths

The approach is fast and, as far as I know, architecture agnostic.

Weaknesses

- The approach requires training a separate model. - 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. - 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. - 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]) - Minor: in eq (2) w_i is normalized to unit length whereas e_i is not. Is this intended? [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. [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.

Questions

- 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. - Could the authors comment on the training dynamics of the approach? How does the result compare to sequential training? - Could the authors comment on location invariance (see weaknesses)? - In l 160 mention that the error prediction model requires the right biases to work well. Can the authors elaborate on this?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Summary

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

Strengths

1. 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. 2. 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. 3. 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. 4. 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. 5. The paper is very clearly written.

Weaknesses

1. 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. 2. 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. 3. 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.

Questions

The paper was very clear and so I have no questions.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Yes

Summary

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

Strengths

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

Weaknesses

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

Questions

- 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. - Is there another way to verify the improvement in uncertainty quantification that the proposed method brings? - Could the proposed method be compared with other image-to-image translation methods of the biological imaging dataset?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

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

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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

yes