Key and Value Weights Are Probably All You Need: On the Necessity of the Query, Key, Value weight Triplet in Self-Attention Transformers

Este artigo demonstra teoricamente e valida empiricamente que as matrizes de pesos de Query, Key ou Value em Transformers são redundantes e podem ser substituídas por matrizes identidade, reduzindo os parâmetros em 25% e simplificando a otimização sem comprometer o desempenho.

O essencial

Imagine que você está construindo um restaurante de luxo (o modelo de IA) onde os chefs (os dados) precisam preparar pratos incríveis (respostas inteligentes). Até hoje, a receita padrão exigia que cada chef tivesse três ferramentas específicas na mão para cada tarefa: um Filtro (Query), uma Lista de Ingredientes (Key) e uma Prateleira de Temperos (Value).

Os autores deste artigo, Marko e Antonij, olharam para essa cozinha e disseram: "E se o Filtro for, na verdade, apenas uma janela aberta? Por que gastar tempo e dinheiro fabricando um Filtro complexo se podemos simplesmente deixar a janela aberta e usar a luz natural?"

Aqui está a explicação do que eles descobriram, usando analogias do dia a dia:

1. A Grande Descoberta: "A Janela Aberta"

No coração dos modelos de IA (chamados Transformers), existe um mecanismo chamado "Atenção". Ele decide quais partes da frase são importantes. Normalmente, ele usa três pesos matemáticos (matrizes) para fazer isso: Query (Q), Key (K) e Value (V).

• A Analogia: Pense na Query como uma lupa que você usa para procurar algo. A Key é o que você está procurando, e o Value é o objeto que você encontra.
• O Problema: A lupa (Query) é feita de um material especial e pesada. Ela custa caro para treinar e usar.
• A Solução: Os autores provaram matematicamente que, na maioria das vezes, você não precisa de uma lupa personalizada. Você pode simplesmente abrir a janela (usar uma "Matriz Identidade", que é basicamente "não fazer nada"). A luz do sol (os dados de entrada) entra diretamente.

Ao fazer essa troca, eles conseguiram eliminar 25% dos parâmetros (peças do quebra-cabeça) da parte de atenção do modelo. É como se você removesse um terço dos utensílios da cozinha e o prato saísse exatamente igual, ou até melhor!

2. Por que isso funciona? (A Magia da "Troca de Base")

Você pode estar pensando: "Mas se eu tirar a lupa, como o chef vai saber o que procurar?"

Aqui entra a parte mágica da matemática. O modelo é inteligente o suficiente para se adaptar.

• A Analogia: Imagine que você tem um mapa (os dados) e uma bússola (a Query). Se você tirar a bússola, o mapa continua lá. O que acontece é que o chef (o resto do modelo) aprende a ler o mapa de um ângulo diferente.
• Em vez de ter uma bússola que aponta para o Norte, o chef aprende a olhar para o mapa e dizer: "Ok, se eu olhar para o canto superior esquerdo, é como se fosse o Norte".
• O modelo "reorganiza" o resto dos pesos (Key e Value) para compensar a falta da Query. É como se você trocasse a roupa do chef, mas ele continuasse cozinhando o mesmo prato saboroso.

3. O Resultado na Prática: Mais Rápido e Mais Estável

Os autores não ficaram só na teoria. Eles construíram modelos do zero (como o GPT) e testaram essa ideia.

• Economia: O modelo ficou menor e mais leve. Menos "peso" para carregar significa que ele é mais rápido e consome menos energia.
• Estabilidade: O mais interessante é que, ao tirar a Query, o modelo ficou mais fácil de treinar.
  • A Analogia: Imagine tentar equilibrar uma torre de pratos. Com a Query, você precisa de um equilíbrio muito delicado (muitos ajustes finos). Sem a Query, a torre se torna mais estável, como se tivesse uma base mais larga. Eles conseguiram treinar o modelo usando 3 vezes menos "freio" (um ajuste matemático chamado weight decay) e ainda assim ele não caiu.

4. O Que Aprendemos Sobre a "Inteligência"

A descoberta mais profunda é que a arquitetura atual dos modelos de IA pode ter um pouco de "gordura" desnecessária.

• A Analogia: É como se tivéssemos estado usando um carro com um motor V8 gigante para ir ao mercado, quando um motor 1.6 faria o mesmo trabalho, gastaria menos gasolina e seria mais fácil de dirigir.
• O papel sugere que a "Query" não é tão essencial quanto pensávamos. A verdadeira inteligência e capacidade de expressão do modelo vêm mais do "cérebro" (a parte que processa a informação, o MLP) do que da "lupa" (a Query).

Resumo para Levar para Casa

1. O Problema: Modelos de IA têm uma peça chamada "Query" que é cara e complexa.
2. A Solução: Eles provaram que podem substituir essa peça por uma "janela aberta" (fazer nada) sem perder qualidade.
3. O Benefício: Modelos menores, mais rápidos, mais baratos e que são mais fáceis de treinar.
4. A Lição: Às vezes, menos é mais. A simplicidade pode ser a chave para a próxima geração de inteligências artificiais.

Em suma, eles mostraram que podemos simplificar drasticamente a "caixa de ferramentas" da IA sem que ela perca sua capacidade de ser inteligente, abrindo caminho para modelos mais eficientes e acessíveis no futuro.

Resumo técnico

Resumo Técnico: Redundância dos Pesos de Query em Transformers

1. O Problema

O treinamento e a implantação de modelos de linguagem baseados em Transformers são computacionalmente caros. A arquitetura padrão utiliza um triplo de matrizes de pesos para o mecanismo de auto-atenção: Query (Q), Key (K) e Value (V). O artigo investiga uma questão fundamental: é estritamente necessário manter os três conjuntos de pesos independentes? Especificamente, os autores questionam se a matriz de pesos de Query ($W_Q$) pode ser eliminada ou substituída sem perda de capacidade expressiva do modelo, visando reduzir parâmetros e complexidade computacional.

