Do All Vision Transformers Need Registers? A Cross-Architectural Reassessment

Resumo Técnico: Reavaliação Cross-Arquitetural dos Registers em Vision Transformers

1. Problema e Motivação

O treinamento de Vision Transformers (ViTs) enfrenta desafios significativos relacionados à interpretabilidade. Especificamente, mapas de atenção e de características (feature maps) de modelos modernos (como DINOv2, DeiT3 e OpenCLIP) frequentemente exibem artefatos visuais.

• Natureza do Problema: Esses artefatos manifestam-se como "tokens de alta norma" (high-norm tokens), que são outliers estatísticos com normas L2 excepcionalmente elevadas.
• Impacto: Eles distorcem a visualização dos mapas de atenção, dificultando a interpretação do que o modelo está "observando" e prejudicando tarefas downstream, como a descoberta de objetos (ex: no framework LOST).
• Hipótese Original (Darcet et al., 2024): Os autores anteriores sugeriram que os ViTs reutilizam patches redundantes da imagem para computação interna e agregação de contexto global. Para resolver isso, propuseram a adição de registers (tokens de entrada vazios) para absorver essa computação, limpando os artefatos.
• Objetivo deste Trabalho: Reproduzir os achados de Darcet et al. (2024), testar a generalização dessas afirmações para uma gama mais ampla de modelos (incluindo menores e arquiteturas hierárquicas) e esclarecer inconsistências terminológicas.

2. Metodologia

Os autores realizaram um estudo de reprodução e extensão rigoroso:

• Modelos Avaliados: DINO, DINOv2 (vários tamanhos), DeiT3, OpenCLIP, Pyramid Vision Transformer (PVT) e Swin Transformer.
• Definições Clarificadas: O trabalho desentrelaça termos usados indiscriminadamente no artigo original:
  • Mapas de Atenção: Scores de atenção entre o token [CLS] e os tokens de entrada.
  • Mapas de Características (Feature Maps): Normas L2 dos tokens de saída.
  • Artefatos vs. Outliers: Nem todo outlier de alta norma gera um artefato visível no mapa de atenção, e vice-versa.
• Experimentos Principais:
  1. Análise de Distribuição: Visualização de mapas de alta resolução e distribuição das normas L2 para identificar outliers (top 2%).
  2. Similaridade Cosine: Verificação se tokens de alta norma residem em regiões de informação visual redundante (alta similaridade com vizinhos).
  3. Tarefas de Avaliação de Informação:
     • Predição de Posição: Capacidade do token de prever sua localização original na imagem.
     • Reconstrução de Patch: Capacidade de reconstruir a imagem a partir do embedding do token.
     • Classificação de Imagem: Uso de tokens normais vs. outliers (substituindo o [CLS]) para classificar a imagem.
  4. Teste de Registers: Avaliação de modelos DINOv2 pré-treinados com e sem registers, e treino de cabeças de classificação lineares utilizando registers.

3. Principais Contribuições e Resultados

A. Refutação e Refinamento das Afirmações Originais:

• Afirmação 1 (Artefatos = Outliers de Alta Norma): Parcialmente Falsa. O trabalho demonstra que artefatos em mapas de atenção nem sempre correspondem a tokens de alta norma em mapas de características. A equivalência depende da arquitetura do modelo.
• Afirmação 2 (Tamanho do Modelo > 300M): Refinada. Embora modelos grandes exibam artefatos consistentemente, modelos menores (<300M) como OpenCLIP-B e DeiT3-M também apresentam outliers. O tamanho não é o único determinante; a arquitetura e o treinamento importam mais.
• Afirmação 3 (Regiões Redundantes): Confirmada para ViTs Clássicos. Em modelos como DINOv2-G, DeiT3-M e OpenCLIP-B, os tokens de alta norma de fato aparecem em regiões de informação redundante. No entanto, em DINOv2-S (sem artefatos) e em arquiteturas hierárquicas, esse padrão não se mantém ou é invertido.
• Afirmação 4 (Informação Local vs. Global): Confirmada para ViTs Clássicos.
  • Tokens de alta norma têm baixa informação local (pior desempenho em reconstrução de patch e predição de posição).
  • Tokens de alta norma têm alta informação global (melhor desempenho em classificação de imagem quando usados como representação global).
  • Exceção: Arquiteturas Hierárquicas (H-ViTs) como Swin e PVT comportam-se de forma oposta ou diferente.

B. Descobertas em Arquiteturas Hierárquicas (H-ViTs):

• Swin Transformer: Os tokens de alta norma concentram-se em informações de aparência fina (textura), enquanto os tokens normais carregam a informação posicional.
• PVT (Pyramid Vision Transformer): Os tokens de alta norma assumem o controle da estrutura posicional grosseira (atingindo 100% de acurácia em predição de posição), enquanto os tokens normais carregam o conteúdo semântico/textura.
• Conclusão: A divisão de trabalho entre tokens "normais" e "outliers" é dependente da arquitetura, não universal.

C. Eficácia dos Registers:

• Limpeza de Artefatos: A adição de registers elimina eficazmente os artefatos nos mapas de atenção e características, confirmando a utilidade prática da proposta original.
• Desempenho: A inclusão de registers mantém ou melhora marginalmente a acurácia de classificação (ex: DINOv2-G com registers: 87.29% vs. 86.90% sem).
• Utilidade para Classificação: Ao tentar usar explicitamente os registers para classificação (concatenando com [CLS] e média de patches), não houve melhoria significativa sobre o uso padrão do [CLS]. Isso sugere que os registers acumulam informação global que já está sobreposta ou redundante com a do token [CLS], limitando seu valor direto para tarefas de classificação, embora sejam vitais para a estabilidade da representação interna.

4. Significado e Conclusões

Este trabalho é fundamental porque:

1. Valida e Generaliza: Confirma a existência de artefatos em ViTs modernos, mas mostra que a solução (registers) e a causa (redundância) não são universais para todas as arquiteturas.
2. Correção Terminológica: Estabelece distinções claras entre "artefatos", "outliers" e "normas altas", evitando generalizações errôneas em futuros estudos.
3. Insight Arquitetural: Revela que mecanismos como janelas desalinhadas (Swin) e redução espacial piramidal (PVT) alteram fundamentalmente como os tokens processam informação, desacoplando a norma do token da redundância local.
4. Implicações Práticas: Recomenda o uso de registers para melhorar a interpretabilidade e a estabilidade de modelos ViTs grandes, mas alerta que a simples adição de registers não é uma "bala de prata" para melhorar a acurácia de classificação além do estado da arte, pois a informação global já é capturada eficientemente pelo token [CLS] em muitos casos.

Em suma, o artigo conclui que, embora os registers sejam uma solução eficaz para limpar artefatos visuais, a origem e o comportamento desses artefatos são mais complexos e dependentes da arquitetura do que originalmente proposto.