🤖 machine learning

Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention

O artigo propõe o Grouped Query Experts (GQE), uma camada de mistura de especialistas aplicada à atenção de consultas agrupadas que seleciona dinamicamente um subconjunto de especialistas de cabeça de consulta por token enquanto mantém as cabeças de chave-valor densas, reduzindo assim a computação ativa e melhorando a eficiência sem sacrificar a precisão em tarefas subsequentes.

Autores originais: Vishesh Tripathi, Abhay Kumar

Publicado 2026-06-23

📖 4 min de leitura☕ Leitura rápida

CC0 1.0

Autores originais: Vishesh Tripathi, Abhay Kumar

Artigo original dedicado ao domínio público sob CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine uma biblioteca imensa (um modelo de IA Transformer) onde cada livro (um token de texto) precisa ser cruzado com todos os outros livros para entender a história. É assim que funciona a "autoatenção" (self-attention).

O Problema: O Bibliotecário "Tamanho Único"
Nos modelos de IA padrão, há uma equipe de 16 bibliotecários especializados (chamados de "cabeças de atenção" ou attention heads). Não importa qual livro você tire da prateleira, todos os 16 bibliotecários devem lê-lo, analisá-lo e escrever um relatório.

Se o livro for uma palavra simples como "o" ou uma vírgula, você não precisa de 16 especialistas; um ou dois bastariam.
Se o livro for um termo científico complexo, você pode precisar de todos os 16.
Mas, atualmente, o sistema força todos os 16 a trabalhar em cada única palavra. À medida que a história fica mais longa (contexto longo), isso se torna incrivelmente lento e caro, como contratar uma orquestra inteira para tocar uma única nota.

A Solução Existente: Atenção de Consulta Agrupada (GQA - Grouped Query Attention)
Antes deste artigo, pesquisadores tentaram economizar dinheiro agrupando os bibliotecários. Em vez de 16 equipes únicas, havia 8 equipes onde dois bibliotecários compartilhavam as mesmas notas de "Chave e Valor" (o material de referência). Isso economizou memória, mas todos os 16 bibliotecários ainda tinham que ler cada palavra. Era como ter um arquivo menor, mas ainda fazer com que cada funcionário caminhasse por todo o prédio para realizar seu trabalho.

A Nova Solução: Especialistas de Consulta Agrupados (GQE - Grouped Query Experts)
Os autores propõem um sistema mais inteligente chamado Especialistas de Consulta Agrupados (GQE). Pense nisso como transformar esses 16 bibliotecários em um sistema de "Mistura de Especialistas" (Mixture of Experts), mas com um toque especial.

A Configuração: Eles mantêm os 8 grupos de notas de referência (as cabeças KV) exatamente iguais. Esta parte é sempre "densa" (totalmente ativa) porque o material de referência é barato de acessar.
O Roteador: Dentro de cada grupo, existem múltiplos bibliotecários "Especialistas" (cabeças de consulta/query heads). Para cada palavra, um "Roteador" inteligente (um guarda de trânsito) olha para a palavra e pergunta: "Nós realmente precisamos que todos os 4 especialistas deste grupo leiam isso?"
A Seleção: O Roteador escolhe apenas os 1 ou 2 melhores especialistas (k=1 ou k=2) para realizar o trabalho real para aquela palavra específica. Os outros especialistas do grupo tiram um intervalo para o café.
A Rede de Segurança: Para garantir que o sistema não fique confuso ou preguiçoso, eles adicionam dois recursos especiais:
- A Cabeça Compartilhada: Um bibliotecário está sempre de plantão, lendo cada palavra, apenas para manter a equipe estável.
- O Resumo Ponderado: O sistema cria um "relatório de consenso" que combina o trabalho dos especialistas selecionados. Isso é crucial porque dá ao Roteador um sinal claro de quão bem ele fez o seu trabalho, permitindo que ele aprenda qual especialista é o melhor para cada palavra.

Os Resultados: Mais Rápido Sem Perder a Inteligência
O artigo testou isso em um modelo pequeno (250 milhões de parâmetros) treinado em 30 bilhões de palavras.

