Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unsupervised Detection of Distribution Shift in Inverse Problems using Diffusion Models

Shirin Shoushtari, Edward P. Chandler, Yuanhao Wang, M. Salman Asif, Ulugbek S. Kamilov

TL;DR

The paper tackles distribution shifts in diffusion-prior imaging inverse problems where clean test data are unavailable. It introduces a measurement-domain KL divergence estimator that uses corrupted measurements and pretrained diffusion-model scores, proving under mild assumptions that it equals the image-domain $D_{KL}(p||q)$ and enabling shift quantification without ground-truth images. Empirically, the metric tracks the image-domain KL in inpainting and MRI tasks and the authors show that adapting OOD diffusion models using only corrupted measurements reduces the estimated shift and improves reconstruction quality. This unsupervised approach provides a practical tool for evaluating and mitigating distribution shifts in high-stakes imaging applications, with potential to guide lightweight adaptation strategies in real-world inverse problems.

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.

Unsupervised Detection of Distribution Shift in Inverse Problems using Diffusion Models

TL;DR

The paper tackles distribution shifts in diffusion-prior imaging inverse problems where clean test data are unavailable. It introduces a measurement-domain KL divergence estimator that uses corrupted measurements and pretrained diffusion-model scores, proving under mild assumptions that it equals the image-domain and enabling shift quantification without ground-truth images. Empirically, the metric tracks the image-domain KL in inpainting and MRI tasks and the authors show that adapting OOD diffusion models using only corrupted measurements reduces the estimated shift and improves reconstruction quality. This unsupervised approach provides a practical tool for evaluating and mitigating distribution shifts in high-stakes imaging applications, with potential to guide lightweight adaptation strategies in real-world inverse problems.

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.
Paper Structure (23 sections, 4 theorems, 41 equations, 9 figures, 4 tables)

This paper contains 23 sections, 4 theorems, 41 equations, 9 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Denoising Diffusion Probabilistic Models
  4. Measuring Distribution Shifts with Clean Images using Score Functions
  5. Distribution Shift in Measurement Domain
  6. Problem Formulation
  7. Theoretical Results
  8. Experiments
  9. Computing Distribution Shift
  10. Adaptation Effect on Distribution Shift
  11. Ablation studies
  12. Conclusion
  13. Proof of KL Divergence Metric on Image Domain
  14. Proof of Theorem \ref{['thm:thm1']}
  15. KL Divergence for Noisy Measurements
  16. ...and 8 more sections

Key Result

Theorem 1

Let ${\overline{\bm{y}}}_\sigma = {\bm{P}} {\overline{\bm{x}}} + {\overline{\bm{n}}}$ denote the noisy projected measurements at noise level $\sigma$ according to eq:ybarsigma. Then, the KL divergence between the InD density ${p({\bm{x}})}$ and the OOD density ${q({\bm{x}})}$ can be expressed as where ${\bm{W}} = \mathbb{E}[{\bm{P}}]^{-3/2}$ is a diagonal scaling matrix, ${\bm{V}}$ is the right s

Figures (9)

  • Figure 1: Comparison of the distribution shift (dashed lines), computed using clean images, and our proposed measurement-domain KL metric (solid lines) between an InD model trained on FFHQ and OOD models trained on MetFaces, AFHQ, and Microscopy. Results are shown under inpainting masks with $p \in \{0.2, 0.5, 0.8\}$. The vertical axis shows $D_{\mathrm{KL}}$, evaluated as the integrand in \ref{['eq:mainmetric']} and \ref{['eq:KLforimage']} up to diffusion noise level $\sigma$. Right: Samples from InD and OOD datasets. Note how the proposed metric accurately tracks the KL divergence, even under high-levels of corruption (smaller values of $p$).
  • Figure 2: KL divergence plotted against the noise level $\sigma$ for InD and OOD Gaussian mixture models (GMMs). KL divergence computed in the image domain (blue) and measurement domain (red) under inpainting corruption with probability $p$, using $N$ InD data example. The measurement-domain KL divergence closely tracks its image-domain counterpart, and the approximation improves with increasing $N$ and $p$.
  • Figure 3: Comparison of the distribution shift (dashed lines), computed using clean images, and our proposed measurement-domain KL metric (solid lines) between an InD model trained on Brain slices and OOD models trained on Knee and Prostate slices from fastMRI dataset with acceleration rate $4$. The vertical axis shows $D_{\mathrm{KL}}$, evaluated as the integrand in \ref{['eq:mainmetric']} and \ref{['eq:KLforimage']} up to diffusion noise level $\sigma$. The proposed metric accurately tracks the KL divergence.
  • Figure 4: $D_{\mathrm{KL}}$ between FFHQ and AFHQ, as well as adapted models using 64 and 128 projected measurements, measured in the image domain (dashed) and the measurement domain (solid) for inpainting with p=0.8. Notably, adapting the network using only projected measurements significantly reduces the distributional gap.
  • Figure 5: Visual comparison of inpainting results (DPS chung2023diffusion) on an FFHQ image with mask rate $p = 0.8$ and measurement noise level $\sigma = 0.01$. The top row shows full reconstructions, while the bottom row displays residual maps (left) and zoomed-in regions (right). Note the performance gap between the InD and OOD models, and the improvement achieved by adapting the OOD models using only corrupted measurements.
  • ...and 4 more figures

Theorems & Definitions (9)

  • Theorem 1
  • proof
  • Theorem 1
  • proof
  • Remark 1
  • Corollary 1
  • proof
  • Theorem 2
  • proof