Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

Este artigo identifica o fenômeno de esquecimento de prompts em Transformers de Difusão Multimodais e propõe uma abordagem sem treinamento chamada "reinjeção de prompts", que realimenta representações das camadas iniciais nas posteriores para melhorar a aderência às instruções e a qualidade geral da geração de imagens.

O essencial

Imagine que você está pedindo para um artista muito talentoso, mas um pouco esquecido, desenhar uma cena muito específica para você.

Você diz: "Desenhe um gato preto com bigodes brancos sentado em cima de um balão vermelho."

O artista começa a desenhar. Nas primeiras pinceladas, ele ouve você perfeitamente. Mas, conforme a pintura avança e ele entra em detalhes complexos (como a textura da pele do gato ou o brilho do balão), ele começa a esquecer os detalhes que você pediu. No final, o resultado pode ser um gato marrom, sem bigodes, ou o balão pode ter sumido.

Esse é o problema que os cientistas descobriram nos modelos de IA mais modernos de geração de imagens (chamados MMDiTs, como o SD3, FLUX e Qwen-Image). Eles chamam esse fenômeno de "Esquecimento do Prompt".

O Que Está Acontecendo? (A Analogia da Mensagem)

Pense no processo de criação da imagem como uma equipe de 50 pessoas passando uma mensagem de um lado para o outro em um corredor longo.

1. O Pedido: Você dá a mensagem inicial ("Gato preto no balão vermelho") para a primeira pessoa.
2. O Processo: A mensagem passa de pessoa para pessoa. Cada pessoa adiciona detalhes visuais (cores, sombras, formas).
3. O Problema: Como a equipe só recebe "nota" se a imagem final ficar bonita, as pessoas no meio do corredor começam a focar apenas na beleza da imagem e esquecem de manter a mensagem original intacta. Quando a mensagem chega na última pessoa (a camada final da IA), ela já está tão distorcida que o "gato preto" virou "gato cinza" e o "balão vermelho" virou "balão azul".

A IA está tão focada em fazer a imagem ficar realista que ela perde a fidelidade ao que você pediu.

A Solução: "Reinjeção do Prompt"

Os autores do artigo (Yuxuan Yao e equipe) criaram uma solução inteligente e simples, chamada Prompt Reinjection (Reinjeção do Prompt). Eles não precisaram reensinar a IA (o que seria caro e demorado); eles apenas mudaram como ela trabalha enquanto cria a imagem.

Aqui está a analogia da solução:

Imagine que, a cada 5 passos que a mensagem percorre pelo corredor, alguém coloca um bilhete de lembrete na mão da pessoa atual. Esse bilhete é uma cópia exata do que você pediu no início ("Gato preto no balão vermelho").

• Como funciona na prática: A IA pega a "memória" fresca do seu pedido (que está guardada nas camadas iniciais, onde a informação ainda está pura) e a reinsere (reinjecta) nas camadas profundas, onde a IA está trabalhando nos detalhes finais.
• O Truque: Eles não apenas jogam a mensagem de volta. Eles usam um "tradutor" matemático (chamado alinhamento geométrico) para garantir que a mensagem antiga se encaixe perfeitamente na nova linguagem que a IA está usando naquele momento, sem bagunçar o desenho.

O Que Isso Mudou?

Com essa "lembrança constante", os resultados foram impressionantes:

1. Contagem Correta: Se você pedir "4 cachorros", a IA agora desenha exatamente 4, não 3 ou 5.
2. Posição Certa: Se você pedir "o gato em cima do balão", o gato não fica flutuando ao lado ou embaixo.
3. Cores e Texturas: O "gato preto" continua preto até o final, e não vira marrom.

Por que isso é importante?

Antes, para corrigir esses erros, você teria que treinar a IA do zero (como ir para a escola de novo), o que custa milhões e leva meses. Com essa técnica, você só precisa de um pequeno ajuste no momento da criação. É como dar um "empurrãozinho" na memória da IA para ela não se perder no caminho.

Resumo da Ópera:
A IA moderna é incrível, mas tem uma memória de curto prazo ruim quando está criando coisas complexas. Os autores criaram um sistema de "lembretes" que mantém a IA focada no seu pedido original do início ao fim, garantindo que a imagem final seja exatamente o que você imaginou, sem esquecer os detalhes.

Resumo técnico

Título: Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

Autores: Yuxuan Yao, Yuxuan Chen, Hui Li, et al. (Fudan University, Alibaba, Baidu, etc.)

1. O Problema: Esquecimento de Prompt (Prompt Forgetting)

O artigo identifica um fenômeno crítico em Transformers de Difusão Multimodais (MMDiTs), como SD3, SD3.5, FLUX e Qwen-Image. Diferentemente das arquiteturas anteriores (U-Net) onde o texto atua como uma condição externa estática, os MMDiTs processam tokens de texto e latentes visuais juntos em uma única pilha de transformadores, permitindo interação bidirecional.

No entanto, os autores observam que, à medida que a profundidade da rede aumenta (através das camadas do transformador), a semântica do prompt na ramificação de texto é progressivamente esquecida.

