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

Sintesi Tecnica: KLIP – Rilevamento di Shift di Distribuzione Localizzato tramite Divergenza KL con Prior di Diffusione

Problema
Il documento affronta la sfida di rilevare immagini Out-of-Distribution (OOD) e, cosa più critica, di localizzare spazialmente le caratteristiche OOD localizzate all'interno delle immagini durante la risoluzione di problemi inversi (ad esempio, la ricostruzione di immagini da misurazioni indirette). I metodi esistenti per il rilevamento OOD si affidano spesso a dataset di calibrazione contenenti esempi OOD, che sono invece assenti per la "coda lunga" di anomalie rare negli scenari del mondo reale (ad esempio, specifiche anomalie mediche). Inoltre, molti approcci attuali operano su immagini complete piuttosto che sulle misurazioni indirette disponibili nella tomografia computazionale, o non riescono a rilevare shift localizzati e sottili (come piccoli tumori o lacerazioni) che sono essenziali per il valore diagnostico.

Metodologia
Gli autori propongono KLIP, una metrica di rilevamento OOD priva di calibrazione basata sulla divergenza di Kullback-Leibler (KL) tra la distribuzione prior di diffusione $p(x)$ e la distribuzione posteriore $p(x|y)$ dato il set di misurazioni $y$.

1. Intuizione Centrale: Il documento postula che se un'immagine contiene caratteristiche OOD, la misurazione $y$ "tirerà" la distribuzione posteriore significativamente lontano dalla distribuzione prior. Di conseguenza, la divergenza KL $D_{KL}(p(x|y) \parallel p(x))$ dovrebbe essere maggiore per le immagini OOD e, specificamente, maggiore per le regioni OOD all'interno di un'immagine.
2. Stima tramite Campionamento Posteriore: Poiché la divergenza KL esatta è intrattabile, gli autori la approssimano utilizzando le funzioni di score generate durante il campionamento posteriore della diffusione. Decomponendo lo score posteriore in score prior e score di verosimiglianza ($\nabla_x \log p_t(y|x)$), la divergenza KL viene stimata integrando la norma al quadrato dello score di verosimiglianza lungo la traiettoria di campionamento (Equazione 11).
3. Restrizione di Blocco e di Timestep: Per rilevare shift localizzati piuttosto che globali, gli autori introducono due restrizioni chiave:
   • Restrizione Spaziale di Blocco: L'immagine viene divisa in blocchi ($B_i$). La divergenza KL viene calcolata localmente per ogni blocco per evitare che il segnale di una piccola caratteristica OOD venga diluito dal resto dell'immagine In-Distribution (ID).
   • Restrizione di Timestep: L'integrazione è limitata a finestre temporali specifiche $[t_0, t_1]$ della traiettoria di campionamento. Basandosi sull'osservazione che i modelli di diffusione generano caratteristiche a diverse scale in diversi timestep, ciò consente alla metrica di concentrarsi sulle fasi specifiche in cui le caratteristiche OOD sono più distinte.

La metrica finale, KLIP, è definita come la divergenza KL ristretta su un determinato blocco spaziale e una finestra temporale (Equazione 12).

Contributi Chiave

• Rilevamento Senza Calibrazione: KLIP non richiede dati di calibrazione OOD o la conoscenza della distribuzione alterata; si affida esclusivamente a un modello di diffusione addestrato su dati In-Distribution (ID).
• Applicabilità ai Problemi Inversi: Il metodo opera direttamente sulle misurazioni indirette disponibili al momento dell'inferenza, rendendolo adatto a compiti di imaging computazionale come la ricostruzione TC e la deblurring gaussiana.
• Capacità di Localizzazione: A differenza dei rilevatori OOD globali, KLIP può identificare l'esatta posizione spaziale di sottili shift di distribuzione (ad esempio, piccoli tumori o artefatti) all'interno di un'immagine.
• Generalizzazione: La metrica dimostra di generalizzare attraverso diversi architetture di modelli di diffusione (predictor-corrector e patch-based), diversi dataset (TC addominali e volti umani) e diversi problemi inversi (CT a vista sparsa e deblurring gaussiano).

Risultati Sperimentali
Gli autori hanno valutato KLIP su due compiti primari:

1. Ricostruzione TC a Vista Sparsa: Utilizzando scansioni TC addominali sane come dati ID, il modello è stato testato su immagini con artefatti stellari sintetici e tumori epatici sintetici realistici. KLIP ha ottenuto punteggi AUC elevati sia per il rilevamento a livello di dataset che a livello di immagine, superando i baseline come DiffPath, la Log-Verosimiglianza Negativa (NLL) e i rilevatori di anomalie non basati su diffusione (CutPaste, SimpleNet) nel contesto di misurazioni indirette.
2. Deblurring Gaussiano di Volti: Utilizzando CelebA come dati ID, il metodo ha rilevato caratteristiche OOD localizzate come cicatrici sintetiche e protesi facciali distintive su personaggi di film (ad esempio, Star Trek, Harry Potter).

I risultati qualitativi (Figure 4 e 5) dimostrano che le mappe di calore KLIP localizzano con successo le caratteristiche OOD (stelle, tumori, cicatrici, protesi) con alta precisione. Gli studi di ablazione hanno confermato che sia la restrizione di blocco che quella di timestep sono necessarie per prestazioni ottimali; l'utilizzo dell'intera immagine o dell'intero intervallo di timestep riduce significativamente la capacità di rilevamento.

Significatività e Rivendicazioni
Il documento afferma che KLIP fornisce una soluzione robusta e priva di calibrazione per rilevare e localizzare sottili shift di distribuzione nei problemi inversi, un contesto in cui i metodi tradizionali spesso falliscono a causa della mancanza di dati di addestramento OOD. Sfruttando la dinamica temporale del campionamento di diffusione e la partizione spaziale, KLIP offre una "finestra ricca" sul processo di campionamento sotto shift di distribuzione. Gli autori sottolineano che questo approccio è particolarmente prezioso per l'imaging medico e altri casi di uso critico per la sicurezza, dove l'identificazione di anomalie rare e localizzate è cruciale e dove l'acquisizione di dati OOD rappresentativi per la calibrazione è impossibile.

Il documento riconosce i limiti, notando che KLIP è sensibile agli iperparametri (dimensione del blocco e finestra temporale) che dipendono dalla scala e dalla forma delle caratteristiche OOD, e che le attuali valutazioni si affidano a setup di "crimine inverso" (utilizzando lo stesso modello forward per simulazione e ricostruzione). Tuttavia, i risultati suggeriscono un forte potenziale di generalizzazione attraverso diversi compiti di imaging.