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

Résumé Technique : KLIP – Détection de décalage de distribution localisée via la divergence de KL avec des priors de diffusion

Énoncé du problème
L'article traite du défi de la détection d'images hors distribution (OOD - Out-of-Distribution) et, plus crucialement, de la localisation spatiale de caractéristiques OOD au sein des images lors de la résolution de problèmes inverses (par exemple, la reconstruction d'images à partir de mesures indirectes). Les méthodes de détection OOD existantes reposent souvent sur des ensembles de données de calibration contenant des exemples OOD, qui sont indisponibles pour la « longue traîne » d'anomalies rares dans des scénarios en monde ouvert (par exemple, des anomalies médicales spécifiques). De plus, de nombreuses approches actuelles opèrent sur des images complètes plutôt que sur les mesures indirectes disponibles dans l'imagerie computationnelle, ou échouent à détecter des décalages localisés subtils (tels que de petites tumeurs ou des déchirures) qui sont essentiels pour la valeur diagnostique.

Méthodologie
Les auteurs proposent KLIP, une métrique de détection OOD sans calibration basée sur la divergence de Kullback-Leibler (KL) entre la distribution de l'a priori de diffusion $p(x)$ et la distribution a posteriori $p(x|y)$ étant donné les mesures $y$.

1. Intuition centrale : L'article postule que si une image contient des caractéristiques OOD, la mesure $y$ « tirera » la distribution a posteriori de manière significative loin de la distribution a priori. Par conséquent, la divergence KL $D_{KL}(p(x|y) \parallel p(x))$ devrait être plus grande pour les images OOD et spécifiquement plus grande pour les régions OOD au sein d'une image.
2. Estimation via l'échantillonnage a posteriori : Puisque la divergence KL exacte est intractale, les auteurs l'approchent en utilisant les fonctions de score générées lors de l'échantillonnage a posteriori de diffusion. En décomposant le score a posteriori en un score a priori et un score de vraisemblance ($\nabla_x \log p_t(y|x)$), la divergence KL est estimée en intégrant la norme au carré du score de vraisemblance sur la trajectoire d'échantillonnage (Équation 11).
3. Restriction par bloc et par pas de temps : Pour détecter des décalages localisés plutôt que globaux, les auteurs introduisent deux restrictions clés :
   • Restriction spatiale par blocs : L'image est divisée en blocs ($B_i$). La divergence KL est calculée localement pour chaque bloc afin d'éviter que le signal d'une petite caractéristique OOD ne soit dilué par le reste de l'image en distribution (ID).
   • Restriction par pas de temps : L'intégration est restreinte à des fenêtres temporelles spécifiques $[t_0, t_1]$ de la trajectoire d'échantillonnage. En s'appuyun sur l'observation que les modèles de diffusion génèrent des caract factuelles à différentes échelles à différents pas de temps, cela permet à la métrique de se concentrer sur les étapes spécifiques où les caractéristiques OOD sont les plus distinctes.

La métrique finale, KLIP, est définie comme la divergence KL restreinte sur un bloc spatial et une fenêtre temporelle spécifiques (Équation 12).

Contributions clés

• Détection sans calibration : KLIP ne nécessite aucune donnée de calibration OOD ni connaissance de la distribution décalée ; il repose uniquement sur un modèle de diffusion entraîné sur des données en distribution (ID).
• Applicabilité aux problèmes inverses : La méthode opère directement sur les mesures indirectes disponibles au moment de l'inférence, ce qui la rend adaptée aux tâches d'imagerie computationnelle telles que la reconstruction CT et le défloutage (deblurring).
• Capacité de localisation : Contrairement aux détecteurs OOD globaux, KLIP peut identifier l'emplacement spatial précis des décalages de distribution subtils (par exemple, de petites tumeurs ou des artefacts) au sein d'une image.
• Généralisation : La métrique montre qu'elle se généralise à travers différentes architectures de modèles de diffusion (prédictions-correcteurs et basées sur des patchs), différents jeux de données (CT abdominaux et visages humains) et différents problèmes inverses (CT à vue sparse et défloutage gaussien).

Résultats expérimentaux
Les auteurs ont évalué KLIP sur deux tâches principales :

1. Reconstruction CT à vue sparse (Sparse-View CT) : En utilisant des scanners CT abdominaux sains comme données ID, le modèle a été testé sur des images présentant des artefacts stellaires synthétiques et des tumeurs hépatiques synthétiques réalistes. KLIP a obtenu des scores AUC élevés pour la détection au niveau du jeu de données et au niveau de l'image, surpassant les bases de référence comme DiffPath, la log-vraisemblance négative (NLL) et les détecteurs d'anomalies non-diffusion (CutPaste, SimpleNet) dans le contexte de mesures indirectes.
2. Défloutage gaussien de visages : En utilisant CelebA comme données ID, la méthode a détecté des caractéristiques OOD localisées telles que des cicatrices synthétiques et des prothèses faciales distinctives sur des personnages de films (ex: Star Trek, Harry Potter).

Les résultats qualitatifs (Figures 4 et 5) démontrent que les cartes thermiques (heatmaps) de KLIP localisent avec précision les caractéristiques OOD (étoiles, tumeurs, cicatrices, prothèses). Les études d'ablation ont confirmé que les restrictions par bloc et par pas de temps sont nécessaires pour une performance optimale ; utiliser l'image complète ou toute la plage de pas de temps réduit considérablement la capacité de détection.

Signification et affirmations
L'article affirme que KLIP fournit une solution robuste et sans calibration pour détecter et localiser des décalages de distribution subtils dans les problèmes inverses, un cadre où les méthodes traditionnelles échouent souvent en raison de l'absence de données d'entraînement OOD. En exploitant la dynamique temporelle de l'échantillonnage de diffusion et le partitionnement spatial, KLIP offre une « fenêtre riche » sur le processus d'échantillonnage sous décalage de distribution. Les auteurs soulignent que cette approche est particulièrement précieuse pour l'imagerie médicale et d'autres applications critiques pour la sécurité où l'identification d'anomalies rares et localisées est cruciale, et où l'acquisition de données OOD représentatives pour la calibration est impossible.

L'article reconnaît des limites, notant que KLIP est sensible aux hyperparamètres (taille du bloc et fenêtre temporelle) qui dépendent de l'échelle et de la forme des caractéristiques OOD, et que les évaluations actuelles reposent sur des configurations de « crime inverse » (utilisant le même modèle direct pour la simulation et la reconstruction). Cependant, les résultats suggèrent un fort potentiel de généralisation à travers diverses tâches d'imagerie.