Resumo Técnico: Amostragem MAP Local para Modelos de Difusão

Declaração do Problema

A Amostragem de Posterior Difusiva (DPS) surgiu como um framework poderoso para resolver problemas inversos amostrando a partir da distribuição posterior $p(x_0 | y)$, oferecendo quantificação de incerteza fundamentada e multimodalidade. No entanto, muitas aplicações práticas em imageamento e ciência priorizam a estimativa pontual precisa sobre a amostragem distribucional completa. O estimador de Máxima A Posteriori (MAP), definido como $\arg \max p(x_0 | y)$, permanece um objetivo padrão de reconstrução.

Métodos existentes baseados em otimização (por exemplo, Resample, DiffPIR, DCDP) tentam melhorar o desempenho de reconstrução alternando entre etapas de remoção de ruído e otimização. No entanto, a conexão teórica entre esses procedimentos iterativos, a solução global de MAP e a amostragem posterior permanece obscura. Não está estabelecido se esses métodos convergem para o MAP global ou se são consistentes com a DPS. Além disso, a natureza específica das etapas de otimização nesses métodos — se resolvem um objetivo global ou local — requer esclarecimento.

Metodologia: Amostragem MAP Local (LMAPS)

Os autores introduzem a Amostragem MAP Local (LMAPS), um novo framework de inferência que resolve iterativamente subproblemas de máxima a posteriori local ao longo da trajetória de difusão.

1. Framework Teórico

A LMAPS opera resolvendo uma sequência de problemas de MAP local em vez de um único problema global. Em cada etapa de difusão $t$, dado um estado ruidoso $x_t$ e uma medição $y$, o método busca:
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
Isso é seguido por uma etapa de transição DDIM padrão:
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

Distinções em Relação aos Métodos Existentes:

• vs. MAP Global: O MAP global busca o modo da posterior completa $p(x_0 | y)$. A LMAPS busca o modo da posterior condicional $p(x_0 | x_t, y)$ em cada etapa. Os autores argumentam que a sequência de maximizadores locais não necessariamente converge para o maximizador global de $p(x_0 | y)$.
• vs. DPS: A DPS utiliza a média condicional $E[x_0 | x_t, y]$ na atualização. A LMAPS substitui essa média pelo modo condicional. Estes são equivalentes apenas se a posterior condicional for Gaussiana; caso contrário, a LMAPS introduz um viés de busca pelo modo, evitando regiões de baixa densidade entre os modos.
• vs. DMAP: Embora a DMAP também utilize uma perspectiva de MAP local, ela otimiza no espaço de variáveis de difusão latente ($x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$). A LMAPS formula o problema diretamente no espaço do sinal limpo ($x^*_0$), fornecendo uma visão complementar que conecta explicitamente à verossimilhança da medição $p(y|x_0)$.

2. Algoritmo Prático e Aproximações

Para tornar o problema de MAP local tratável para problemas inversos não lineares gerais, os autores propõem aproximações e reformulações específicas:

• Aproximação de Covariância: Os autores aproximam a covariância condicional $\Sigma_{0|t}$ do prior $p(x_0 | x_t)$ usando uma forma isotrópica motivada por uma suposição de prior Gaussiano. À medida que o nível de ruído $\sigma_t \to 0$, a covariância posterior torna-se assintoticamente isotrópica. Eles aproximam $\Sigma_{0|t} \approx \frac{1}{SNR} I$, onde $SNR = \alpha_t^2 / \sigma_t^2$.

• Reformulação do Objetivo: O objetivo de MAP local é formulado como a minimização de uma soma ponderada da distância do prior e da fidelidade aos dados:
  $$x^*_0 = \arg \min_{x_0} \left\{ \frac{SNR}{k} \|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2} \|y - H(x_0)\|^2 \right\}$$
  Para melhorar a estabilidade numérica e a interpretabilidade, isso é reformulado em uma combinação convexa invariante à escala:
  $$x^*_0 = \arg \min_{x_0} \left\{ (1 - \mu_t) \frac{1}{2} \|x_0 - m_{0|t}\|^2 + \mu_t k_2 \|y - H(x_0)\|^2 \right\}$$
  onde $\mu_t = \sigma_t^2 / (\sigma_t^2 + k_1^2)$. Esta formulação permite um recozimento automático: à medida que $t$ diminui ($\sigma_t \to 0$), o objetivo muda de ser impulsionado pela medição para ser impulsionado pelo prior.

• Linear vs. Não Linear: Para problemas inversos lineares, a solução de MAP local pode ser derivada em forma fechada, recuperando a Difusão Projetada de Momentos de Tweedie (TMPD) como um caso especial. Para problemas não lineares, a solução é obtida via descida de gradiente.

Contribuições Principais

1. Unificação Teórica: O artigo formula a LMAPS para unificar métodos de problemas inversos baseados em otimização e a Difusão Projetada de Momentos de Tweedie (TMPD) sob um único framework probabilístico. Ele esclarece que métodos baseados em otimização inerentemente resolvem problemas de MAP local em vez de MAP global ou amostragem de posterior completa.
2. Avanços Metodológicos: Os autores introduzem uma aproximação de covariância motivada por Gaussiana e uma função objetivo reformulada que aprimora a estabilidade numérica e a interpretabilidade através de recozimento automático.
3. Validação Empírica: O framework é validado em 10 tarefas de restauração de imagem (incluindo operadores lineares, não lineares e não diferenciáveis) e 3 problemas inversos científicos.

Resultados Experimentais

Os autores avaliaram a LMAPS nos conjuntos de dados FFHQ e ImageNet para restauração de imagem e em benchmarks para problemas inversos científicos (Espalhamento Inverso Linear, Ressonância Magnética por Sensação Compressiva e Imageamento de Buraco Negro).

• Restauração de Imagem: A LMAPS alcançou desempenho state-of-the-art em 43 dos 60 casos nas métricas PSNR, SSIM e LPIPS no FFHQ e ImageNet. Ela consistentemente superou a linha de base DMAP relacionada, mostrando melhorias de mais de 2 dB de PSNR em tarefas como desfocagem de movimento, restauração JPEG e quantização.
• Problemas Inversos Científicos: A LMAPS alcançou consistentemente o maior PSNR em todas as tarefas científicas. Notavelmente, alcançou ganhos superiores a 1,5 dB de PSNR em três tarefas de Espalhamento Inverso Linear (LIS) em comparação com métodos anteriores.
• Eficiência: A LMAPS demonstrou eficiência superior em comparação com a DAPS, alcançando melhores resultados com menos recursos computacionais (por exemplo, menos etapas de otimização por etapa de difusão).

Significado e Alegações

O artigo afirma que a LMAPS fornece uma interpretação probabilística fundamentada para métodos baseados em otimização, fechando a lacuna entre amostragem posterior (DPS) e estimativa pontual (MAP). Ao resolver explicitamente subproblemas de MAP local, o método oferece uma visão unificada que explica o comportamento dos solucionadores baseados em otimização existentes.

Os autores enfatizam que, embora a LMAPS não amostre da distribuição posterior completa (e, portanto, não capture a multimodalidade da mesma forma que a DPS), ela é especificamente projetada para cenários onde a estimativa pontual precisa é o objetivo principal. O framework oferece uma alternativa robusta que equilibra a fidelidade de reconstrução com os priores estruturais dos modelos de difusão, alcançando resultados state-of-the-art tanto em restauração de imagem padrão quanto em tarefas desafiadoras de reconstrução científica. O trabalho sugere que pesquisas futuras poderiam explorar ainda mais a utilidade da estimativa de MAP global com priores de difusão, o que permanece um desafio aberto.