Your Pre-trained Diffusion Model Secretly Knows Restoration

Título: Your Pre-trained Diffusion Model Secretly Knows Restoration

Autores: Sudarshan Rajagopalan e Vishal M. Patel (Johns Hopkins University)

1. Problema e Motivação

O trabalho aborda o desafio da Restauração Tudo-em-Um (AiOR - All-in-One Restoration), que visa recuperar imagens ou vídeos corrompidos por diversas degradações (como neblina, chuva, neve, desfoque de movimento e baixa luminosidade) usando um único modelo unificado.

• Limitações das Abordagens Atuais: Métodos baseados em modelos de difusão pré-treinados (como Stable Diffusion, FLUX, WAN) geralmente dependem de:

  1. Ajuste Fino (Fine-tuning): Que é custoso computacionalmente e pode enviesar o modelo para a distribuição de dados de treinamento, reduzindo a robustez em cenários do mundo real.
  2. Módulos de Controle (ex: ControlNet): Que introduzem parâmetros adicionais e podem enfraquecer os prios (conhecimentos prévios) fortes do modelo base.
  3. Métodos de Edição (ex: SDEdit, Inversão): Tentativas de usar prompts de texto naturais ou otimização de tokens para restauração falham. O modelo tende a apenas remover o ruído adicionado, retornando à imagem degradada original, em vez de restaurar o conteúdo limpo.

• A Questão Central: Os prios de grandes modelos de difusão contêm conhecimento relevante para restauração? Se sim, por que não conseguimos acessá-lo através de prompts de texto padrão ou otimização de tokens?

2. Metodologia Proposta

Os autores propõem uma abordagem que mantém o "backbone" (núcleo) do modelo de difusão congelado e aprende apenas prompts leves. A metodologia divide-se em duas descobertas principais:

A. Otimização no Espaço de Embedding (Não em Tokens)

• Falha do Espaço de Tokens: A otimização direta de tokens de texto (mesmo com textual inversion) ou o uso de prompts de linguagem natural não consegue elicitar comportamento de restauração. O modelo ignora a instrução de "remover neblina" e apenas denoise a imagem, mantendo a degradação.
• Solução no Espaço de Embedding: Os autores demonstram que otimizar diretamente as representações de contexto no espaço de embedding (na saída do codificador de texto, antes de entrar no mecanismo de atenção cruzada) desbloqueia o potencial de restauração. Ao aprender vetores de embedding ($p$) em vez de tokens, o modelo consegue mapear consistentemente estados ruidosos e degradados para alvos limpos.

B. Resolução do Desalinhamento de Trajetória (Diffusion Bridge)

• O Problema: Treinar prompts de forma ingênua cria uma incompatibilidade entre o treinamento e a inferência.
  • Treinamento: O estado intermediário é construído como $z_t = (1-t)z_{deg} + t\epsilon$ (ancorado na imagem degradada $z_{deg}$).
  • Inferência: O processo de amostragem reversa deve mover-se de $z_{deg}$ para $z_{clean}$.
  • Isso faz com que o prompt aprendido seja aplicado a estados intermediários que ele nunca viu durante o treinamento, causando instabilidade e artefatos.
• A Solução (Ponte de Difusão - Diffusion Bridge): Para alinhar a dinâmica de treinamento e inferência, os autores utilizam uma formulação de ponte de difusão.
  • Eles definem estados intermediários que transitam suavemente da imagem degradada para a limpa, adicionando ruído de forma controlada.
  • Especificamente, eles adotam uma Ponte Orientada a Energia (EBR - Energy-Oriented Bridge). Diferente de pontes padrão (DDBM) que podem ter baixa variância nas extremidades (limitando a capacidade do modelo de "reorientar" a trajetória cedo), a EBR garante que o estado inicial já contenha ruído sobre a entrada degradada e que o nível de ruído diminua monotonicamente à medida que o sinal transita para o estado limpo.
  • Isso garante que o prompt aprenda a guiar o modelo através de estados intermediários idênticos aos encontrados durante a amostragem real.

3. Contribuições Principais

1. Descoberta de Capacidade Latente: Evidenciam que modelos de difusão pré-treinados possuem conhecimento de restauração intrínseco, mas este é inacessível via prompts de texto ou otimização de tokens, sendo acessível apenas via otimização de embeddings de contexto.
2. Identificação e Correção de Desalinhamento: Identificam o problema de desalinhamento de trajetória (train-test mismatch) em métodos de aprendizado de prompts e propõem o uso de uma formulação de ponte de difusão (EBR) para alinhar as dinâmicas de treinamento e inferência.
3. Eficiência e Generalização: Implementam o framework nos modelos WAN (vídeo, 1.3B parâmetros) e FLUX (imagem, 12B parâmetros). O método utiliza apenas prompts leves aprendidos (sem ajuste fino do backbone e sem módulos de controle específicos), alcançando desempenho competitivo com muito menos parâmetros treináveis.

4. Resultados Experimentais

Os autores realizaram extensas avaliações em tarefas de restauração de imagem e vídeo:

• Desempenho Quantitativo:
  • No modelo FLUX (imagem), a abordagem supera ou compete com métodos state-of-the-art (SOTA) como AutoDIR, PixWizard e FoundIR em métricas de distorção (PSNR, SSIM) e percepção (LPIPS, CLIPIQA, MUSIQ).
  • No modelo WAN (vídeo), supera métodos dedicados como ViWS-Net e AverNet, especialmente em métricas de qualidade de vídeo (DOVER).
• Generalização (OOD e Mistos):
  • O método demonstra forte generalização em conjuntos de dados fora da distribuição (OOD), misturas de degradações (ex: chuva + neve) e degradações nunca vistas (ex: câmeras sob display), mantendo alta qualidade perceptual onde outros métodos falham ou produzem artefatos.
• Ablações:
  • Confirmam que a otimização no espaço de embedding supera a otimização de tokens.
  • Demonstram que a ponte EBR supera abordagens ingênuas e pontes DDBM padrão, eliminando artefatos visuais e melhorando a correção de degradações severas.
  • A injeção residual do prompt (aprendendo um resíduo sobre o contexto nulo) mostrou-se mais eficiente em parâmetros e performance do que otimizar o embedding completo.

5. Significado e Impacto

Este trabalho é significativo porque:

• Desmistifica a Restauração com Difusão: Mostra que não é necessário re-treinar massivamente ou adicionar controladores complexos para usar modelos generativos em tarefas de baixo nível (restauração). O conhecimento já existe no modelo.
• Eficiência de Recursos: Oferece uma solução leve (apenas alguns milhares de parâmetros de prompts) que pode ser aplicada a modelos massivos de 1B+ parâmetros, tornando a restauração de alta qualidade mais acessível.
• Novo Paradigma de Controle: Estabelece que o espaço de embedding de saída do codificador de texto é um local muito mais rico para controlar o comportamento de difusão do que o espaço de tokens de entrada, abrindo novas direções para o controle de modelos generativos em tarefas de visão computacional.

Em resumo, o artigo demonstra que, com a engenharia correta de prompts (espaço de embedding + alinhamento de trajetória via ponte de difusão), modelos de difusão pré-treinados podem se tornar restauradores poderosos e generalistas sem perder seus prios originais.