Conformalized Generative Bayesian Imaging: An Uncertainty Quantification Framework for Computational Imaging

TL;DR

This work presents a scalable framework to jointly quantify aleatoric and epistemic uncertainties in computational imaging by augmenting existing generative-model posterior samplers with Bayesian neural networks that include latent variables and deep ensembling, followed by conformal calibration to guarantee marginal coverage. The method formulates a conditional distribution with a latent-variable-enhanced generator, uses an ensemble of parameter realizations to approximate a predictive distribution, and decomposes total uncertainty into epistemic and aleatoric components. It connects GM posterior sampling and Bayesian neural network approaches, enabling robust uncertainty characterization beyond what either paradigm offers alone. Empirical results across CT, MRI, and image inpainting demonstrate that the framework captures true uncertainty patterns, improves predictive uncertainty quality, and can be conformally calibrated to achieve frequentist coverage guarantees, with practical considerations on computation and ensemble design.

Abstract

Uncertainty quantification plays an important role in achieving trustworthy and reliable learning-based computational imaging. Recent advances in generative modeling and Bayesian neural networks have enabled the development of uncertainty-aware image reconstruction methods. Current generative model-based methods seek to quantify the inherent (aleatoric) uncertainty on the underlying image for given measurements by learning to sample from the posterior distribution of the underlying image. On the other hand, Bayesian neural network-based approaches aim to quantify the model (epistemic) uncertainty on the parameters of a deep neural network-based reconstruction method by approximating the posterior distribution of those parameters. Unfortunately, an ongoing need for an inversion method that can jointly quantify complex aleatoric uncertainty and epistemic uncertainty patterns still persists. In this paper, we present a scalable framework that can quantify both aleatoric and epistemic uncertainties. The proposed framework accepts an existing generative model-based posterior sampling method as an input and introduces an epistemic uncertainty quantification capability through Bayesian neural networks with latent variables and deep ensembling. Furthermore, by leveraging the conformal prediction methodology, the proposed framework can be easily calibrated to ensure rigorous uncertainty quantification. We evaluated the proposed framework on magnetic resonance imaging, computed tomography, and image inpainting problems and showed that the epistemic and aleatoric uncertainty estimates produced by the proposed framework display the characteristic features of true epistemic and aleatoric uncertainties. Furthermore, our results demonstrated that the use of conformal prediction on top of the proposed framework enables marginal coverage guarantees consistent with frequentist principles.

Figures (7)

• Figure 1: Variation in the epistemic uncertainty estimates offered by the proposed framework for a randomly chosen CT test measurement vector as the number of examples in the training dataset is changed. The fractions of the training data used (0.125, 0.25, 0.5, and 1.0) are indicated above each epistemic uncertainty map. For this example, the generative model-based posterior sampling method employed by the proposed framework is deep posterior sampling Adler2019CWGAN (see \ref{['ssec:experiments-setup']} for further details).
• Figure 2: Average epistemic uncertainty as a function of training dataset size for the CT problem. Average values are calculated across all pixels in the test dataset. The generative model-based posterior sampling method used within the proposed framework is deep posterior sampling Adler2019CWGAN (see \ref{['ssec:experiments-setup']} for details).
• Figure 3: Effect of the test measurement vectors that are not well-represented by the training dataset on the epistemic uncertainty estimates provided by the proposed framework. The first row shows the output of filtered backprojection (FBP), the output of zero-filling (ZF), the predictive mean provided by the proposed framework, and the epistemic uncertainty estimate offered by the proposed framework for both the CT and MRI problems, where there are no abnormalities present in the test measurement vectors. The second row shows the results of the cases where abnormalities are introduced into the same test measurement vectors (please refer to Section \ref{['ssec:experiments-abnormality-epistemic']} for the details of these abnormalities). Note: For the MRI problem, epistemic uncertainty maps are computed before the final data consistency layer of UQVAE Edupuganti2021VAEMRI.
• Figure 4: Negative predictive log-likelihood (NPLL) results for deep posterior sampling Adler2019CWGAN (DPS), the variational autoencoder-based posterior sampling method Edupuganti2021VAEMRI (UQVAE), and the proposed deep ensembling Lakshminarayanan2017ensembling-based framework (DE). The results are presented for the CT (top) and MRI (bottom) problems.
• Figure 5: SSIM results for filtered backprojection (FBP), FBPConvNet Jin2017FBPConvNet (FBPCN), deep posterior sampling Adler2019CWGAN (DPS), zero-filling (ZF), BPConvNet Jin2017FBPConvNet (BPCN), the variational autoencoder-based posterior sampling method Edupuganti2021VAEMRI (UQVAE), and the proposed deep ensembling Lakshminarayanan2017ensembling-based framework (DE). Results shown for the CT (top) and MRI (bottom) problems.