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/
翻译:不确定性量化对于图像恢复模型在自动驾驶和生物成像等安全关键领域的部署至关重要。迄今为止,不确定性可视化方法主要集中在逐像素估计上。然而,逐像素方差的热图通常实用性有限,因为它未能捕捉像素之间的强相关性。一种更自然的不确定性度量对应于后验分布主成分(PCs)上的方差。理论上,通过对输入图像的条件生成模型生成的样本进行主成分分析(PCA),即可计算这些主成分。但这种方法需要在测试时生成大量样本,目前最先进模型(如扩散模型)执行此操作的速度异常缓慢。本研究提出一种方法,可通过单次神经网络前向传播预测任意输入图像的后验主成分。该方法既可封装于预训练的均方误差(MSE)最小化模型,也可从头训练,同时输出预测图像和后验主成分。我们在图像处理的多类逆问题(包括去噪、修复、超分辨率和生物图像到图像的翻译)中展示了该方法。该方法能可靠地传递实例自适应不确定性方向,实现与后验采样器相当的不确定性量化,同时速度提升数个数量级。代码和示例见 https://eliasnehme.github.io/NPPC/