Resumo Técnico: Cabeças Laplacianas Melhoram Transformers Suavizando Representações de Tokens

Declaração do Problema

A arquitetura Transformer padrão atualiza representações de tokens por meio de atenção multi-cabeça e conexões residuais. Nessa formulação, cada cabeça de atenção $i$ calcula uma matriz de atenção softmax $P^{(i)}$ e atualiza a entrada $X$ via uma soma ponderada de projeções de valor: $X \leftarrow X + \sum_i P^{(i)}XW_{Vi}W_{oi}$.

Embora eficaz, esse mecanismo carece de controle explícito sobre a variância dentro da sequência das representações de tokens. A visão predominante na literatura, particularmente em relação a Redes Neurais de Grafos (GNNs) e Transformers profundos, trata o "excesso de suavização" — onde as representações de tokens tornam-se muito semelhantes, perdendo distintividade — como um modo de falha prejudicial. Consequentemente, a maioria dos trabalhos anteriores foca em prevenir a suavização por meio de normalização, descarte de arestas ou funções de perda de diversidade. Este artigo desafia a suposição de que o excesso de suavização de tokens é inerentemente prejudicial, propondo, em vez disso, que a suavização controlada pode ser benéfica para a aprendizagem de representações.

Metodologia

Os autores propõem uma modificação arquitetônica sem parâmetros: substituir um subconjunto de cabeças de atenção padrão por cabeças Laplacianas.

O Operador Laplaciano

Na atualização padrão, uma cabeça calcula $P^{(i)}X$, que é uma média ponderada por atenção dos tokens. A cabeça Laplaciana proposta substitui a matriz de atenção $P^{(i)}$ pelo operador Laplaciano do grafo $I - P^{(i)}$, onde $I$ é a matriz identidade. A equação de atualização torna-se:
$$X \leftarrow X + \sum_{i \in A} P^{(i)}XW_{Vi}W_{oi} + \sum_{i \in L} (I - P^{(i)})XW_{Vi}W_{oi}$$
onde $A$ é o conjunto de cabeças de atenção padrão e $L$ é o conjunto de cabeças Laplacianas.

Motivações

O artigo fornece duas motivações teóricas para essa modificação:

1. Controle Direto da Média e Variância: Cabeças de atenção padrão atualizam a média das representações de tokens. Em contraste, o termo $(I - P^{(i)})X$ calcula a diferença entre cada token e sua média ponderada por atenção. Adicionar essa diferença permite que o modelo controle diretamente a variância dentro da sequência das representações de tokens, uma capacidade ausente em Transformers padrão.
2. Difusão em Grafos: Se os tokens forem vistos como nós em um grafo com pesos de aresta definidos por $P^{(i)}$, então $I - P^{(i)}$ é o Laplaciano do grafo normalizado por passeio aleatório. Aplicar esse operador corresponde a um único passo de difusão de calor no grafo de tokens, suavizando progressivamente o sinal (representações de tokens) ao longo da sequência.

Principais Contribuições

1. Proposta Arquitetônica: Uma modificação simples e sem parâmetros à arquitetura Transformer que integra cabeças Laplacianas ao lado de cabeças de atenção padrão.
2. Validação Empírica: Demonstração de que essa modificação melhora o desempenho em três paradigmas de aprendizagem distintos: classificação supervisionada, modelagem de linguagem e aprendizagem auto-supervisionada.
3. Caracterização Geométrica: Uma análise detalhada de como as cabeças Laplacianas alteram a geometria das representações de tokens aprendidas, mostrando que elas induzem formas benéficas de "colapso" e suavização em vez de excesso de suavização prejudicial.

Resultados Experimentais

1. Classificação de Imagens (Aprendizagem Supervisionada)

Usando a família DeiT-3 (ViT-B) no CIFAR-10, CIFAR-100 e ImageNet-1k:

• Desempenho: Adicionar cabeças Laplacianas melhorou consistentemente a precisão Top-1. A melhoria foi quase monotônica com o número de cabeças Laplacianas.
  • No ImageNet, o modelo com 12 cabeças Laplacianas (ViT-B-12L) alcançou 82,32% de precisão, comparado aos 81,45% da linha de base.
  • No CIFAR-100, o ganho foi de quase 2% (77,04% vs. 75,12%).
