Especialistas de Consulta Agrupada: Mistura de Especialistas na Autoatenção GQA

Resumo

A autoatenção é central para o desempenho do Transformer e frequentemente é a parte mais dispendiosa do Transformer em contextos longos, pois suas interações pareadas de tokens escalam quadraticamente com o comprimento da sequência. A atenção densa padrão também aplica o mesmo conjunto de cabeças de atenção a cada token, independentemente da dificuldade do token ou do conteúdo da informação. Essa ativação uniforme pode desperdiçar poder computacional, especialmente à medida que as sequências se alongam e o custo da atenção aumenta rapidamente. Propomos os Especialistas de Consulta Agrupada (Grouped Query Experts, GQE), uma camada de mistura de especialistas sobre a atenção de consulta agrupada (Grouped-Query Attention, GQA). Dentro de cada grupo GQA, um roteador seleciona k especialistas de cabeça de consulta por token, enquanto todas as cabeças chave-valor (KV) permanecem densas e inalteradas. Assim, o GQE mantém os benefícios do cache KV do GQA e reduz apenas o cálculo ativo das cabeças de consulta. Com um orçamento fixo de 30B tokens na escala de 250M parâmetros, o GQE iguala a linha de base do GQA totalmente ativo em precisão downstream, enquanto ativa metade das cabeças de consulta por token.

English

Self-attention is central to Transformer performance and is often the most expensive part of the Transformer at long context lengths because its pairwise token interactions scale quadratically with sequence length. Standard dense attention also applies the same set of attention heads to every token regardless of token difficulty or information content. This uniform activation can waste compute, especially as sequences grow longer and attention cost increases rapidly. We propose Grouped Query Experts (GQE), a mixture-of-experts layer on top of grouped-query attention (GQA). Within each GQA group, a router selects k query-head experts per token while all key-value (KV) heads remain dense and unchanged. Thus, GQE keeps the KV cache benefits of GQA and reduces only the active query-head computation. On a fixed 30B token budget at the 250M parameter scale, GQE matches the all-active GQA baseline in downstream accuracy while activating half the query heads per token.