Grouped Query Attention (GQA)

A atenção de consulta agrupada (GQA) é uma forma de atenção multi-cabeça para Transformers, na qual grupos de cabeças de consulta partilham uma cabeça de chaves e uma cabeça de valores. Por exemplo, o Llama 3.1 8B combina 32 cabeças de consulta com 8 cabeças de chave-valor em cada camada. Situa-se entre a atenção multi-cabeça (MHA), na qual cada cabeça de consulta tem as suas próprias chaves e valores, e a atenção multi-consulta (MQA), na qual todas as cabeças de consulta partilham uma única cabeça de chaves e uma única cabeça de valores. Joshua Ainslie e colegas da Google Research apresentaram-na em GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, na EMNLP 2023, mostrando que os modelos GQA alcançam uma qualidade próxima da MHA e uma velocidade de descodificação próxima da MQA. Modelos de linguagem de grande escala (LLM) abertos, como o Llama 2 70B, a família Llama 3, o Mistral 7B e o Gemma 2, usam esta arquitetura; os modelos de visão-linguagem (VLM) construídos com base neles, como o PaliGemma 2 da Google, baseado no Gemma 2, herdam-na nos descodificadores de linguagem.

Como funciona a GQA#

Na autoatenção, cada token é projetado numa consulta, numa chave e num valor para cada cabeça, e cada cabeça compara as respetivas consultas com as chaves. A GQA mantém o conjunto completo de cabeças de consulta, mas projeta menos cabeças de chaves e valores; cada grupo de cabeças de consulta presta atenção às chaves e aos valores partilhados pelo seu grupo. Com 8 cabeças de consulta e 2 cabeças de chave-valor, cada cabeça de chave-valor serve 4 cabeças de consulta:

O cálculo da atenção dentro de cada cabeça continua a ser atenção de produto escalar com escala, sem alterações. Com um grupo, a camada torna-se MQA; com tantos grupos quantas cabeças de consulta, torna-se MHA. O PyTorch suporta diretamente esta disposição através do argumento enable_gqa de scaled_dot_product_attention:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

Imprime Output (1, 8, 16, 64), same as copied heads: True e 256 with GQA, 1024 with MHA: as 8 cabeças de consulta produzem resultados usando um quarto das chaves e dos valores que a MHA armazenaria.

Por que razão a GQA reduz a cache KV#

Durante a geração de texto, um modelo mantém as chaves e os valores de cada token anterior na sua cache KV e lê toda a cache para cada novo token, pelo que a velocidade de descodificação é limitada mais pela largura de banda da memória do que pela aritmética. O tamanho da cache aumenta com o número de cabeças de chave-valor. O Llama 3.1 8B tem 32 camadas, cada uma com 8 cabeças de chave-valor de 128 dimensões, pelo que armazena 128 KiB por token com precisão de 16 bits, ou 16 GiB para a sua janela de contexto completa de 131.072 tokens. Com MHA, o mesmo modelo precisaria de 512 KiB por token e 64 GiB para o comprimento máximo. A memória libertada pode ser usada para contextos mais longos, lotes maiores em batching contínuo e menor latência de inferência por token. A interface LLM da Ultralytics aceita um base_url para qualquer servidor compatível com OpenAI, como uma instância autoalojada do Llama 3.1 8B; com as suas 8 cabeças de chave-valor, o servidor consegue manter na cache quatro vezes mais tokens de conversa do que uma disposição MHA com 32 cabeças, usando a mesma quantidade de memória.

Converter modelos para GQA#

A GQA não requer treino do zero. Os autores converteram um checkpoint MHA existente calculando a média das cabeças de chaves e valores de cada grupo e continuando depois o pré-treino com cerca de 5% do cálculo original, um processo a que chamam treino adicional. A MQA, o caso extremo de uma única cabeça partilhada, foi apresentada por Noam Shazeer em 2019 e é a disposição mais rápida, mas degradava a qualidade e podia tornar o treino instável; um número intermédio de grupos recupera a maior parte dessa qualidade.

GQA vs. técnicas relacionadas#

• Atenção latente multi-cabeça (MLA): Apresentada com o DeepSeek-V2, comprime as chaves e os valores de cada token num pequeno vetor latente partilhado.
• PagedAttention: Aloca a cache KV em blocos de tamanho fixo para evitar a fragmentação, independentemente do número de cabeças que a preenche.
• Atenção de janela deslizante: Limita o número de tokens em cache por camada; o Mistral 7B combina-a com a GQA.
• Flash attention: Um kernel de GPU que calcula a atenção exata mais rapidamente e suporta nativamente disposições GQA.