Resumo Técnico: Registradores Importam para Transformers de Difusão no Espaço de Pixels

Declaração do Problema

Os Vision Transformers (ViTs) são conhecidos por exibir outliers de tokens de patch com alta norma que degradam a qualidade dos mapas de características e criam artefatos nos mapas de atenção. Essa questão é efetivamente mitigada em ViTs pela introdução de "tokens registradores", que atuam como sumidouros para esses outliers. À medida que os Diffusion Transformers (DiTs) adotam cada vez mais arquiteturas de transformer e migram para o treinamento diretamente no espaço de pixels (afastando-se do espaço latente), eles se assemelham estruturalmente mais aos ViTs. Isso levanta duas questões críticas:

1. Os DiTs herdam os outliers de tokens de patch com alta norma observados em ViTs?
2. Os tokens registradores são benéficos para DiTs, particularmente em regimes de treinamento no espaço de pixels?

Enquanto sumidouros de atenção e comportamentos de tokens especiais foram estudados em Grandes Modelos de Linguagem (LLMs) e modelos de linguagem de difusão, sua presença e papel funcional em transformers de difusão de imagem permanecem pouco explorados.

Metodologia

Os autores realizam uma análise sistemática de tokens registradores em diferentes espaços de modelos de difusão (espaço de pixels, espaço VAE e espaço RAE) e os comparam com ViTs padrão (especificamente DINOv2).

1. Análise Empírica de Outliers: O estudo investiga as normas de características token a token e os mapas de atenção em DiTs (especificamente pDiTs treinados com fluxo de correspondência no ImageNet) versus ViTs. Eles analisam se os tokens de patch exibem outliers de alta norma na ausência de registradores.
2. Integração de Tokens Registradores: Tokens registradores são implementados como tokens adicionais aprendíveis anexados à sequência de tokens de patch, seguindo a abordagem usada em ViTs. Esses tokens são excluídos do objetivo de treinamento (função de perda).
3. Análise de Representações Internas: Os autores analisam mapas de características intermediários usando Variação Total (TV) para medir a suavidade espacial e sondagem linear para avaliar o conteúdo semântico dos tokens registradores. Eles examinam como os registradores afetam as normas de características e a coerência das representações em diferentes passos de difusão (níveis de ruído).
4. Exploração Arquitetural: Motivados pelos papéis distintos de registradores e tokens de patch, o artigo propõe uma arquitetura de fluxo duplo eficiente em parâmetros. Este design processa tokens registradores e de patch através de caminhos especializados (especificamente desacoplando camadas RMSNorm, MLP e adaLN), enquanto compartilha mecanismos de atenção e a maioria dos parâmetros do transformer.
5. Estudos de Ablação: O estudo abla o número de registradores, a profundidade na qual são introduzidos e os componentes específicos a serem dualizados para determinar configurações ótimas.

Principais Descobertas e Contribuições

1. Os DiTs Diferem Fundamentalmente dos ViTs em Relação aos Outliers

Ao contrário dos ViTs, os DiTs padrão (tanto latentes quanto no espaço de pixels) não exibem outliers de tokens de patch com alta norma. Em ViTs, os outliers estão localizados em tokens específicos (frequentemente em regiões de baixa informação), enquanto em DiTs, as normas dos tokens de patch permanecem relativamente uniformes, e os mapas de atenção carecem dos artefatos de baixa informação vistos em ViTs.

2. Registradores Melhoram DiTs no Espaço de Pixels Apesar da Ausência de Outliers Preexistentes

Surpreendentemente, adicionar tokens registradores a DiTs melhora significativamente a convergência e a qualidade de geração, particularmente em modelos no espaço de pixels.

• Mecanismo: Ao contrário dos ViTs, onde os registradores absorvem outliers preexistentes, os registradores em DiTs criam tokens de alta norma dentro deles mesmos.
• Efeito: A presença de registradores reduz consistentemente as normas de características de todos os tokens de patch. Isso leva a mapas de características intermediários mais limpos e suaves, especialmente em altos níveis de ruído (passos iniciais de difusão).
• Função Dupla: Os tokens registradores se especializam em dois papéis distintos:
  • Sumidouros de Norma: Tokens com normas extremamente altas que absorvem magnitude dos tokens de patch.
  • Portadores de Informação Global: Tokens com normas moderadas que codificam informações semânticas diversas sobre a imagem (por exemplo, objetos em primeiro plano, contexto de fundo), verificados via sondagem linear e visualização de mapas de atenção.