Precisão: O modelo GQE teve um desempenho tão bom quanto o modelo antigo que usava todos os 16 bibliotecários para cada palavra. Ele não ficou "mais burro" ao pular pessoas.
Velocidade: Como ele pulou cerca de metade do trabalho das cabeças de consulta, tornou-se mais rápido.
- Para histórias curtas, foi ligeiramente mais rápido (1.15x).
- Para histórias muito longas (como um romance inteiro), foi 1.7 a 1.8 vezes mais rápido.

Por Que Isso Importa
O artigo afirma que, ao tornar a parte da "leitura" desta IA condicional (fazendo o trabalho apenas quando necessário) enquanto mantém a parte de "referência" constante, você pode acelerar significamente conversas longas ou análise de documentos sem sacrificar a qualidade das respostas.

A Ressalva (Limitações)
Os autores são cuidadosos ao notar que isso foi testado em um modelo relativamente pequeno. Eles ainda não provaram que funciona em modelos massivos de trilhões de parâmetros. Além disso, o "Roteador" precisa ser treinado com muito cuidado; se você apenas escolher especialistas aleatoriamente sem a "Rede de Segurança" (a cabeça compartilhada e o resumo ponderado), o modelo na verdade piora.

Em resumo: O GQE é como contratar uma equipe de especialistas onde apenas os certos aparecem para a tarefa específica em questão, tornando todo o processo muito mais rápido para tarefas longas, sem perder a qualidade do trabalho.

Resumo Técnico: Grouped Query Experts (GQE)

Definição do Problema

A autoatenção é o gargalo computacional nas arquiteturas Transformer, particularmente para sequências de contexto longo, onde as interações par a par entre tokens escalam quadraticamente. Os mecanismos de atenção densa padrão aplicam o mesmo conjunto de cabeças de atenção a cada token, independentemente do conteúdo de informação ou da dificuldade do token. Essa ativação uniforme leva a um uso ineficiente de computação, pois tokens de baixa informação (ex: pontuação, palavras de parada/stop words) consomem os mesmos recursos que palavras portadoras de conteúdo. Embora o Mixture-of-Experts (MoE) tenha introduzido com sucesso a computação condicional nos blocos de feed-forward (MLP) do Transformer, aplicar uma esparsidade semelhante ao mecanismo de atenção permanece um desafio em aberto. O Grouped-Query Attention (GQA) existente reduz a largura de banda de memória e o tamanho do KV-cache ao compartilhar cabeças de Key-Value (KV) entre grupos de cabeças de query, mas ainda avalia todas as cabeças de query para cada token.

Metodologia: Grouped Query Experts (GQE)

Os autores propõem o Grouped Query Experts (GQE), um método que integra os princípios de MoE ao bloco de atenção de modelos GQA. A filosofia central de design é tornar a computação das cabeças de query condicional, preservando os benefícios de cache KV denso do GQA.

Arquitetura e Roteamento

Pool de Especialistas Agrupados (Grouped Expert Pool): Dentro de cada grupo fixo de GQA (onde múltiplas cabeças de query compartilham uma única cabeça KV), as cabeças de query são tratadas como um pool de "especialistas".
Ativação Condicional: Para cada token, um roteador seleciona um subconjunto esparso de $k$ especialistas de cabeças de query de cada grupo de GQA.
Caminho KV Denso: Crucialmente, as projeções de Key-Value e o KV-cache permanecem densos e inalterados. A esparsidade é aplicada apenas à computação do lado da query.
Construção da Saída: A saída é construída através de um mecanismo específico para garantir estabilidade no treinamento:
- Concatenação Hard (Hard Concatenation): As saídas dos $k$ especialistas selecionados por grupo são concatenadas.
- Soma Ponderada Renormalizada: Uma soma ponderada das saídas dos especialistas selecionados é computada usando as probabilidades do roteador. Isso fornece um caminho diferenciável para que a perda de modelagem de linguagem alcance o roteador, resolvendo a não-diferenciabilidade da seleção top- $k$ discreta.
- Cabeça Compartilhada Sempre Ativa (Always-On Shared Head): Uma cabeça de atenção adicional é computada para cada token, independentemente do roteamento. Isso atua como uma âncora estável durante o treinamento, evitando que o roteador colapse para um único especialista e garantindo que um padrão de atenção base esteja sempre disponível.

