Uncertainty Quantification via Neural Posterior Principal Components

TL;DR

This paper introduces a neural network method to efficiently predict the principal components of the posterior distribution for image restoration tasks, enabling fast and reliable uncertainty quantification in safety-critical applications.

Contribution

The authors propose a neural approach to predict posterior principal components in a single forward pass, improving speed and practicality over traditional sampling methods for uncertainty estimation.

Findings

Achieves uncertainty quantification comparable to posterior sampling methods.

Provides instance-adaptive uncertainty directions in various imaging tasks.

Operates orders of magnitude faster than existing sampling-based approaches.

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 at https://eliasnehme.github.io/NPPC/