VIPaint: Image Inpainting with Pre-Trained Diffusion Models via Variational Inference

1. Formulação do Problema

O artigo aborda o desafio do preenchimento de imagens (image inpainting) (preenchimento de regiões ausentes ou mascaradas de uma imagem) utilizando Modelos de Difusão Pré-treinados (DMs), particularmente Modelos de Difusão Latente (LDMs) como o Stable Diffusion.

• A Dificuldade Central: Embora os DMs se destaquem na geração de imagens novas a partir de ruído, condicioná-los a imagens corrompidas ou parcialmente observadas (inpainting) é não trivial. O objetivo é amostrar a partir da verdadeira distribuição posterior condicional $p(x|y)$, onde $y$ é a imagem observada (mascarada) e $x$ é a imagem completa.
• Limitações dos Métodos Existentes:
  • Métodos Heurísticos/Iterativos: Abordagens como RePaint, CoPaint e DPS (Diffusion Posterior Sampling) frequentemente dependem de refinamento iterativo ou heurísticas de "viagem no tempo". Estas frequentemente falham em capturar a estrutura global quando grandes regiões são mascaradas, levando a resultados incoerentes ou irreais.
  • Métodos Variacionais (REDdiff): Tentativas anteriores de inferência variacional (por exemplo, REDdiff) frequentemente otimizam uma única imagem no nível livre de ruído ($t=0$) usando uma aproximação Gaussiana simples. Isso falha em capturar incerteza, é propenso a instabilidade de otimização (especialmente em regimes de baixo ruído) e luta para generalizar para LDMs.
  • Requisitos de Treinamento: Muitos métodos do estado da arte exigem o ajuste fino (fine-tuning) do modelo generativo em conjuntos de dados ou máscaras específicos, o que é computacionalmente caro e impede a aplicação "zero-shot" a novas tarefas.

2. Metodologia: VIPaint

Os autores propõem o VIPaint, um algoritmo de inferência variacional (VI) hierárquica que trata o modelo de difusão pré-treinado como um prior. Em vez de otimizar uma única imagem, o VIPaint otimiza uma distribuição sobre variáveis latentes em múltiplos níveis de ruído.

A. Posterior Variacional Hierárquica

O VIPaint define uma aproximação não-Gaussiana e Markoviana da verdadeira posterior de difusão. Ele introduz um conjunto de $K$ "pontos-chave" (instantes intermediários) $\{h(1), \dots, h(K)\}$ dentro da trajetória de difusão reversa (onde $h(K)$ é alto ruído e $h(1)$ é baixo ruído).

A distribuição variacional é definida como:
$$q_\lambda(z_{h(1):h(K)}) = q_\lambda(z_{h(K)}) \prod_{i=1}^{K-1} q_\lambda(z_{h(i)} | z_{h(i+1)})$$

• Estrutura:
  • Alto Ruído ($h(K)$): Modelado como um Gaussiano fatorado.
  • Ruído Intermediário/Baixo: Modelado como um Gaussiano onde a média é uma combinação convexa da previsão de difusão prior (do modelo pré-treinado) e um parâmetro variacional aprendido específico da imagem.
  • Variância: A variância é aprendida e específica por dimensão, permitindo que o modelo aumente dinamicamente a incerteza em regiões mascaradas enquanto mantém certeza em regiões observadas.

B. A Função Objetivo

O VIPaint otimiza um limite inferior variacional sobre a verossimilhança marginal da observação $y$. A função de perda consiste em três termos:

1. Perda de Reconstrução: Garante que as amostras geradas no nível de ruído mais baixo ($h(1)$) sejam consistentes com os pixels observados $y$. Isso é aproximado usando uma previsão de remoção de ruído de um passo.
2. Perda de Prior: Regulariza o nível de ruído mais alto ($h(K)$) para alinhar com a semântica global codificada pelo prior de difusão no regime de alto ruído.
3. Perda Hierárquica: Garante consistência entre pontos-chave adjacentes na hierarquia, capturando detalhes e estrutura de nível médio da imagem.

C. Otimização e Amostragem

