Resumo Técnico: Modulação Consciente de Qualidade para Transformers de Difusão

Declaração do Problema

Os modelos modernos de difusão de texto para imagem, particularmente os Diffusion Transformers (DiT) como o Stable Diffusion 3.5 (SD3.5), dependem de sinais de condicionamento — especificamente embeddings de tempo (timestep) e embeddings de prompt agrupados (pooled) — para modular o processo de denoising via Adaptive Layer Normalization (AdaLN). Embora esses sinais comuniquem efetivamente o nível atual de ruído e a intenção semântica, eles carecem de informação consciente da qualidade. Consequentemente, o mecanismo de modulação não consegue reagir diretamente a modos de falha específicos da amostra, como acúmulo de erro, interpretação incorreta do prompt ou características alucinadas que emergem conforme a imagem latente evolui. O condicionamento base é independente do estado atual do latente ($\partial M_{base} / \partial x_t = 0$), impedindo que o modelo corrija problemas de fidelidade baseados na qualidade real da imagem intermediária.

Metodologia

Os autores propõem o Quality Representation Module (QRM), um adaptador leve baseado em transformer projetado para injetar sinais conscientes de qualidade no processo de difusão sem modificar o backbone pré-treinado.

Arquitetura

O QRM é implementado como um decodificador transformer que prevê offsets de modulação aditivos ($M_{qrm}$) aos parâmetros da AdaLN base ($M_{base}$).

• Entradas: O módulo recebe a imagem latente atual ($x_t$), o embedding de prompt agrupado ($p$), o embedding de tempo ($t$) e os parâmetros de modulação base ($M_{base}$).
• Processamento:
  1. A imagem latente $x_t$ é projetada em tokens latentes via uma convolução $1\times1$ aprendida e pooling adaptativo.
  2. Os parâmetros de modulação base são normalizados e projetados para formar tokens de consulta (query tokens) ($Z_0$).
  3. Uma memória de chave-valor ($KV$) é construída a partir dos tokens latentes, do prompt e dos embeddings de tempo.
  4. Através de camadas de cross-attention e self-attention, os tokens de consulta atendem à memória $KV$ para gerar uma representação refinada.
  5. A saída é projetada de volta para as larguras nativas da AdaLN para produzir o vetor de offset $M_{qrm}$.
• Aplicação: A modulação final aplicada aos blocos do transformer é $M_t = M_{base,t} + M_{qrm,t}$.

Estratégia de Treinamento (ReFL)

O QRM é treinado usando uma estratégia de Aprendizado por Feedback de Recompensa (ReFL):

• Componentes Congelados: O backbone de difusão (denoiser, codificadores de texto, VAE), o sampler e o modelo de recompensa permanecem congelados. Apenas os parâmetros do QRM são atualizados.
• Objetivo: O módulo aprende a maximizar uma pontuação de recompensa escalar derivada de um modelo de recompensa congelado (ex: HPSv2.1 ou CLIPScore) aplicado a uma imagem decodificada de um passo intermediário de denoising.
• Funções de Perda: O artigo utiliza uma perda de maximização de recompensa ou uma perda de busca de margem (margin-seeking loss), que empurra a imagem modulada pelo QRM para superar a base por uma margem específica.
• Seleção de Timestep: Para melhorar a eficiência e focar na formação estrutural, o QQRM é ativado apenas durante um subconjunto predefinido de timesteps de alto ruído (ex: $t \in \{50, \dots, 26\}$), em vez de todo o cronograma de difusão.

Inferência

Durante a inferência, o módulo QRM treinado atua como um adaptador plug-in. Ele prevê offsets de modulação baseados no estado latente atual e os adiciona aos parâmetros base. Isso permite que o modelo direcione a trajetória de denoising para estados de maior qualidade sem retreinar o backbone ou alterar o cronograma de amostragem.

Principais Contribuições

1. Arquitetura QRM: Introdução de um módulo transformer leve que aprende representações conscientes de qualidade a partir de entradas existentes do modelo (latente, prompt, tempo) para refinar parâmetros de modulação.
2. Integração Não Invasiva: Um mecanismo para aumentar a modulação do transformer de difusão com sinais de qualidade da imagem latente sem modificar o backbone de difusão, o VAE ou o cronograma de amostragem subjacente.
3. Modulação Guiada por Recompensa: Demonstração de que atualizações de modulação previstas ($M_{qrm}$) podem efetivamente refinar o processo de denoising em cada timestep, melhorando a aderência ao prompt e a fidelidade da imagem.
4. Avaliação Abrangente: Estudos de ablação extensos analisando os efeitos do tamanho da arquitetura QRM, intervalos de timesteps ativos e métricas de recompensa de treinamento na qualidade da imagem.

Resultados Experimentais

Os autores avaliaram o QRM usando o Stable Diffusion 3.5 (SD3.5) no benchmark Parti-Prompts.

• Desempenho Quantitativo:
  • Comparado ao baseline do SD3.5, o QRM treinado com HPSv2.1 como sinal de recompensa alcançou as maiores pontuações entre todos os métodos avaliados (incluindo refinamento de prompt, fine-tuning e abordagens ReNeg).
  • Melhorias: +0,20 em Aesthetic score e +0,82 em HPSv2.1 score.
  • Quando treinado com CLIPScore, o QRM melhorou o alinhamento com a preferência humana (HPSv2.1) mantendo uma qualidade perceptual comparável.
• Análise Qualitativa:
  • Visualizações das trajetórias de denoising mostram que aplicar o QRM durante os timesteps iniciais de alto ruído leva a estruturas semânticas reconhecíveis emergindo significativamente mais cedo do que no baseline.
  • Aplicar o QRM apenas nos estágios intermediários ou tardios resultou em melhorias menores, confirmando que a estrutura global é estabelecida precocemente no processo de difusão.
• Descobertas de Ablação:
  • Capacidade: Modelos QRM maiores otimizados mais fortemente para a recompensa de treinamento específica (ex: HPSv2.1) mostraram ligeiros trade-offs em outras métricas (ex: CLIPScore).
  • Intervalo de Timestep: A modulação durante timesteps iniciais ($50 \to 26$) foi consistentemente mais eficaz para o alinhamento do que a modulação em estágios posteriores.
  • Métrica de Recompensa: A escolha do modelo de recompensa (CLIPScore vs. HPSv2.1) moldou diretamente os vieses perceptuais aprendidos pelo QRM.

Significância e Alegações

O artigo alega que intervenções ao nível de modulação fornecem um caminho escalável e eficaz para alinhar grandes modelos de difusão com sinais de preferência humana. Ao contrário de métodos que exigem um fine-tuning caro de todo o backbone (ex: LoRA ou fine-tuning total), o QRM alcança melhorias consistentes de qualidade operando inteiramente através de atualizações de modulação aditivas condicionadas à representação latente em evolução.

Os autores posicionam o QRM como um adaptador flexível e consciente da qualidade que pode ser direcionado por diferentes sinais de recompensa e aplicado a várias arquiteturas de diffusion transformer. Eles reconhecem que o QRM herda os vieses dos modelos de recompensa de treinamento e sugerem que trabalhos futuros poderiam explorar o treinamento de múltiplas recompensas e a seleção de timestep adaptativa para aumentar ainda mais a robustez e a generalização. O trabalho demonstra que ganhos significativos na fidelidade da imagem e no alinhamento do prompt podem ser alcançados sem alterar os pesos pré-treinados de modelos de última geração.