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

Resumen Técnico: KLIP – Detección de Desplazamiento de Distribución Localizado mediante la Divergencia KL con Priores de Difusión

Planteamiento del Problema
El artículo aborda el desafío de detectar imágenes fuera de la distribución (OOD, Out-of-Distribution) y, de manera más crítica, de localizar características OOD espacialmente localizadas dentro de las imágenes durante la resolución de problemas inversos (por ejemplo, la reconstrucción de imágenes a partir de mediciones indirectas). Los métodos actuales de detección de OOD a menudo dependen de conjuntos de datos de calibración que contienen ejemplos OOD, los cuales no están disponibles para la "cola larga" (long tail) de anomalías raras en escenarios de mundo abierto (por ejemplo, anomalías médicas específicas). Además, muchos enfoques actuales operan sobre imágenes completas en lugar de las mediciones indirectas disponibles en la imagen computacional, o bien no logran detectar desplazamientos sutiles y localizados (como pequeños tumores o desgarros) que son esenciales para el valor diagnóstico.

Metodología
Los autores proponen KLIP, una métrica de detección de OOD libre de calibración basada en la divergencia de Kullback-Leibler (KL) entre la distribución del prior de difusión $p(x)$ y la distribución posterior $p(x|y)$ dada las mediciones $y$.

1. Intuición Central: El artículo postula que si una imagen contiene características OOD, la medición $y$ "tirará" de la distribución posterior significativamente lejos de la distribución del prior. En consecuencia, la divergencia KL $D_{KL}(p(x|y) \parallel p(x))$ debería ser mayor para las imágenes OOD y, específicamente, mayor para las regiones OOD dentro de una imagen.
2. Estimación mediante Muestreo Posterior: Dado que la divergencia KL exacta es intratable, los autores la aproximan utilizando las funciones de puntuación (score functions) generadas durante el muestreo posterior de difusión. Al descomponer la puntuación posterior en la puntuación del prior y la puntuación de la verosimilitud ($\nabla_x \log p_t(y|x)$), la divergencia KL se estima integrando la norma al cuadrado de la puntuación de la verosimilitud a lo largo de la trayectoria de muestreo (Ecuación 11).
3. Restricción de Bloque y de Paso de Tiempo: Para detectar desplazamientos localizados en lugar de globales, los autores introducen dos restricciones clave:
   • Restricción de Bloque Espacial: La imagen se divide en bloques ($B_i$). La divergencia KL se calcula localmente para cada bloque para evitar que la señal de una pequeña característica OOD se diluya por el resto de la imagen que está dentro de la distribución (ID).
   • Restricción de Paso de Tiempo: La integración se restringe a ventanas de tiempo específicas $[t_0, t_1]$ de la trayectoria de muestreo. Basándose en la observación de que los modelos de difusión generan características a diferentes escalas en diferentes pasos de tiempo, esto permite que la métrica se concentre en las etapas específicas donde las características OOD son más distintivas.

La métrica final, KLIP, se define como la divergencia KL restringida sobre un bloque espacial y una ventana de tiempo específicos (Ecuación 12).

Contribuciones Clave

• Detección Libre de Calibración: KLIP no requiere datos de calibración OOD ni conocimiento de la distribución desplazada; se basa únicamente en un modelo de difusión entrenado con datos dentro de la distribución (ID).
• Aplicabilidad a Problemas Inversos: El método opera directamente sobre las mediciones indirectas disponibles en el momento de la inferencia, lo que lo hace adecuado para tareas de imagen computacional como la reconstrucción de TC y el desenfoque (deblurring).
• Capacidad de Localización: A diferencia de los detectores de OOD globales, KLIP puede identificar la ubicación espacial precisa de desplazamientos de distribución sutiles (por ejemplo, pequeños tumores o artefactos) dentro de una imagen.
• Generalización: Se demuestra que la métrica generaliza a través de diferentes arquitecturas de modelos de difusión (predictor-corrector y basados en parches), diferentes conjuntos de datos (TC abdominales y rostos humanos) y diferentes problemas inversos (TC de vista dispersa y desenfoque Gaussiano).

Resultados Experimentales
Los autores evaluaron KLIP en dos tareas principales:

1. Reconstrucción de TC de Vista Dispersa: Utilizando tomografías computarizadas abdominales sanas como datos ID, el modelo fue probado en imágenes con artefactos de estrella sintéticos y tumores hepáticos sintéticos realistas. KLIP logró altos puntajes de Área Bajo la Curva (AUC) tanto para la detección a nivel de conjunto de datos como a nivel de imagen, superando a las líneas base como DiffPath, la Log-Verosimilitud Negativa (NLL) y detectores de anomalías que no son de difusión (CutPaste, SimpleNet) en el contexto de mediciones indirectas.
2. Desenfoque Gaussiano de Rostos: Utilizando CelebA como datos ID, el método detectó características OOD localizadas como cicatrices sintéticas y prótesis faciales distintivas en personajes de películas (por ejemplo, Star Trek, Harry Potter).

Los resultados cualitativos (Figuras 4 y 5) demuestran que los mapas de calor de KLIP localizan con éxito las características OOD (estrellas, tumores, cicatrices, prótesis) con alta precisión. Los estudios de ablación confirmaron que tanto la restricción de bloque como la de paso de tiempo son necesarias para un rendimiento óptimo; el uso de la imagen completa o del rango de tiempo completo reduce significativamente la capacidad de detección.

Significancia y Reivindicaciones
El artículo afirma que KLIP proporciona una solución robusta y libre de calibración para detectar y localizar desplazamientos de distribución sutiles en problemas inversos, un escenario donde los métodos tradicionales suelen fallar debido a la falta de datos de entrenamiento OOD. Al aprovechar la dinámica temporal del muestreo de difusión y la partición espacial, KLIP ofrece una "ventana rica" al proceso de muestreo bajo un desplazamiento de distribución. Los autores enfatizan que este enfoque es particularmente valioso para la imagen médica y otras aplicaciones de seguridad crítica donde la identificación de anomalías localizadas y raras es crucial, y donde la adquisición de datos OOD representativos para la calibración es imposible.

El artículo reconoce limitaciones, señalando que KLIP es sensible a los hiperparámetros (tamaño del bloque y ventana de tiempo), los cuales dependen de la escala y la forma de las características OOD, y que las evaluaciones actuales dependen de configuraciones de "crimen inverso" (usando el mismo modelo de propagación hacia adelante para la simulación y la reconstrucción). Sin embargo, los resultados sugieren un fuerte potencial de generalización a través de diversas tareas de imagen.