Treinamento e Perda

Além da perda padrão de modelagem de linguagem, o GQE emprega uma perda auxiliar de balanceamento de carga para encorajar o roteador a distribuir os tokens uniformemente pelo pool de especialistas dentro de cada grupo, evitando o "colapso do especialista".

Principais Contribuições

Formulação de MoE em GQA: O artigo formula o MoE especificamente dentro dos grupos de GQA, ativando $k$ especialistas por token enquanto mantém todas as cabeças KV densas. Isso alinha a comparação com a configuração subjacente do GQA e preserva o perfil de memória do GQA.
Redução de Computação sem Perda de Qualidade: Os autores demonstram que o GQE pode igualar a precisão downstream de uma linha de base GQA totalmente ativa, enquanto ativa apenas um subconjunto esparso de especialistas de cabeças de query.
Mecanismos de Estabilização: O artigo identifica que o roteamento esparso bem-sucedido na atenção requer dois componentes específicos: um slot de soma ponderada renormalizada para fornecer um sinal de aprendizado adequado ao roteador, e uma cabeça compartilhada sempre ativa para ancorar a camada.

Resultados Experimentais

Os experimentos foram conduzidos com um orçamento fixo de 30 bilhões de tokens na escala de 250 milhões de parâmetros utilizando o dataset FineWeb-Edu.

Precisão: Na configuração principal (16 cabeças de query, 8 grupos KV, $k=1$ ), o GQE ativa 8 especialistas roteados mais 1 cabeça compartilhada, totalizando 9 computações de query-attention de 16. Apesar dessa redução (ativando ~56% do total de operações de query), o GQE iguala a precisão da linha de base GQA totalmente ativa (HellaSwag, PIQA, ARC-Easy), atingindo uma pontuação média de 56,04 comparado aos 55,86 da linha de base.
Estudos de Ablação: O artigo mostra que remover o slot de soma ponderada renormalizada ou a cabeça compartilhada resulta em degradação de desempenho (caindo para ~55,18 e 55,43 respectivamente), confirmando que a esparsidade sozinha é insuficiente sem os estabilizadores arquiteturais específicos.
Throughput e Aceleração (Speedup): O GQE demonstra acelerações significativas em regimes de contexto longo.
- Em sequências curtas (2k tokens), o ganho de velocidade é modesto (~1,15×) devido ao overhead de roteamento.
- Em regimes de contexto longo (4k a 1024k tokens), o GQE alcança um aceleração de 1,67× a 1,80× sobre a linha de base GQA. Isso ocorre porque o método pula computações de query desnecessárias, e as economias escalam com o comprimento da sequência conforme o custo de atenção cresce quadraticamente.

Significância e Alegações

O artigo afirma que o GQE oferece um caminho viável para reduzir os custos de computação de autoatenção em Transformers de contexto longo sem sacrificar a qualidade do modelo. A significância reside em:

Eficiência: Reduz a computação de cabeças de query ativas em aproximadamente metade na configuração testada, mantendo a paridade de desempenho com o GQA denso.
Escalabilidade: A aceleração torna-se cada vez mais valiosa à medida que o comprimento do contexto aumenta, abordando o gargalo de escalonamento quadrático da atenção.
Insight de Design: O trabalho destaca que transferir MoE para a atenção requer um design arquitetural cuidadoso (especificamente a cabeça compartilhada e o sinal de roteamento renormalizado) para evitar as penalidades de precisão frequentemente vistas em tentativas de atenção esparsa ingênuas.

Os autores mantêm a modéstia, observando que os resultados estão na escala de 250M e devem ser confirmados em escalas maiores. Eles enquadram o resultado como um "empate" com a linha de base, em vez de uma melhoria clara, enfatizando a conquista de redução de computação para um desempenho equivalente.