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. However, 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. Examples are available at https://eliasnehme.github.io/NPPC/
翻译:不确定性量化对于在安全关键领域(如自动驾驶和生物成像)中部署图像恢复模型至关重要。迄今为止,不确定性可视化方法主要集中于逐像素估计。然而,逐像素方差的热力图通常缺乏实际用途,因为它未能捕捉像素间的强相关性。更自然的不确定性度量对应于后验分布主成分上的方差。理论上,通过对输入图像的条件生成模型生成的样本进行主成分分析,可以计算主成分。然而,这需要在测试时生成大量样本,而使用当前最先进的扩散模型时,这一过程极为缓慢。在本工作中,我们提出了一种方法,能够通过神经网络单次前向传播,为任意输入图像预测后验分布的主成分。我们的方法既可以包装在预训练的均方误差最小化模型周围,也可以从头开始训练,同时输出预测图像和后验主成分。我们在多种成像逆问题(包括去噪、修复、超分辨率和生物图像到图像的转换)上展示了该方法。我们的方法能可靠地传递实例自适应不确定性方向,在实现与后验采样器相当的不确定性量化效果的同时,速度快几个数量级。示例请参见 https://eliasnehme.github.io/NPPC/