o problema
Um decoder transformer gera um token por vez, e a conta que domina não é a de multiplicação de matriz. A cada passo o chip precisa carregar da memória os pesos do decoder e todas as keys e values já acumuladas na sequência. Quanto mais longa a saída, mais bytes atravessam o barramento por token gerado. É um problema de banda de memória, não de FLOP, e por isso ele não some quando você compra um acelerador mais rápido.
Multi-query attention, proposto por Shazeer em 2019, atacou isso da forma mais direta possível: mantém as várias query heads e reduz key e value a uma head só. O KV cache encolhe por um fator igual ao número de heads. O preço é que a qualidade cai e o treino fica instável — os autores treinaram vários T5-Large com MQA do zero e todos sofreram picos de loss no pré-treino e divergiram imediatamente no fine-tuning de entrada longa. Havia ainda um problema logístico: PaLM já nascia com MQA, mas T5 e LLaMA não. Quem quisesse inferência rápida teria que pré-treinar um modelo separado só para isso.
a ideia
Dois movimentos, independentes um do outro.
O primeiro: não é preciso treinar do zero. Dá para pegar um checkpoint multi-head pronto, colapsar as projeções de key e value na estrutura desejada e rodar mais um pedaço pequeno de pré-treino para o modelo se acomodar. Os autores chamam isso de uptraining, emprestado do trabalho que converte checkpoints densos em mixture-of-experts.
O segundo: MQA e MHA não são as duas únicas opções, são as duas pontas de um botão. Divida as query heads em G grupos e dê a cada grupo uma head de key e uma de value. Com G igual a 1 você tem MQA; com G igual ao número de heads, MHA. O meio do caminho é o interessante — e ele importa mais conforme o modelo cresce, porque modelos maiores escalam o número de heads, o que faz de MQA um corte cada vez mais agressivo em capacidade. Com grupos, a proporção do corte se mantém.
como funciona
A conversão é mean pooling: as matrizes de projeção de key e value das heads de um grupo viram uma só, pela média. O ablation compara com escolher a primeira head e com inicializar do zero, e a ordem do resultado segue a quantidade de informação preservada — média ganha, seleção fica no meio, aleatório perde.
Depois vem o pré-treino adicional, por uma proporção α dos passos originais, com a mesma receita e o mesmo dataset. Para α igual a 0,05 no T5-XXL, isso custou cerca de 600 TPUv3 chip-days.
MQA e GQA entram na self-attention do decoder e na cross-attention, mas não na self-attention do encoder: o encoder computa tudo em paralelo, então banda de memória não é o gargalo lá.
Nos experimentos com T5.1.1, o MHA-XXL leva 1,51 s por amostra e faz 47,2 de média nos benchmarks; o MQA-XXL uptreinado cai para 0,24 s e 46,6; o GQA de 8 grupos fica em 0,28 s e 47,1. Oito grupos foi escolhido olhando a curva de tempo por número de grupos: sair de 1 para 8 custa pouco, e o custo sobe rápido depois disso.
o que isso custou
Os autores listam os limites sem enfeite. A qualidade em geração longa é medida com Rouge, que eles descrevem como uma avaliação falha que não conta a história toda — ou seja, não dá para ter certeza de que o trade-off está calibrado certo. Por falta de compute, não existe comparação entre o XXL com GQA uptreinado e um GQA equivalente treinado do zero: ninguém sabe quanto se perde pelo atalho. E toda a avaliação é em modelos encoder-decoder, num momento em que o mundo já tinha virado para decoder-only.
A instabilidade do MQA sobrevive parcialmente ao uptraining: os modelos convertidos são mais estáveis que os treinados do zero, mas ainda têm variância alta, ao ponto de os números reportados em tarefas instáveis serem média de três fine-tunings. GQA se mostrou estável, e por isso a causa raiz da instabilidade do MQA ficou sem investigação — um bug conhecido e contornado, não resolvido.
onde isso aparece hoje
A previsão dos próprios autores era que GQA teria vantagem ainda maior em modelos decoder-only, que não têm cross-attention separada. Foi o que aconteceu: GQA virou configuração padrão em famílias abertas, incluindo os modelos maiores do Llama 2 e o Mixtral.
O trabalho se encaixa num conjunto de técnicas que atacam o mesmo gargalo por ângulos diferentes e compõem entre si: FlashAttention reorganiza a computação para não materializar a matriz de scores, decodificação especulativa esconde a latência propondo tokens com um modelo menor, e o vLLM trata a fragmentação do KV cache como problema de gerência de memória. GQA é o que reduz o tamanho do cache na origem.