2. Metodologia e Abordagem Teórica

Os autores adotam uma abordagem "teoria primeiro", provando a redundância sob certas suposições e validando empiricamente em arquiteturas complexas.

• Lema de Reparametrização: O ponto central da análise é a observação de que a atenção depende da entrada $X$ apenas através dos produtos $XW_Q$, $XW_K$ e $XW_V$. Isso permite uma construção telescópica onde a transformação de base de uma camada pode ser absorvida pela próxima.
• Eliminação de Pesos: Eles provam que, sob condições específicas (como ausência de camadas de normalização ou uso de conexões de salto específicas), a matriz $W_Q$ pode ser definida como a matriz identidade ($I$) em todas as camadas, desde que as matrizes $W_K$ e $W_V$ sejam ajustadas (reparametrizadas) para compensar a mudança de base.
  • Caso de Camada Única: Em qualquer transformador sem normalização, a $W_Q$ de uma única camada pode ser eliminada sem modificações arquitetônicas.
  • Múltiplas Camadas: Para eliminar $W_Q$ em todas as camadas simultaneamente, o modelo deve satisfazer uma das seguintes condições:
    1. Conexões de salto (skip connections) apenas ao redor do bloco de atenção (e não ao redor do MLP).
    2. Compartilhamento de pesos entre todas as camadas (como no ALBERT).
• Normalização (LayerNorm): O artigo analisa as obstruções introduzidas pela LayerNorm. Eles derivam condições suficientes para que as transformações de base comutem com a normalização, sugerindo que, embora a eliminação exata seja teoricamente complexa com normalização, uma equivalência aproximada é viável com ajustes de hiperparâmetros.
• Análise de MLP e Conexões de Salto: Um resultado fundamental (Teorema 8.4) caracteriza quando conexões de salto podem ser absorvidas por MLPs com ativação ReLU. Eles provam que, em geral, MLPs com e sem conexões de salto pertencem a classes de funções disjuntas, a menos que condições algébricas muito específicas sejam atendidas. Isso motiva a necessidade de manter a expressividade do MLP quando se remove a complexidade da atenção.

3. Contribuições Principais

1. Prova Teórica de Redundância: Demonstração de que os pesos de Query são redundantes em mecanismos de atenção multi-head, permitindo sua substituição pela identidade, reduzindo os parâmetros de atenção em 25% por camada (cerca de 8% dos parâmetros totais do bloco do transformador).
2. Novas Limites de Expressividade: Caracterização precisa de quando conexões de salto podem ser absorvidas por MLPs ReLU, revelando que a eliminação de $W_Q$ empurra o modelo para uma classe de funções onde o MLP deve compensar a perda de expressividade.
3. Validação Empírica: Treinamento de modelos do tipo GPT (decoder-only) do zero com $W_Q = I$, demonstrando que a eliminação é viável na prática.
4. Ajustes Práticos: Identificação de duas modificações necessárias para a estabilidade do treinamento:
   • Correção de Escala de Atenção: Ajuste do fator de escala para $\frac{1}{2\sqrt{d_k}}$ (em vez de $\frac{1}{\sqrt{d_k}}$) para compensar a maior variância inicial das pontuações de atenção quando $W_Q = I$.
   • Redução de Weight Decay: O modelo reduzido permanece estável com um weight decay 3x menor, sugerindo que a eliminação de $W_Q$ atua como uma regularização implícita.

4. Resultados Experimentais

Os autores treinaram modelos GPT-style (117M a 124M parâmetros) no conjunto de dados OpenWebText:

• Desempenho Paritário: O modelo reduzido (117M parâmetros, sem $W_Q$) alcançou uma perda de validação comparável ao modelo baseline completo (124M parâmetros), apesar de ter 8% menos parâmetros não-embedding.
• Melhoria com Realocação: Quando os parâmetros economizados pela eliminação de $W_Q$ foram realocados para aumentar a dimensão oculta do MLP (criando um modelo de 124M com MLP maior), o modelo reduzido superou o baseline completo, alcançando a melhor perda de validação (3.004 vs 3.016).
• Estabilidade: O treinamento foi estável e convergiu suavemente com um weight decay significativamente reduzido, indicando que a arquitetura simplificada possui uma estabilidade inerente.

5. Significado e Implicações

• Eficiência Arquitetônica: A descoberta de que $W_Q$ é redundante oferece uma rota direta para reduzir a memória e a computação em modelos de linguagem grandes (LLMs), especialmente em cenários de inferência com cache KV (Key-Value), onde a eliminação de $W_Q$ simplifica ainda mais o processo.
• Regularização Implícita: A estabilidade do modelo com weight decay reduzido sugere que a estrutura de identidade em $W_Q$ impõe uma regularização benéfica, evitando a divergência de perda comum em modelos grandes.
• Reavaliação de Design: O trabalho desafia a suposição de que todos os componentes do triplo Q-K-V são essenciais, sugerindo que a complexidade atual pode ser, em parte, um artefato histórico de design em vez de uma necessidade matemática para expressividade.
• Futuro: Os resultados motivam a investigação da eliminação de outros pesos (como $W_K$ ou $W_V$) e a aplicação dessas simplificações em diferentes modalidades e escalas maiores.

Em suma, o artigo demonstra que é possível simplificar fundamentalmente o mecanismo de atenção de Transformers eliminando os pesos de Query, mantendo ou até melhorando o desempenho através de ajustes de hiperparâmetros e realocação inteligente de parâmetros para as camadas MLP.