• Geometria:
  • ANOVA: Cabeças Laplacianas reduziram significativamente a Variância Dentro da Sequência enquanto aumentavam a Variância Entre Classes. Isso indica que tokens dentro da mesma sequência colapsam em direção a uma média comum, enquanto as médias de diferentes sequências se separam mais efetivamente.
  • Colapso Neural: Os modelos exibiram uma convergência mais forte para a geometria de Colapso Neural. As representações de tokens formaram clusters bem separados, semelhantes a simplexos (Simplex ETF) para cada classe, com médias de sequência alinhadas à geometria global da classe.

2. Modelagem de Linguagem (Autoregressiva)

Usando modelos no estilo GPT-2 (561M parâmetros) pré-treinados no FineWeb-Edu e ajustados finamente em benchmarks de raciocínio/código:

• Desempenho: Modelos com cabeças Laplacianas superaram a linha de base em média em ARC-Easy, ARC-Challenge, MMLU, GSM8K e HumanEval.
  • O modelo com 5 cabeças Laplacianas alcançou o melhor desempenho médio (30,06%), com ganhos significativos no ARC-Easy (+16%) e ARC-Challenge (+7%).
• Geometria:
  • Decaimento Espectral: Cabeças Laplacianas produziram representações de tokens com espectros de decaimento mais rápidos. O número de componentes principais necessários para explicar 99% da variância caiu 21% (de 547,7 para 430,4) quando todas as cabeças eram Laplacianas.
  • Colapso Linguístico: Os modelos mostraram maior separabilidade entre representações de tokens que compartilham a mesma previsão do próximo token, sugerindo uma forma mais forte de "colapso linguístico" onde tokens com papéis semânticos semelhantes agrupam-se juntos.

3. Aprendizagem Auto-Supervisionada (DINO)

Usando modelos ViT-S no ImageNet-1k para pré-treinamento DINO:

• Desempenho: Cabeças Laplacianas melhoraram a precisão de sondagem linear no ImageNet e o desempenho de segmentação no ADE20K.
  • O ViT-S-5L alcançou o melhor mIoU no ADE20K (32,27%) comparado à linha de base (29,81%).
• Geometria:
  • Segmentação: Visualizações de PCA mostraram que cabeças Laplacianas produziram limites de objetos mais nítidos e preservaram detalhes mais finos (por exemplo, formas de capacetes, números de camisas) em comparação com a linha de base.
  • Decaimento Espectral: Similar à modelagem de linguagem, esses modelos exibiram decaimento espectral mais rápido, concentrando representações de tokens dentro da imagem em um subespaço de dimensão inferior.

Significado e Alegações

O artigo argumenta que a visão predominante — de que o excesso de suavização de tokens é universalmente prejudicial — é incompleta. Em vez disso, os autores afirmam que a suavização é um mecanismo para aprendizagem de representações que pode ser benéfica quando controlada.

• Benefícios Específicos da Tarefa: A "suavização" induzida por cabeças Laplacianas não é um modo de falha, mas um recurso que se alinha aos objetivos da tarefa:
  • Em classificação supervisionada, facilita o Colapso Neural ao colapsar tokens dentro de sequências e alinhar médias de sequências com a geometria da classe.
  • Em modelagem de linguagem, aprimora o Colapso Linguístico, melhorando a separabilidade de tokens com alvos de próximo token idênticos.
  • Em aprendizagem auto-supervisionada, concentra representações em subespaços de dimensão inferior, auxiliando tarefas subsequentes como segmentação.
• Mecanismo: Ao incorporar explicitamente a difusão em grafos (via o Laplaciano) na arquitetura, o modelo ganha a capacidade de controlar diretamente a variância dentro da sequência, levando a representações mais robustas e geometricamente estruturadas.

Os autores concluem que, embora sua implementação de estratégia fixa seja provavelmente subótima, os resultados sugerem que a suavização deve ser investigada mais a fundo como um componente fundamental da aprendizagem de representações, em vez de apenas como um problema a ser mitigado.