🤖 machine learning

KLIP: localized distribution shift detection via KL-divergence with diffusion priors in Inverse Problems

The paper proposes KLIP, a novel out-of-distribution detection metric for inverse problems that leverages the Kullback-Leibler divergence between diffusion priors and posterior distributions to identify and localize subtle distribution shifts without requiring calibration data or knowledge of the shifted distribution.

Original authors: Alireza Kheirandish, Jihoon Hong, Sara Fridovich-Keil

Published 2026-06-01

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Alireza Kheirandish, Jihoon Hong, Sara Fridovich-Keil

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Finding the "Odd One Out" in a Blurry Photo

Imagine you are trying to solve a puzzle, but the pieces are missing, and you only have a few blurry clues (like a few scattered puzzle pieces or a smudged photo). In the world of medical imaging, this is called an inverse problem. Doctors often have to reconstruct a clear picture of a patient's insides (like a liver) from limited, noisy data (like a few X-ray angles).

Usually, AI models are trained to fill in the missing pieces based on what they've seen before. If the model was trained on healthy livers, it will try to "guess" a healthy liver when it sees the blurry data.

The Challenge: What if the patient actually has a tumor?

The AI might still try to draw a healthy liver because that's what it knows.
The tumor is a localized problem (it's only in one small spot), not a problem with the whole image.
Existing AI tools often fail here. They might say, "This whole picture looks weird," or they need to be shown examples of tumors first to know what to look for. But in the real world, you can't show the AI every possible rare disease beforehand.

The Solution: KLIP (The "Suspicion Meter")

The authors propose a new tool called KLIP. Think of it as a suspicion meter that doesn't need to have seen a crime before to know one is happening.

It works by comparing two things:

The "Expectation" (The Prior): What the AI thinks the image should look like based on its training (e.g., "I expect a healthy liver").
The "Reality" (The Posterior): What the image actually looks like when forced to match the blurry clues (e.g., "The data says there's a weird lump here").

If the "Reality" is very different from the "Expectation," the meter goes off. The bigger the difference, the more likely it is that something unusual (Out-of-Distribution or OOD) is happening.

The Secret Sauce: Zooming In and Time-Traveling

The paper's main breakthrough is that KLIP doesn't just look at the whole picture at once. It uses two clever tricks to find small, hidden problems:

1. The "Grid Search" (Spatial Restriction)

Imagine you are looking at a large map to find a tiny lost coin. If you look at the whole map, the coin is invisible.

Old way: Look at the whole map and say, "It looks mostly normal."
KLIP way: It divides the map into a grid of small squares. It checks each square individually. If one tiny square has a huge difference between "what we expect" and "what the data says," KLIP flags that specific square.
Result: It can pinpoint a small tumor in a liver or a scar on a face, even if the rest of the image is perfect.

2. The "Time-Travel" Trick (Timestep Restriction)

Diffusion models (the AI used here) build images like a sculptor. They start with a rough, blurry blob and slowly add fine details over time.

Early stages: The AI is just figuring out the big shapes (the outline of a face or the general shape of a liver).
Later stages: The AI adds the fine details (the eyes, the texture of the skin, the specific shape of a tumor).

KLIP realizes that different types of "weirdness" show up at different times.

A big, global change might be obvious early on.
A tiny, specific detail (like a scar or a small tumor) might only become obvious when the AI is adding the fine details near the end of the process.
KLIP way: It checks the "suspicion meter" at specific moments in time, not just the whole process. This helps it catch subtle details that would otherwise be missed.

What They Tested It On

The authors tested this "suspicion meter" on two main scenarios:

Medical Scans (CT Scans):
- Training: The AI only saw healthy livers.
- The Test: They gave it blurry scans of livers that had synthetic tumors added to them.
- The Result: KLIP didn't just say "This is a weird scan." It drew a heat map that highlighted the exact location of the tumor, even though the AI had never seen a tumor before.
Faces (De-blurring):
- Training: The AI was trained on photos of celebrities (healthy faces).
- The Test: They gave it blurry photos of faces with added scars or photos of sci-fi characters with prosthetics (like a metal eye implant).
- The Result: KLIP successfully pointed to the scar on a human face or the metal implant on a character's face, distinguishing them from normal features.

Why This Matters (According to the Paper)

No "Training on Bad Data" Needed: Most AI needs to be shown examples of "bad" things (like tumors) to learn to spot them. KLIP doesn't need this. It only needs to know what "good" looks like.
Works on Blurry Data: It works directly on the indirect measurements (the blurry clues), not just the final clear image.
Finds the Needle in the Haystack: It is great at finding small, localized problems without getting confused by the rest of the image.

The Limitations Mentioned

The authors are honest about a few things:

Tuning: You have to adjust the "size of the grid squares" and "when to check the time" carefully. If you pick the wrong settings, it might not work as well.
Simulation: They tested this using computer simulations (creating fake blurry images from perfect ones). They haven't tested it on real-world, messy hospital data yet, though they believe it will work there too.

In summary: KLIP is a smart way to ask an AI, "Does this blurry clue match what you expect?" If the answer is "No," it zooms in to find exactly where the mismatch is, without needing to have seen that specific mismatch before.

Technical Summary: KLIP – Localized Distribution Shift Detection via KL-Divergence with Diffusion Priors

Problem Statement
The paper addresses the challenge of detecting Out-of-Distribution (OOD) images and, more critically, localizing spatially localized OOD features within images during the solution of inverse problems (e.g., image reconstruction from indirect measurements). Existing OOD detection methods often rely on calibration datasets containing OOD examples, which are unavailable for the "long tail" of rare abnormalities in open-world scenarios (e.g., specific medical anomalies). Furthermore, many current approaches operate on full images rather than the indirect measurements available in computational imaging, or they fail to detect subtle, localized shifts (such as small tumors or tears) that are essential for diagnostic value.

Methodology
The authors propose KLIP, a calibration-free OOD detection metric based on the Kullback-Leibler (KL) divergence between the diffusion prior distribution $p(x)$ and the posterior distribution $p(x|y)$ given measurements $y$ .

Core Intuition: The paper posits that if an image contains OOD features, the measurement $y$ will "pull" the posterior distribution significantly away from the prior distribution. Consequently, the KL divergence $D_{KL}(p(x|y) \parallel p(x))$ should be larger for OOD images and specifically larger for OOD regions within an image.
Estimation via Posterior Sampling: Since the exact KL divergence is intractable, the authors approximate it using the score functions generated during diffusion posterior sampling. By decomposing the posterior score into the prior score and the likelihood score ( $\nabla_x \log p_t(y|x)$ ), the KL divergence is estimated by integrating the squared norm of the likelihood score over the sampling trajectory (Equation 11).
Block and Timestep Restriction: To detect localized shifts rather than global ones, the authors introduce two key restrictions:
- Spatial Block Restriction: The image is divided into blocks ( $B_i$ ). The KL divergence is computed locally for each block to prevent the signal of a small OOD feature from being diluted by the rest of the in-distribution image.
- Timestep Restriction: The integration is restricted to specific time windows $[t_0, t_1]$ of the sampling trajectory. Drawing on the observation that diffusion models generate features at different scales at different timesteps, this allows the metric to focus on the specific stages where OOD features are most distinct.

The final metric, KLIP, is defined as the restricted KL divergence over a specific spatial block and time window (Equation 12).

Key Contributions

Calibration-Free Detection: KLIP requires no OOD calibration data or knowledge of the shifted distribution; it relies solely on a diffusion model trained on In-Distribution (ID) data.
Applicability to Inverse Problems: The method operates directly on indirect measurements available at inference time, making it suitable for computational imaging tasks like CT reconstruction and deblurring.
Localization Capability: Unlike global OOD detectors, KLIP can identify the precise spatial location of subtle distribution shifts (e.g., small tumors or artifacts) within an image.
Generalization: The metric is shown to generalize across different diffusion model architectures (predictor-corrector and patch-based), different datasets (abdominal CTs and human faces), and different inverse problems (sparse-view CT and Gaussian deblurring).

Experimental Results
The authors evaluated KLIP on two primary tasks:

Sparse-View CT Reconstruction: Using healthy abdominal CT scans as ID data, the model was tested on images with synthetic star artifacts and realistic synthetic liver tumors. KLIP achieved high Area Under the Curve (AUC) scores for both dataset-level and image-level detection, outperforming baselines like DiffPath, Negative Log Likelihood (NLL), and non-diffusion anomaly detectors (CutPaste, SimpleNet) in the context of indirect measurements.
Gaussian Deblurring of Faces: Using CelebA as ID data, the method detected localized OOD features such as synthetic scars and distinctive facial prosthetics on characters from movies (e.g., Star Trek, Harry Potter).

Qualitative results (Figures 4 and 5) demonstrate that KLIP heatmaps successfully localize the OOD features (stars, tumors, scars, prosthetics) with high precision. Ablation studies confirmed that both block and timestep restrictions are necessary for optimal performance; using the full image or full timestep range significantly reduces detection capability.

Significance and Claims
The paper claims that KLIP provides a robust, calibration-free solution for detecting and localizing subtle distribution shifts in inverse problems, a setting where traditional methods often fail due to the lack of OOD training data. By leveraging the temporal dynamics of diffusion sampling and spatial partitioning, KLIP offers a "rich window" into the sampling process under distribution shift. The authors emphasize that this approach is particularly valuable for medical imaging and other safety-critical applications where identifying rare, localized anomalies is crucial, and where acquiring representative OOD data for calibration is impossible.

The paper acknowledges limitations, noting that KLIP is sensitive to hyperparameters (block size and time window) which depend on the scale and shape of the OOD features, and that current evaluations rely on "inverse crime" setups (using the same forward model for simulation and reconstruction). However, the results suggest strong potential for generalization across diverse imaging tasks.