3. Eficácia Específica do Espaço

Os benefícios dos tokens registradores dependem altamente do espaço de treinamento:

• Espaço de Pixels: Registradores produzem as melhorias mais significativas na qualidade de geração (FID). Modelos no espaço de pixels produzem inerentemente representações intermediárias mais ruidosas com normas de características mais altas, tornando o efeito de regularização dos registradores crucial.
• Espaço Latente (VAE/RAE): Registradores fornecem ganhos moderados ou, no caso de modelos baseados em RAE, degradam ligeiramente o desempenho. Espaços latentes são mais estruturados e de menor dimensão, reduzindo a necessidade da forte regularização fornecida por registradores.

4. Registradores Implícitos em Arquiteturas Existentes

Os autores observam que arquiteturas recentes de DiT no espaço de pixels (por exemplo, JiT) que usam condicionamento de classe em contexto (anexando embeddings de classe à sequência de tokens) incorporam implicitamente mecanismos semelhantes a registradores. Esses tokens em contexto exibem comportamento semelhante a registradores explícitos: alguns atuam como sumidouros de norma, enquanto outros codificam informações semânticas globais. Isso sugere que os ganhos de desempenho em tais modelos podem decorrer em grande parte desse comportamento semelhante a registradores, em vez da própria informação de classe.

5. Arquitetura de Fluxo Duplo Eficiente em Parâmetros

Para aproveitar essas percepções, os autores propõem uma arquitetura compacta de fluxo duplo que especializa o processamento de tokens registradores.

• Design: Registradores são introduzidos apenas em camadas mais profundas (começando do bloco 4), onde a estrutura semântica é estabelecida. A arquitetura usa atenção compartilhada, mas desacopla RMSNorm, MLP e adaLN para registradores e patches.
• Eficiência: Esta abordagem aumenta o tamanho do modelo em apenas ~14% (comparado a ~77% para duplicação completa ingênua) enquanto mantém custos de inferência semelhantes (GFLOPs).
• Desempenho: Esta arquitetura melhora consistentemente a qualidade de geração em todas as escalas de modelo (B, L, H) e resoluções (256x256, 512x512) em comparação com bases de linha de único fluxo.

Resultados

• Qualidade de Geração: No ImageNet 256x256, pDiTs com registradores superam consistentemente aqueles sem. Por exemplo, um modelo pDiT-B/16 melhorou seu FID de 7,39 para 5,30 em 200 épocas.
• Suavidade de Características: Modelos com registradores mostram razões de Variação Total (TV) mais baixas em características intermediárias em altos níveis de ruído, indicando representações mais suaves e coerentes.
• Ablação: Registradores são mais eficazes quando introduzidos após as primeiras 4 camadas e quando se usa um número maior de tokens (por exemplo, 32 registradores para pDiT-B, comparado a 4 para ViTs).

Significado e Alegações

O artigo alega que tokens registradores são um componente crítico, mas anteriormente negligenciado, para geração de difusão de alta qualidade no espaço de pixels. O significado reside em:

1. Redefinir o Papel dos Registradores: Demonstrar que registradores não são meramente uma correção para outliers específicos de ViT, mas servem uma função de regularização mais ampla em modelos de difusão, gerenciando normas de características e aprimorando a coerência espacial.
2. Explicar Sucessos Recentes: Fornecer uma racionalização para o forte desempenho empírico de DiTs recentes no espaço de pixels que utilizam condicionamento em contexto, atribuindo parte de seu sucesso ao comportamento implícito de registradores.
3. Design de Arquitetura Eficiente: Oferecer um design de fluxo duplo eficiente em parâmetros que capitaliza os papéis distintos de registradores e patches, alcançando qualidade de geração state-of-the-art com sobrecarga de tempo de execução negligenciável.

Os autores concluem que, embora seus designs capturem a promessa da especialização de tokens, o potencial pleno de utilizar registradores em transformers de difusão no espaço de pixels permanece uma questão aberta para pesquisas futuras.