Towards Distribution-Shift Uncertainty Estimation for Inverse Problems with Generative Priors

TL;DR

The paper tackles the problem of distribution-shift-induced hallucinations when using generative priors for severely undersampled inverse problems. It introduces an instance-level, calibration-free uncertainty indicator that relies on reconstruction stability across random measurement subsets to detect out-of-distribution targets, demonstrated with a learned proximal network trained on digit '0' evaluated on all MNIST digits in sparse-view CT. The key contributions are the proposed OOD indicator, its efficient computation without external calibration data, and empirical validation showing higher reconstruction variance for out-of-distribution images correlating with larger errors. This approach enables a practical deployment strategy where aggressive measurement reduction can be used for in-distribution cases while automatically warning practitioners when priors encounter distribution shift, guiding additional measurements or method switching.

Abstract

Generative models have shown strong potential as data-driven priors for solving inverse problems such as reconstructing medical images from undersampled measurements. While these priors improve reconstruction quality with fewer measurements, they risk hallucinating features when test images lie outside the training distribution. Existing uncertainty quantification methods in this setting (i) require an in-distribution calibration dataset, which may not be available, (ii) provide heuristic rather than statistical estimates, or (iii) quantify uncertainty from model capacity or limited measurements rather than distribution shift. We propose an instance-level, calibration-free uncertainty indicator that is sensitive to distribution shift, requires no knowledge of the training distribution, and incurs no retraining cost. Our key hypothesis is that reconstructions of in-distribution images remain stable under random measurement variations, while reconstructions of out-of-distribution (OOD) images exhibit greater instability. We use this stability as a proxy for detecting distribution shift. Our proposed OOD indicator is efficiently computable for any computational imaging inverse problem; we demonstrate it on tomographic reconstruction of MNIST digits, where a learned proximal network trained only on digit "0" is evaluated on all ten digits. Reconstructions of OOD digits show higher variability and correspondingly higher reconstruction error, validating this indicator. These results suggest a deployment strategy that pairs generative priors with lightweight guardrails, enabling aggressive measurement reduction for in-distribution cases while automatically warning when priors are applied out of distribution.

Figures (4)

• Figure 1: Our generative prior is trained on MNIST "0" and used for CT imaging on all digits, to test distribution shift detection.
• Figure 2: Per-digit mean PSNR and SSIM (averaged over 10 images with 10 random seeds per image) with error bars indicating the min–-max spread. Solid curves correspond to LPN reconstructions and dashed curves to the FBP baseline; LPN consistently outperforms FBP. The performance gap is largest for the in-distribution digit "0", whose PSNR/SSIM is higher than those of OOD digits when reconstruction leverages the LPN. This is especially pronounced in the 11-view experiment that is most undersampled and thus relies most heavily on the learned prior.
• Figure 3: Visualizing distribution shift detection on MNIST: (a) mean reconstructions, (b) pixel-wise standard deviation.
• Figure 4: Average pixel-wise standard deviation is lower for the in-distribution digit "0" than the OOD digits, confirming our hypothesis that reconstruction instability across random measurements can detect distribution shift.