Unsupervised Detection of Distribution Shift in Inverse Problems using Diffusion Models

TL;DR

This paper introduces an unsupervised method to detect distribution shifts in inverse problems using diffusion models, which improves reconstruction quality without requiring clean test images.

Contribution

It proposes a novel score-based metric that estimates distribution shifts solely from corrupted measurements, aligning with the KL divergence between distributions.

Findings

The metric accurately estimates KL divergence from corrupted data.

Aligning scores reduces distribution shift and improves reconstruction.

Method works across multiple inverse problem types.

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 almost never available while solving inverse problems (at test time). We propose a fully unsupervised metric for estimating distribution shifts using only indirect (corrupted) measurements and score functions from diffusion models trained on different datasets. We theoretically show that this metric estimates the KL divergence between the training and test image distributions. Empirically, we show that our score-based metric, using only corrupted measurements, closely approximates the KL divergence computed from clean images. Motivated by this result, we show that aligning the out-of-distribution score with the in-distribution score -- using only corrupted measurements -- reduces the KL divergence and leads to improved reconstruction quality across multiple inverse problems.