Measuring Distribution Shifts in Inverse Problems without clean data

ICLR 2026 Conference Submission21206 Authors

Keywords: Imaging inverse problems, distribution shift, model selection.

TL;DR: We propose a metric to identify the best diffusion model prior for solving imaging inverse problems in test-time,provide adaptation based on the proposed metric, and measure KL divergence.

Abstract: Diffusion models are widely used as priors in imaging inverse problems. However, their performance often degrades under distribution shifts between the training and test-time images. Existing methods for identifying and quantifying distribution shifts typically require access to clean test images, which are never available at test time when solving inverse problems. We propose a flexible framework for measuring distribution shift using *only* corrupted test measurements and candidate diffusion model scores. Our framework enables three complementary capabilities. First, in the general case with only a pool of diffusion models, it supports a principled model selection by identifying the model whose prior best matches the test data. Second, when an in-distribution model is available, our metric provides a theoretically guaranteed estimator of KL divergence that closely matches the image-domain KL. Third, the metric serves as a tool for adaptation guidance: aligning score functions with corrupted measurements reduces the estimated shift and improves reconstruction quality. Experiments on inpainting and MRI confirm that our method (i) achieves robust model selection, (ii) reliable estimates KL divergence in the presence of an in-distribution model, and (iii) enables effective adaptation to mitigate distribution shift.

Primary Area: generative models

Submission Number: 21206