Grouped Query Attention (GQA)

L'attenzione a query raggruppate (GQA) è una forma di attenzione multi-testa per i transformer, in cui gruppi di teste query condividono una testa key e una testa value. Llama 3.1 8B, per esempio, abbina 32 teste query a 8 teste key-value in ogni livello. GQA si colloca tra l'attenzione multi-testa (MHA), in cui ogni testa query ha le proprie key e value, e l'attenzione multi-query (MQA), in cui tutte le teste query condividono una sola testa key e una sola testa value. Joshua Ainslie e colleghi di Google Research l'hanno introdotta in GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints a EMNLP 2023, mostrando che i modelli GQA raggiungono una qualità vicina a MHA con una velocità di decodifica vicina a MQA. La usano modelli linguistici di grandi dimensioni (LLM) open source come Llama 2 70B, la famiglia Llama 3, Mistral 7B e Gemma 2; inoltre, i modelli visione-linguaggio (VLM) che si basano su questi modelli, come PaliGemma 2 di Google basato su Gemma 2, la ereditano nei rispettivi decoder linguistici.

Come funziona GQA#

Nell'autoattenzione, ogni token viene proiettato in una query, una key e un value per ogni testa, e ogni testa confronta le proprie query con le proprie key. GQA mantiene l'intero insieme di teste query, ma proietta un numero inferiore di teste key e value; ogni gruppo di teste query applica l'attenzione alle key e ai value condivisi dal proprio gruppo. Con 8 teste query e 2 teste key-value, ogni testa key-value serve 4 teste query:

Il calcolo dell'attenzione all'interno di ogni testa resta invariato: è un'attenzione a prodotto scalare scalato. Con un solo gruppo, il livello diventa MQA; con tanti gruppi quante sono le teste query, diventa MHA. PyTorch supporta direttamente questa disposizione tramite l'argomento enable_gqa di 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")

Stampa Output (1, 8, 16, 64), same as copied heads: True e 256 with GQA, 1024 with MHA: tutte le 8 teste query producono output usando un quarto delle key e dei value che MHA memorizzerebbe.

Perché GQA riduce la cache KV#

Durante la generazione di testo, un modello conserva le key e i value di ogni token precedente nella cache KV e legge l'intera cache per ogni nuovo token; la velocità di decodifica è quindi limitata più dalla larghezza di banda della memoria che dai calcoli. Le dimensioni della cache aumentano con il numero di teste key-value. Llama 3.1 8B ha 32 livelli, ciascuno con 8 teste key-value da 128 dimensioni, e memorizza quindi 128 KiB per token con precisione a 16 bit, ovvero 16 GiB per la sua finestra di contesto completa di 131.072 token. Con MHA, lo stesso modello avrebbe bisogno di 512 KiB per token e 64 GiB per la lunghezza massima. La memoria liberata permette di gestire contesti più lunghi, batch più grandi per il batching continuo e una latenza di inferenza inferiore per token. L'interfaccia LLM di Ultralytics accetta un base_url per qualsiasi server compatibile con OpenAI, per esempio un Llama 3.1 8B ospitato autonomamente; con 8 teste key-value, il server può conservare nella stessa quantità di memoria della cache una conversazione quattro volte più lunga rispetto a una configurazione MHA con 32 teste.

Conversione dei modelli a GQA#

GQA non richiede l'addestramento da zero. Gli autori hanno convertito un checkpoint MHA esistente calcolando la media delle teste key e value all'interno di ciascun gruppo, poi hanno continuato il preaddestramento usando circa il 5% del calcolo originale, in un passaggio che chiamano uptraining. MQA, l'estremo con una sola testa condivisa, è stata introdotta da Noam Shazeer nel 2019 ed è la configurazione più veloce, ma riduceva la qualità e poteva rendere instabile l'addestramento; un numero intermedio di gruppi recupera gran parte di tale qualità.

GQA e tecniche correlate#

• Attenzione latente multi-testa (MLA): introdotta con DeepSeek-V2, comprime le key e i value di ogni token in un piccolo vettore latente condiviso.
• PagedAttention: alloca la cache KV in blocchi di dimensioni fisse per evitare la frammentazione, indipendentemente dal numero di teste che la riempiono.
• Attenzione a finestra scorrevole: limita il numero di token memorizzati nella cache per livello; Mistral 7B la combina con GQA.
• Attenzione Flash: un kernel GPU che calcola l'attenzione esatta più rapidamente e supporta nativamente le configurazioni GQA.