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

Resumo Técnico: KLIP – Detecção de Desvio de Distribuição Localizado via Divergência KL com Priores de Difusão

Definição do Problema
O artigo aborda o desafio de detectar imagens Fora da Distribuição (OOD - Out-of-Distribution) e, mais criticamente, localizar características OOD espacialmente localizadas em imagens durante a solução de problemas inversos (por exemplo, reconstrução de imagem a partir de medições indiretas). Os métodos existentes de detecção de OOD frequentemente dependem de conjuntos de dados de calibração contendo exemplos OOD, que não estão disponíveis para a "cauda longa" de anormalidades raras em cenários de mundo aberto (por exemplo, anomalias médicas específicas). Além disso, muitos métodos atuais operam sobre imagens completas em vez das medições indiretas disponíveis na imagem computacional, ou falham em detectar desvios sutis e localizados (como pequenos tumores ou rasgos) que são essenciais para o valor diagnóstico.

Metodologia
Os autores propõem o KLIP, uma métrica de detecção de OOD livre de calibração baseada na divergência de Kullback-Leibler (KL) entre a distribuição prior de difusão $p(x)$ e a distribuição posterior $p(x|y)$ dado as medições $y$.

1. Intuição Central: O artigo postula que, se uma imagem contém características OOD, a medição $y$ "puxará" a distribuição posterior significativamente para longe da distribuição prior. Consequentemente, a divergência KL $D_{KL}(p(x|y) \parallel p(x))$ deve ser maior para imagens OOD e, especificamente, maior para regiões OOD dentro de uma imagem.
2. Estimativa via Amostragem Posterior: Como a divergência KL exata é intratável, os autores a aproximam usando as funções de score geradas durante a amostragem posterior de difusão. Ao decompor o score da posterior no score da prior e no score da verossimilhança ($\nabla_x \log p_t(y|x)$), a divergência KL é estimada integrando a norma ao quadrado do score da verossimilhança ao longo da trajetória de amostragem (Equação 11).
3. Restrição de Bloco e de Passo de Tempo (Timestep): Para detectar desvios localizados em vez de globais, os autores introduzem duas restrições fundamentais:
   • Restrição de Bloco Espacial: A imagem é dividida em blocos ($B_i$). A divergência KL é computada localmente para cada bloco para evitar que o sinal de uma pequena característica OOD seja diluído pelo restante da imagem que está dentro da distribuição (In-Distribution - ID).
   • Restrição de Passo de Tempo (Timestep): A integração é restrita a janelas de tempo específicas $[t_0, t_1]$ da trajetória de amostragem. Baseando-se na observação de que modelos de difusão geram características em diferentes escalas em diferentes passos de tempo, isso permite que a métrica foque nos estágios específicos onde as características OOD são mais distintas.

A métrica final, KLIP, é definida como a divergência KL restrita sobre um bloco espacial e uma janela de tempo específicos (Equação 12).

Principais Contribuições

• Detecção Livre de Calibração: O KLIP não requer dados de calibração OOD ou conhecimento da distribuição deslocada; ele depende exclusivamente de um modelo de difusão treinado em dados ID.
• Aplicabilidade a Problemas Inversos: O método opera diretamente nas medições indiretas disponíveis no momento da inferência, tornando-o adequado para tarefas de imagem computacional como reconstrução de CT e desfoque (deblurring).
• Capacidade de Localização: Diferente dos detectores de OOD globais, o KLIP pode identificar a localização espacial precisa de desvios de distribuição sutis (ex: pequenos tumores ou artefatos) dentro de uma imagem.
• Generalização: A métrica mostra-se capaz de generalizar através de diferentes arquiteturas de modelos de difusão (preditor-corretor e baseados em patches), diferentes conjuntos de dados (CTs abdominais e rostos humanos) e diferentes problemas inversos (CT de visão esparsa e desfoque Gaussiano).

Resultados Experimentais
Os autores avaliaram o KLIP em duas tarefas principais:

1. Reconstrução de CT de Visão Esparsa: Usando tomografias computadorizadas abdominais saudáveis como dados ID, o modelo foi testado em imagens com artefatos de estrela sintéticos e tumores hepáticos sintéticos realistas. O KLIP alcançou altos escores de Área Sob a Curva (AUC) tanto para detecção ao nível do conjunto de dados quanto ao nível da imagem, superando baselines como DiffPath, Log-Verossimilhança Negativa (NLL) e detectores de anomalia não baseados em difusão (CutPaste, SimpleNet) no contexto de medições indiretas.
2. Desfoque Gaussiano de Rostos: Usando o CelebA como dados ID, o método detectou características OOD localizadas, como cicatrizes sintéticas e próteses faciais distintas em personagens de filmes (ex: Star Trek, Harry Potter).

Resultados qualitativos (Figuras 4 e 5) demonstram que os mapas de calor do KLIP localizam com sucesso as características OOD (estrelas, tumores, cicatrizes, próteses) com alta precisão. Estudos de ablação confirmaram que tanto a restrição de bloco quanto a de passo de tempo são necessárias para o desempenho ideal; usar a imagem completa ou o intervalo total de passos de tempo reduz significativamente a capacidade de detecção.

Significância e Alegações
O artigo afirma que o KLIP fornece uma solução robusta e livre de calibração para detectar e localizar desvios de distribuição sutis em problemas inversos, um cenário onde os métodos tradicionais frequentemente falham devido à falta de dados de treinamento OOD. Ao alavancar a dinâmica temporal da amostragem de difusão e a partição espacial, o KLIP oferece uma "janela rica" para o processo de amostragem sob desvio de distribuição. Os autores enfatizam que esta abordagem é particularmente valiosa para imagens médicas e outras aplicações críticas de segurança, onde identificar anomalias localizadas e raras é crucial e onde adquirir dados OOD representativos para calibração é impossível.

O artigo reconhece limitações, observando que o KLIP é sensível a hiperparâmetros (tamanho do bloco e janela de tempo) que dependem da escala e forma das características OOD, e que as avaliações atuais dependem de configurações de "crime inverso" (usando o mesmo modelo direto para simulação e reconstrução). No entanto, os resultados sugerem um forte potencial de generalização para diversas tarefas de imagem.