• Causa Raiz: Existe um desequilíbrio de supervisão. A função de perda (loss) é definida apenas no espaço latente visual (para prever ruído ou imagem limpa). Os tokens visuais recebem supervisão direta, enquanto os tokens de texto são atualizados apenas indiretamente através do mecanismo de atenção conjunta.
• Consequência: Para minimizar o erro de reconstrução visual, o modelo pode sacrificar a preservação de detalhes finos do prompt (como atributos específicos, contagem ou relações espaciais) nas camadas mais profundas. Isso resulta em uma deriva (drift) das representações textuais, onde informações token a token tornam-se irrecoveráveis.

2. Metodologia e Análise

Os autores validaram esse fenômeno através de duas etapas de análise em várias arquiteturas (SD3, SD3.5, FLUX, Qwen-Image):

1. Análise Observacional:
   • CKNNA (Conditional K-Nearest Neighbor Alignment): Mediu a preservação da estrutura semântica local. Os resultados mostraram um declínio monotônico na similaridade entre os vizinhos dos tokens de texto conforme a profundidade aumenta.
   • Projeção PCA: Visualizou o deslocamento global das distribuições, mostrando que os tokens de texto colapsam em regiões altamente concentradas do espaço latente, perdendo separabilidade.
2. Quantificação Funcional (Probing):
   • Treinaram classificadores leves (MLPs) para decodificar atributos linguísticos (substantivos, adjetivos, relações espaciais, numerais) a partir das representações intermediárias de cada camada.
   • Resultado: A precisão de recuperação dos atributos cai drasticamente com a profundidade, especialmente para relações espaciais, que sofrem a degradação mais severa.

3. Solução Proposta: Prompt Reinjection (Reinjeção de Prompt)

Para mitigar esse esquecimento sem necessidade de re-treinamento (training-free), os autores propõem o Prompt Reinjection. A ideia é re-injetar representações textuais de camadas rasas (onde a semântica é fiel) nas camadas profundas durante a inferência.

O método consiste em duas fases principais para garantir uma fusão estável:

1. Ancoragem de Distribuição (Distribution Anchoring):
   • Como as estatísticas (média e variância) das camadas rasas e profundas diferem, aplica-se Normalização de Camada (Layer Normalization) às características de origem e destino antes da fusão.
   • Após a adição, as características são projetadas de volta para a distribuição estatística original da camada alvo.
2. Alinhamento Geométrico (Geometry Alignment):
   • Para corrigir discrepâncias no sistema de coordenadas latentes (rotação), utiliza-se uma Transformação Ortogonal de Procrustes.
   • Um mapa de rotação ótimo ($R$) é calculado uma única vez (usando um conjunto de dados de calibração como COCO-5K) para alinhar as características da camada de origem com o manifold da camada alvo.
   • A fórmula de reinjeção é: $T_{final} = T_{alvo} + w \cdot (T_{origem} \cdot R)$, onde $w$ é um hiperparâmetro de peso.

4. Resultados Principais

Os experimentos foram realizados em benchmarks padrão (GenEval, DPG-Bench, T2I-CompBench++) e métricas de qualidade humana (HPSv2, ImageReward, PickScore, CLIP).

• Melhoria na Seguir Instruções (Instruction Following):
  • O método melhorou consistentemente a capacidade de seguir instruções em todos os modelos testados.
  • GenEval: SD3.5 teve um aumento de 6,48% e FLUX de 5,64% na pontuação geral.
  • Tarefas Específicas: Houve ganhos significativos em tarefas de posicionamento espacial (onde o esquecimento é mais crítico), contagem de objetos e ligação de atributos (cor/textura).
• Preservação da Qualidade Visual:
  • As métricas de preferência humana e alinhamento global (CLIP) foram mantidas ou ligeiramente melhoradas, indicando que a correção semântica não introduz artefatos ou degrada a fidelidade da imagem.
• Análise Qualitativa:
  • Imagens geradas com Prompt Reinjection obedecem muito melhor a restrições complexas (ex: "quatro cães", "um gato branco e bigodes pretos", "um balão em cima de um pássaro") comparado aos modelos base.

5. Contribuições e Significância

• Descoberta Teórica: É a primeira análise rigorosa que quantifica e explica o "esquecimento de prompt" em MMDiTs, atribuindo-o ao desequilíbrio de supervisão entre os ramos de texto e imagem.
• Método Prático e Eficiente: O Prompt Reinjection é uma intervenção de inferência sem treinamento (training-free). Não requer ajuste de pesos do modelo, apenas uma calibração inicial de rotação e a adição de características durante a geração.
• Custo Computacional: O overhead é mínimo (apenas ~8% a mais de FLOPs por bloco devido à rotação), tornando-o viável para uso em produção.
• Impacto Geral: O método demonstra que a fidelidade semântica em modelos de difusão modernos pode ser restaurada simplesmente mantendo o "sinal" original do prompt vivo ao longo da profundidade da rede, resolvendo um gargalo fundamental na geração de imagens complexas e composicionais.

Em resumo, o trabalho oferece uma solução elegante para um problema fundamental na evolução dos modelos de geração de imagem, permitindo que modelos como FLUX e SD3.5 atendam a prompts complexos com uma precisão muito superior, sem a necessidade de re-treinamento custoso.