• Otimização: Os parâmetros variacionais $\lambda$ (médias, variâncias e pesos de mistura) são otimizados por imagem usando descida de gradiente (Adam). Crucialmente, isso é feito sem annealing (amostragem aleatória de tempo é usada para a perda), o que fornece um limite não tendencioso e evita a instabilidade vista em métodos que fazem annealing de $T$ para $0$.
• Amostragem: Após a otimização, reconstruções diversas são geradas por:
  1. Amostragem Ancestral: Amostragem a partir da posterior hierárquica otimizada de $h(K)$ até $h(1)$.
  2. Refinamento: Refinamento adicional da amostra de $h(1)$ até $t=0$ usando a rede de remoção de ruído pré-treinada, guiada pelo gradiente de verossimilhança (semelhante ao DPS, mas começando de uma posterior melhor inicializada).

3. Contribuições Principais

1. Aproximação Não-Gaussiana Hierárquica: Ao contrário de métodos VI anteriores que aproximam a posterior em um único ponto ($t=0$), o VIPaint modela a posterior através de uma hierarquia de níveis de ruído. Isso permite que ele capture simultaneamente semântica global (em alto ruído) e detalhes locais (em baixo ruído).
2. Adaptação Zero-Shot: O VIPaint não requer o ajuste fino do modelo de difusão pré-treinado. Ele adapta o modelo a novas máscaras e imagens otimizando parâmetros variacionais para cada consulta, tornando-o aplicável a qualquer DM pré-treinado (baseado em pixels ou latente).
3. Estabilidade e Modelagem de Incerteza: Ao evitar o regime de baixo ruído durante a fase principal de otimização (onde redes de remoção de ruído são não suaves e instáveis) e modelar explicitamente a variância, o VIPaint produz imputações diversas e plausíveis que refletem a incerteza dos dados ausentes.
4. Generalização para LDMs: O método é explicitamente projetado para funcionar com Modelos de Difusão Latente, permitindo preenchimento de alta resolução (por exemplo, 1024x1024) com menor custo computacional do que abordagens baseadas em pixels.

4. Resultados Experimentais

Os autores avaliaram o VIPaint em ImageNet64, ImageNet256, LSUN-Churches e Stable Diffusion 3.5 (condicionado a texto).

• Desempenho Quantitativo:

  • Máscaras Grandes: O VIPaint supera significativamente as linhas de base (DPS, RePaint, CoPaint, REDdiff) em regiões mascaradas grandes (40-80% da imagem) em termos de PSNR, SSIM e LPIPS.
  • Modelos Condicionados a Texto: No Stable Diffusion 3.5, o VIPaint alcança PSNR superior (26,69 vs 25,22 para DPS) e LPIPS (0,050 vs 0,061), produzindo imagens globalmente coerentes que preservam detalhes finos.
  • Outros Problemas Inversos: O método também mostra forte desempenho em desembaçamento e super-resolução, superando linhas de base especializadas como PSLD e ReSample.

• Desempenho Qualitativo:

  • Resultados visuais mostram que o VIPaint gera estruturas globalmente coerentes (por exemplo, formas de objetos consistentes e iluminação) onde outros métodos produzem artefatos desfocados ou incoerentes.
  • Ele gera com sucesso múltiplas completions plausíveis e diversas para a mesma imagem mascarada, demonstrando sua capacidade de modelar a natureza multimodal da posterior.

• Eficiência:

  • O VIPaint é computacionalmente eficiente. Para $K=2$, requer apenas ~150 chamadas à rede de remoção de ruído para otimização (comparado a 1000+ para métodos baseados em amostragem) e é mais rápido do que muitas linhas de base enquanto entrega maior qualidade.

5. Significância

O VIPaint representa um avanço significativo em problemas inversos com modelos de difusão. Ao mudar da amostragem heurística ou otimização de ponto único para um framework de inferência variacional hierárquica principial, ele resolve a questão crítica da incoerência global no preenchimento de grandes máscaras.

Sua capacidade de funcionar zero-shot com modelos fundacionais do estado da arte (como o Stable Diffusion 3.5) sem re-treinamento torna-o uma ferramenta altamente prática para tarefas reais de edição, restauração e completamento de imagens. Ele preenche a lacuna entre o rigor teórico da inferência variacional e as capacidades práticas da IA generativa moderna.