Adaptive Moments are Surprisingly Effective for Plug-and-Play Diffusion Sampling

Resumo Técnico: Momentos Adaptativos para Amostragem de Difusão Plug-and-Play

1. O Problema

Os modelos de difusão tornaram-se o estado da arte em geração de conteúdo (imagens, áudio, vídeo). Uma técnica popular é a geração condicional "Plug-and-Play", que permite amostrar de uma distribuição condicional $p(x|y)$ usando um modelo de difusão treinado apenas na distribuição marginal $p(x)$, sem necessidade de re-treinamento específico para a tarefa.

O desafio central reside na aproximação da pontuação de verossimilhança (likelihood score), $\nabla_{x_t} \log p(y|x_t)$. Na prática, este termo é intratável e deve ser aproximado. Métodos existentes, como Diffusion Posterior Sampling (DPS) e Classifier Guidance (CG), utilizam aproximações que introduzem ruído significativo nas dinâmicas de amostragem. Esse ruído surge devido à limitação das informações de condicionamento (especialmente em tarefas difíceis como super-resolução extrema ou desfoque severo) e à propagação de gradientes através de redes neurais grandes. O ruído leva a trajetórias de amostragem instáveis, oscilações e degradação na qualidade da imagem gerada, especialmente quando as condições são esparsas.

2. Metodologia Proposta

Os autores propõem uma abordagem simples, porém eficaz: aplicar Estimação de Momentos Adaptativos (inspirada no algoritmo de otimização Adam) para estabilizar os gradientes de orientação (guidance) ruidosos durante o processo de amostragem.

• Mecanismo: Em vez de desenvolver aproximações de verossimilhança mais complexas, os autores mantêm médias móveis exponenciais dos primeiros e segundos momentos das estimativas de pontuação de verossimilhança ao longo dos passos de amostragem.
  • O primeiro momento (momento) suaviza a trajetória de otimização acumulando informações de gradiente de passos anteriores.
  • O segundo momento escala adaptativamente as atualizações com base na variância histórica de cada componente do gradiente.
• Implementação:
  • AdamDPS: Aplica a estimativa adaptativa ao método DPS (onde o gradiente é calculado sobre a estimativa de imagem limpa $x_{0|t}$).
  • AdamCG: Aplica a mesma lógica ao Classifier Guidance (onde o gradiente é calculado diretamente sobre o latente ruidoso $x_t$).
• Fórmula: O gradiente estabilizado $\hat{g}_t$ é calculado como:
  $$\hat{g}_t = \frac{\hat{m}_k}{\sqrt{\hat{v}_k} + \delta}$$
  Onde $\hat{m}_k$ e $\hat{v}_k$ são os momentos corrigidos por viés, e $\delta$ é uma constante de estabilidade numérica.

3. Contribuições Principais

1. Melhoria Significativa com Complexidade Mínima: Demonstra-se que a estimativa de momentos adaptativos melhora substancialmente métodos de orientação existentes (DPS e CG), alcançando resultados state-of-the-art com apenas uma modificação de poucas linhas de código.
2. Análise Empírica Robusta: Fornecem análises detalhadas em dados sintéticos e reais, mostrando que a mitigação de ruído de gradiente via momentos adaptativos melhora a alinhamento com a condição desejada.
3. Descoberta sobre Dificuldade da Tarefa: Revelam que a dificuldade da tarefa impacta drasticamente o desempenho relativo dos métodos. Enquanto métodos complexos (como TFG - Training-Free Guidance) funcionam bem em tarefas fáceis, eles degradam rapidamente em cenários difíceis (condicionamento limitado), onde a abordagem simples de momentos adaptativos supera consistentemente os baselines.

4. Resultados Experimentais

Os autores avaliaram seus métodos em tarefas de restauração de imagem (super-resolução, desfoque, inpainting) e geração condicional de classes (CIFAR-10, ImageNet).

• Restauração de Imagem:
  • O AdamDPS superou todos os métodos concorrentes (incluindo DPS, TFG, UGD, MPGD) em todas as tarefas de restauração no ImageNet e no conjunto de dados "Cats".
  • Em super-resolução 16x e desfoque gaussiano intenso, onde outras métodos produziam artefatos visuais ou imagens borradas, o AdamDPS gerou detalhes mais nítidos e realistas.
  • Em tarefas de inpainting (máscara de 90%), o AdamDPS manteve a superioridade, enquanto o TFG (que usa iterações recursivas) falhou em superar o DPS em cenários mais difíceis.
• Amostragem Condicional de Classes:
  • No CIFAR-10, o AdamDPS superou o próximo melhor método (DPS) em 9,86% de precisão de classificação.
  • No ImageNet, métodos padrão atingiram precisão próxima ao acaso aleatório (~1%), enquanto o AdamDPS alcançou 10,49%.
  • Na orientação com classificadores "time-aware" (sensíveis ao tempo), o AdamCG melhorou a precisão em mais de 19 pontos percentuais em relação ao CG padrão.
• Análise de Dificuldade e Passos de Amostragem:
  • À medida que a dificuldade aumenta (ex: super-resolução de 4x para 16x), métodos complexos como o TFG degradam-se rapidamente, ficando abaixo do DPS simples. O AdamDPS, no entanto, mantém uma melhoria consistente sobre o DPS em todos os níveis de dificuldade.
  • O método adiciona overhead computacional negligenciável (tempo de parede) em comparação ao DPS, sendo mais rápido que métodos recursivos como o TFG.
• Análise de Trajetória:
  • Visualizações mostram que as trajetórias do DPS frequentemente divergem da verdade fundamental (ground truth) devido a gradientes conflitantes (similaridade cossena negativa entre passos adjacentes).
  • O AdamDPS mantém uma similaridade cossena positiva consistente, indicando progressos estáveis e coerentes em direção à condição desejada.

5. Significado e Conclusão

O trabalho demonstra que a mitigação de ruído é um fator crítico e subexplorado na orientação de difusão plug-and-play. Em vez de tentar criar aproximações de verossimilhança cada vez mais sofisticadas e computacionalmente caras, estabilizar os gradientes existentes através de momentos adaptativos (Adam) oferece uma solução simples e poderosa.

A principal implicação é que, para problemas inversos mal condicionados (onde a informação de condicionamento é escassa e ruidosa), a robustez da otimização é mais importante do que a complexidade do modelo de orientação. Isso abre novas direções para a aplicação de modelos de difusão em domínios científicos e inversos onde os dados são limitados, oferecendo uma alternativa eficiente e de fácil implementação aos métodos atuais. O código está disponível publicamente.