Resumo Técnico: Os Transformers Precisam de Três Projeções?

Declaração do Problema

Desde a introdução da arquitetura Transformer, o mecanismo de atenção Query, Key e Value (QKV) tornou-se o padrão para agregar informações em sequências. No entanto, a necessidade de manter três projeções lineares distintas e independentes para esses componentes permanece uma questão em aberto. Enquanto outras arquiteturas, como CNNs e Modelos de Espaço de Estados (SSMs), frequentemente utilizam representações internas mais unificadas, os Transformers persistem com uma estrutura tripartida. À medida que as janelas de contexto se expandem e a demanda por inferência em tempo real cresce, a comunidade de pesquisa tem focado na eficiência arquitetônica para aliviar o gargalo quadrático da autoatenção. Este artigo investiga se as três projeções são verdadeiramente necessárias ou se representam uma redundância que pode ser reduzida para diminuir a contagem de parâmetros e, mais criticamente, o overhead de memória de inferência sem degradação significativa de desempenho.

Metodologia

Os autores propõem e avaliam sistematicamente três restrições de "Compartilhamento de Projeção" que reduzem progressivamente o número de transformações aprendidas no mecanismo de atenção:

1. Q=K-V (Query-Key Unificada): As projeções de query e key são unificadas ($Q=K$), enquanto a projeção de value permanece separada. Isso resulta em uma matriz de atenção simétrica ($KK^T$). Para abordar as limitações da simetria em tarefas sequenciais, os autores introduzem uma variante $(Q=K-V)^+$ que injeta assimetria via codificações posicionais 2D.
2. Q-K=V (Key-Value Unificada): As projeções de key e value são unificadas ($K=V$), enquanto a query permanece separada. Isso preserva mapas de atenção assimétricos ($QK^T$) e impõe uma forma de amarra de pesos entre chaves e valores.
3. Q=K=V (Projeção Única): Todos os três papéis compartilham uma única matriz de projeção. Isso combina a simetria da primeira variante com o gargalo representacional da segunda.

O estudo também explora a ortogonalidade dessas estratégias de compartilhamento de projeção com técnicas existentes de "Compartilhamento de Cabeças", como Grouped Query Attention (GQA) e Multi-Query Attention (MQA). Os autores hipotetizam que combinar o compartilhamento de projeção (restringindo $K=V$) com o compartilhamento de cabeças (reduzindo o número de cabeças KV) pode gerar ganhos compostos em eficiência de memória.

Experimentos foram conduzidos em 12 tarefas diversas:

• Raciocínio Sintético: Cinco tarefas envolvendo manipulação de listas (Reverse, Sort, Sub, Swap, Copy).
• Visão Computacional: Classificação de imagens em MNIST, FashionMNIST, CIFAR-10, CIFAR-100 e TinyImageNet, bem como detecção de anomalias em conjuntos e segmentação semântica.
• Modelagem de Linguagem: Pré-treinamento de modelos do tipo GPT com 300M e 1.2B de parâmetros em 10 bilhões de tokens do dataset SlimPajama.

Resultados Principais

1. Trocas entre Desempenho e Eficiência (Trade-offs)

• Q-K=V (Equilíbrio Ótimo): Esta variante emergiu como a vencedora clara. Em modelagem de linguagem, ela alcançou uma redução de 50% no tamanho do cache KV com apenas uma degradação de 3,1% na perplexidade (para modelos de 300M) e 2,48% (para modelos de 1.2B). Manteve um desempenho competitivo em tarefas de visão e sintéticas.
• Q=K-V (Limitação de Simetria): Embora esta variante tenha alcançado qualidade de treinamento comparável ao baseline, ela ofereceu zero benefícios de memória de inferência porque ainda requer o cache de tensores K e V separados. Além disso, os mapas de atenção simétricos limitaram sua eficidade em modelagem de linguagem causal, resultando em uma degradação de 4,9% na perplexidade.
• Q=K=V (Super-restrita): Esta simplificação agressiva causou uma perda de desempenho catastrófica (25,4% de degradação de perplexidade em modelagem de linguagem), indicando que forçar uma única projeção para todos os três papéis é excessivamente restritivo para tarefas complexas.

2. Sinergia com Compartilhamento de Cabeças

O estudo demonstra que o compartilhamento de projeção e o compartilhamento de cabeças são estritamente complementares.

• Combinar Q-K=V com GQA-4 (4 grupos) resultou em uma redução de cache de 87,5% com uma degradação de 3,9% na perplexidade.
• Combinar Q-K=V com MQA (1 grupo) alcançou uma redução de cache de 96,9% com uma degradação de 4,8% na perplexidade.
• Na escala de 1.2B, a configuração Q-MQA reduziu o cache KV em 67× (de 5,9 GB para 88 MB por usuário para 32k de contexto) mantendo paridade próxima ao baseline na precisão de tarefas downstream.

3. Comportamento de Escala

O ranking relativo das variantes permaneceu estável conforme o tamanho do modelo aumentava de 300M para 1.2B de parâmetros. Notavelmente, modelos maiores pareceram ligeiramente mais robustos às restrições de projeção, com a degradação de perplexidade para Q-K=V caindo de 3,1% em 300M para 2,48% em 1.2B.

Significância e Alegações

O artigo afirma que o compartilhamento de projeção é uma instância subexplorada de amarra de pesos (weight tying) que oferece benefícios de memória de inferência diretos e quantificáveis, que são particularmente valiosos para implantação em dispositivos de borda (edge) e aplicações de contexto longo.

• Por que Q-K=V Funciona: Os autores argumentam que chaves e valores podem ocupar espaços representacionais semelhantes, e a atenção opera efetivamente em um regime de baixo rank onde a distinção entre K e V é menos crítica do que a assimetria entre Q e a representação compartilhada K-V. A análise empírica mostrou uma alta similaridade de cosseno (0,73) entre as matrizes de projeção padrão de K e V, sugerindo uma redundância inerente.
• Por que Q=K-V Falha: Forçar $Q=K$ cria mapas de atenção simétricos, o que quebra a direcionalidade necessária para dependências causais em modelagem de linguagem.
• Impacto Prático: A principal contribuição não é apenas a redução de parâmetros (que é modesta, ~7-14% de redução no tamanho total do modelo), mas a redução de 50% no cache de memória KV. Isso permite:
  • Janelas de contexto 2× mais longas para o mesmo orçamento de memória.
  • Throughput 2× maior (usuários simultâneos por GPU).
  • Reduções significativas de custo (ex: uma redução projetada de 43% nos custos de serviço para um modelo de autocompletar código).

Os autores concluem que, embora nenhuma arquitetura única vença em todos os domínios, a configuração Q-K=V fornece o melhor equilíbrio entre qualidade e eficiência para tarefas sequenciais, e sua combinação com técnicas de compartilhamento de cabeças permite a inferência prática em dispositivos para modelos de escala de bilhões de parâmetros. O código está publicamente disponível para facilitar pesquisas futuras.