Resumo Técnico: Problemas Inversos com Modelos de Difusão via Perda de Busca de Modo

Definição do Problema

O artigo aborda o desafio de resolver problemas inversos utilizando modelos de difusão incondicionais pré-treinados. Um problema inverso envolve estimar uma amostra de dados original $x_0$ dada uma medição parcialmente degradada $y$, governada pela equação $y = A(x_0) + \eta$, onde $A$ é um operador de degradação conhecido e $\eta$ é um ruído Gaussiano.

Embora os modelos de difusão pré-treinados possam resolver esses problemas de maneira zero-shot (plug-and-play), as abordagens existentes enfrentam limitações significativas:

1. Métodos de Amostragem de Posterior: Estes visam extrair amostras proporcionais à probabilidade posterior $p(x_0|y)$. No entanto, eles frequentemente dependem de aproximações de modelagem (ex: aproximações Gaussianas do score condicional) ou métodos de Monte Carlo sequencial, que podem ser computacionalmente exigentes e podem não aderir estritamente ao posterior verdadeiro sem validação de ground-truth.
2. Métodos de Estimativa MAP: Estes visam encontrar a estimativa de Máxima Posterior (MAP) (o modo da posterior). Os resolvedores existentes frequentemente utilizam aproximações heurísticas, dependem de divisão de variáveis (ex: HQS, ADMM) ou requerem modelos de consistência. Essas abordagens frequentemente envolvem computações de gradiente caras e suposições de modelagem que podem não se sustentar para todos os operadores lineares.

Os autores argumentam que muitas tarefas de restauração de imagem priorizam a geração de uma imagem plausível consistente com as medições, em vez de cobrir toda a distribuição posterior. Isso se alinha mais de perto com a estimativa MAP, motivando a necessidade de uma estratégia MAP mais fundamentada e livre de aproximações.

Metodologia: Perda de Busca de Modo Variacional (VML)

A contribuição central é a Perda de Busca de Modo Variacional (VML), uma nova função objetivo projetada para guiar as amostras geradas em direção à estimativa MAP durante o processo de difusão reversa.

Fundamentação Teórica

Os autores derivam a VML de uma perspectiva variacional ao minimizar a divergência de Kullback-Leibler (KL) entre o posterior de difusão $p(x_0|x_t)$ e o posterior de medição $p(x_0|y)$ em cada passo de difusão reversa $t$:
$$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$$

A intuição é que, conforme $t \to 0$, a distribuição $p(x_0|x_t)$ torna-se cada vez mais concentrada em torno de $x_t$. Ao encontrar um $x_t^*$ ótimo que minimiza essa divergência KL em cada passo, o processo guia a amostra em direção aos modos da posterior $p(x_0|y)$.

Derivação em Forma Fechada para Operadores Lineares

Um marco teórico fundamental é a derivação de uma expressão em forma fechada para a VML no caso de problemas inversos lineares ($A(x) = Hx$) sem qualquer aproximação de modelagem. A expressão completa da VML inclui:

1. Um termo de consistência de medição: $\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$.
2. Termos de prior: Envolvendo o denoiser $D(x_t, t)$, o cronograma de ruído $\sigma_t$ e a covariância da posterior $Cov[x_0|x_t]$.

Objetivos Simplificados e Reponderados

A VML completa contém termos de ordem superior envolvendo a matriz de covariância $Cov[x_0|x_t]$, que são computacionalmente caros para calcular e diferenciar. Os autores provam que esses termos de ordem superior convergem uniformemente para zero (ou uma constante) conforme $t \to 0$. Consequentemente, eles propõem uma VML Simplificada ($VML_S$) que descarta esses termos, retendo apenas os termos de consistência de medição e de prior.

Para melhorar ainda mais a estabilidade da otimização, eles introduzem uma VML simplificada reponderada ($VML_S^\rho$):
$$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$$
onde $\rho_t$ é um cronograma de peso de prior tal que $\rho_t \to 1$ conforme $t \to 0$. Esta formulação permite a computação eficiente de gradientes usando diferenciação automática padrão.

Algoritmos

Baseados na minimização da VML, os autores propõem:

1. VML-MAP: Um algoritmo que utiliza o gradiente descendente para minimizar $VML_S^\rho$ em cada passo de difusão reversa.
2. VML-MAP$_{pre}$: Uma variante pré-condicionada para operadores lineares mal condicionados. Ela utiliza a Decomposição em Valores Singulares (SVD) da matriz de degradação $H$ para construir um pré-condicionador $P$, acelerando a convergência para tarefas difíceis como deblurring (desfoque).
3. VML-MAP$^\tau$: Uma estratégia para restauração de imagem com ruído. Para evitar artefatos agudos causados por desafios de otimização em níveis baixos de ruído, o algoritmo minimiza a VML apenas até um limiar de tempo de difusão $\tau$, e então muda para a amostragem DDIM padrão para gerar a imagem final.
4. LatentVML-MAP: Uma extensão para Modelos de Difusão Latente (LDMs), aproximando a VML no espaço latente.

Principais Contribuições

1. Nova Função de Perda: Introdução da VML, que teoricamente direciona as amostras para estimativas MAP ao minimizar a divergência KL entre os posteriores de difusão e de medição.
2. Derivação Livre de Aproximação: Para problemas inversos lineares, a VML é derivada em forma fechada sem aproximações de modelagem heurísticas, fornecendo uma base teórica formal para a estimativa MAP com modelos de difusão.
3. Análise de Convergência: Prova formal de que os objetivos de VML simplificados e reponderados convergem para o log-posterior negativo (o objetivo MAP) conforme $t \to 0$ sob suposições moderadas.
4. Algoritmos Práticos: Desenvolvimento da VML-MAP e de sua variante pré-condicionada, oferecendo uma solução empiricamente eficaz para diversas tarefas de restauração de imagem.
5. Validação Extensa: Demonstração de desempenho superior sobre baselines existentes (ex: DDRM, $\Pi$GDM, MAPGA, DAPS) em múltiplos datasets (ImageNet, FFHQ, CelebA) e tarefas (inpainting, super-resolução, deblurring).

Resultados Experimentais

Os autores avaliaram seus métodos em tarefas de restauração de imagem usando modelos pré-treinados no ImageNet256, FFHQ256 e CelebA.

• Métricas: O artigo enfatiza métricas perceptuais (LPIPS e FID) sobre métricas de pixel-a-pixel (PSNR, SSIM), argumentando que estas últimas frequentemente falham em capturar a qualidade perceptual de reconstruções plausíveis em problemas inversos mal definidos.
• Desempenho:
  • O VML-MAP superou consistentemente métodos de amostragem de posterior (DDRM, $\Pi$GDM) e outros estimadores MAP (MAPGA, DAPS) em termos de LPIPS e FID em tarefas de inpainting, super-resolução e deblurring.
  • O VML-MAP$_{pre}$ mostrou melhorias significativas sobre a versão não pré-condicionada para tarefas mal condicionadas (ex: deblurring), alcançando pontuações de LPIPS e FID menores.
  • Eficiência Computacional: O VML-MAP alcançou melhor qualidade perceptual com menor tempo de execução (menos avaliações de função neural) em comparação com métodos que requerem muitos passos de difusão reversa.
• Restauração com Ruído: A estratégia VML-MAP$^\tau$ lidou com sucesso com medições ruidosas, produzindo reconstruções de alta qualidade onde o VML-MAP padrão apresentava dificuldades com artefatos.
• Modelos Latentes: O LatentVML-MAP demonstrou eficácia no CelebA, embora os autores notem que a otimização permanece mais desafiadora no espaço latente devido à não-linearidade do decoder.

Significância e Alegações

O artigo afirma que a VML representa a primeira modelagem livre de aproximação para estimativa MAP para problemas inversos lineares, utilizando apenas um modelo de difusão incondicional pré-treinado, apoiada por uma análise teórica formal.

Os autores enfatizam que sua abordagem difere fundamentalmente da amostragem de posterior (que visa a cobertura da distribuição) e dos resolvedores MAP existentes (que frequentemente dependem de heurísticas ou modelos de consistência). Ao minimizar a VML, o método efetivamente visa os "modos" da posterior, o que os autores argumentam ser o comportamento desejado para gerar restaurações de imagem plausíveis e de alta qualidade.

O trabalho é modesto quanto às suas limitações, reconhecendo que:

• A eficácia depende da disponibilidade de otimizadores praticamente eficientes para o objetivo VML.
• Desafios de otimização persistem para operadores mal condicionados ou não lineares, embora o pré-condicionamento ajude.
• Trabalhos futuros são necessários para desenvolver otimizadores de ordem superior para mitigar ainda mais os problemas de otimização sem custos computacionais proibitivos.

O artigo conclui que o VML-MAP fornece um framework principado, teoricamente fundamentado e empiricamente eficaz para resolver problemas inversos com modelos de difusão, preenchendo a lacuna entre a estimativa MAP teórica e a restauração prática de imagens.