Resumo Técnico: Domando Tokens Outlier em Transformers de Difusão

Declaração do Problema

O artigo investiga o fenômeno de tokens outlier dentro de Transformers de Difusão (DiTs) utilizados para geração de imagens. Embora trabalhos anteriores (por exemplo, Darcet et al.) tenham estabelecido que Vision Transformers (ViTs) em tarefas de reconhecimento frequentemente produzem um pequeno subconjunto de tokens com normas de características anormalmente grandes, que absorvem atenção desproporcional enquanto carregam informações locais limitadas, o papel desses tokens em modelos generativos permanecia inexplorado.

Os autores identificam duas lacunas críticas:

1. Ubiquidade em Pipelines Generativas: Tokens outlier aparecem não apenas nos codificadores ViT pré-treinados (tokenizadores) de pipelines modernas de Autoencoder de Representação (RAE), mas também internamente dentro dos próprios Transformers de Difusão (denoizers).
2. Natureza do Defeito: A degradação causada por esses tokens não é meramente resultado de alguns valores de perda extremos. Em vez disso, os autores hipotetizam que tokens outlier são sintomáticos de semânticas de patches locais corrompidas. Simplesmente mascarar tokens de alta norma durante o treinamento não melhora o desempenho, sugerindo que o problema reside na perda de estrutura espacial e informação local, e não apenas em instabilidade numérica.

Metodologia: Registradores de Duplo Estágio (DSR)

Para abordar essas questões, os autores propõem Registradores de Duplo Estágio (DSR), uma estratégia de intervenção unificada que aplica tokens registradores a ambos os componentes do pipeline RAE-DiT: o codificador de representação e o gerador de difusão.

1. Registradores em Codificadores Visuais

Os autores abordam tokens outlier no estágio do codificador (o tokenizador) usando duas abordagens dependendo da disponibilidade de registradores pré-treinados:

• Registradores Treinados: Para codificadores como DINOv2 que suportam registradores, os autores utilizam tokens registradores pré-treinados para suprimir artefatos de alta norma nos mapas de características do codificador.
• Registradores Recursivos de Tempo de Teste (TTR): Para codificadores sem registradores incorporados (por exemplo, SigLIP2), os autores anexam tokens adicionais não treinados à sequência de entrada no momento da inferência. Em casos onde um segundo estágio de outliers é detectado na representação resultante, eles aplicam TTR recursivamente para estabilizar a saída.

2. Registradores em Transformers de Difusão

Os autores observam que, mesmo com um codificador estabilizado, o próprio Transformer de Difusão desenvolve tokens outlier internos, predominantemente em camadas intermediárias (diferente dos ViTs padrão, onde outliers frequentemente aparecem nas camadas finais).

• Registradores de Difusão: Eles introduzem um pequeno conjunto de tokens registradores aprendíveis diretamente no Transformer de Difusão. Esses tokens são treinados conjuntamente com o gerador para absorver sumidouros de atenção e estabilizar o espaço de representação.
• Protocolo de Inferência: Durante a inferência, as saídas de todos os tokens registradores (tanto do codificador quanto do modelo de difusão) são descartadas, retendo apenas as saídas de tokens de imagem.

Contribuições Principais

• Descoberta de Outliers Internos: O artigo demonstra que tokens outlier são pervasivos em DiTs, aparecendo tanto no codificador quanto no denoizer, com um padrão distinto de concentração em camadas intermediárias do modelo de difusão.
• Rejeição da Hipótese de "Valor Extremo": Através de experimentos de mascaramento de perda, os autores mostram que simplesmente excluir tokens de alta norma da perda de treinamento não melhora a qualidade da geração. Isso leva à conclusão de que a questão central é a corrupção das semânticas de patches locais, e não apenas a presença de valores extremos.
• Intervenção de Duplo Estágio: A proposta de DSR fornece um mecanismo simples e leve para estabilizar toda a pipeline de geração, abordando outliers na fonte (codificador) e durante o processo de geração (denoizer).
• Generalizabilidade: O método demonstrou ser eficaz em diferentes espaços de entrada (espaço de pixels, latentes VAE e espaços de representação) e várias arquiteturas (SiT, JiT, designs baseados em RAE).

Resultados Experimentais

Os autores avaliam o DSR na geração condicional por classe ImageNet-1K (256×256) e na geração de imagem a partir de texto em grande escala.

• Desempenho ImageNet-1K:
  • Usando SigLIP2-B como codificador, o DSR reduz o FID de 5,89 para 4,58 (com 800 épocas) e melhora as pontuações GenEval em tarefas de texto-para-imagem de 0,426 para 0,466.
  • Quando combinado com codificadores DINOv2-B, o FID melhora de 4,16 para 3,95.
  • O DSR acelera a convergência, alcançando qualidade comparável com 4× menos épocas de treinamento em comparação com a linha de base.
• Escalabilidade: As melhorias são consistentes em diferentes tamanhos de modelo (DiT-S, B, XL), com apenas um aumento marginal no custo computacional (aumento de aproximadamente 10% em GFLOPs).
• Estudos de Ablação:
  • Profundidade de Inserção: Os registradores são mais eficazes quando introduzidos nas camadas iniciais a médias (por exemplo, começando no bloco 8). Introduzi-los muito cedo ou muito tarde produz retornos decrescentes.
  • Contagem de Registradores: Um número moderado de registradores (por exemplo, 36) produz desempenho ótimo; poucos demais (1) ou muitos demais (100) degradam o desempenho.
• Comparação com Alternativas: O DSR supera o "condicionamento em contexto" (que usa tokens dependentes da entrada) à medida que a representação de entrada se torna mais semântica, sugerindo que os registradores fornecem um mecanismo mais robusto para capacidade global.

Significado e Alegações

O artigo afirma que o controle de tokens outlier é um componente crítico, mas anteriormente negligenciado, na construção de pipelines de difusão baseadas em Transformer robustas. Ao demonstrar que esses artefatos degradam as semânticas de patches locais e que intervenções baseadas em registradores podem mitigá-los efetivamente, os autores posicionam o DSR como uma melhoria fundamental para a estabilidade tanto do codificador quanto do gerador.

O trabalho destaca que a flexibilidade dos Transformers, embora benéfica para interações globais, introduz artefatos específicos (tokens outlier) que prejudicam a fidelidade espacial na geração. A solução proposta é apresentada como uma intervenção simples, plug-and-play, que consistentemente melhora a qualidade da geração em diversos cenários sem exigir mudanças arquiteturais complexas ou re-treinamento de toda a pipeline do zero.