🧙 Maestro Yoda Cap. 19 · Architetture

Puntata 208

Puntata 208 — Multi-head attention

Livello: 🧙 Maestro Yoda · Capitolo 19 · Architetture

Una sola testa di attention vede una storia per volta. Otto teste vedono otto storie in parallelo. Trentadue teste cominciano a vedere cose che noi umani manco pensiamo.

schema di 8 “teste” disegnate come paia di occhi diversi (binocolo, microscopio, telescopio, occhiale spia, ecc.), tutte puntate sullo stesso testo. Ognuna evidenzia una parte diversa: una segue il soggetto, una la coreferenza, una la sintassi, una la traduzione di caso/numero/genere. Tag: “ogni testa, una specializzazione emergente”

Perché una testa sola non basta

Nella self-attention con una testa sola, ogni token produce una sola distribuzione sui token vicini e ne fa una sola media pesata. Tutto il contesto rilevante per quel token deve passare da quel singolo collo di bottiglia.

Ma le relazioni tra token sono molteplici e di tipo diverso:

  • chi è il soggetto del verbo (sintassi)
  • chi è il referente di un pronome (anafora)
  • chi è il modificatore semantico (aggettivo→nome)
  • traduzioni di genere/numero/caso
  • relazioni di topica/distanza

Un solo softmax(QK^T) può evidenziare un solo pattern alla volta. Multi-head è la risposta: fai più attention in parallelo, ognuna con parametri propri, poi mescola.


La meccanica

Hai dimensione del modello d_model. Definisci h teste, ciascuna con dimensione d_k = d_v = d_model / h.

Per ogni testa i:

Q_i = X W_Q^i      (n × d_k)
K_i = X W_K^i      (n × d_k)
V_i = X W_V^i      (n × d_v)
head_i = softmax(Q_i K_i^T / √d_k) V_i        (n × d_v)

Poi concateni gli output e proietti:

MHA(X) = Concat(head_1, ..., head_h) · W_O      (n × d_model)

Il totale dei parametri (W_Q, W_K, W_V, W_O) è ~ 4 · d_model², indipendentemente dal numero di teste — perché h · d_k = d_model.

Importante: il costo computazionale di MHA è uguale a quello di una single-head con d_model. Multi-head non è “più costoso”, è “lo stesso compute distribuito su più sottospazi”.


Esempio: GPT-3, Llama-3 70B, Mistral

Modellod_modelh tested_k per testa
GPT-3 175B1228896128
GPT-2 small7681264
Llama-3 8B409632128
Llama-3 70B819264128
Mistral 7B409632128
Gemma-2 9B358416256

Tipicamente d_k = 64, 128, 256. Più teste = pattern più granulari; teste più larghe = pattern più ricchi. C’è un trade-off.


Cosa “imparano” davvero le teste?

Studi di interpretability (Vig 2019, Clark et al. 2019, Olsson et al. 2022) hanno mostrato che teste diverse si specializzano in pattern diversi. Tre esempi celebri:

Induction heads

Coppie di teste (in layer adiacenti) che cercano nel passato il pattern [A][B] ... [A] e copiano [B] come prediction. È il meccanismo base dell’in-context learning: il modello che “completa il pattern” guardando esempi nel prompt.

Identificato in Olsson et al. (Anthropic, 2022) come precursore dell’apparizione del comportamento few-shot durante il training.

Coreference heads

Teste che mappano pronomi al loro antecedente: “Mario disse che lui era stanco” → la testa connette “lui” → “Mario”.

Sintattiche

Teste che ricostruiscono dipendenze grammaticali: soggetto → verbo, oggetto → verbo, modificatore → testa nominale.

Sono emergenti, non scritte a mano. Nessuno ha programmato la testa che fa coreferenza: è il backprop che ce l’ha messa lì perché serve a predire il prossimo token.


Varianti moderne: MQA, GQA, MLA

A inferenza, le matrici K e V di ogni layer vengono cachate (la KV cache) per non ricomputarle. Più teste = più memoria. A context lunghi questa diventa la metrica dominante.

MQA — Multi-Query Attention (Shazeer, 2019)

Una sola coppia K, V condivisa tra tutte le teste; solo Q resta multi-head. Riduce KV cache di un fattore h. Quality perde un po’ rispetto al multi-head pieno.

GQA — Grouped Query Attention (Ainslie et al., Google 2023)

Compromesso: gruppi di teste condividono K, V. Es. 32 teste Q, 8 gruppi → ogni gruppo di 4 teste condivide una K, V. Adottato da Llama-2 70B, Llama-3, Mistral, Gemma.

MLA — Multi-head Latent Attention (DeepSeek-V2, 2024)

Comprime K, V in un vettore latente di dimensione molto più piccola, da cui ricostruisce per ogni testa. Risultato: KV cache drasticamente più piccola (~10×) con qualità preservata. Pattern adottato in DeepSeek-V3/R1.


Tabella sintetica

Schemah Qh KVKV cache sizeQualityUso 2026
MHA classicohhgrandegold standardpiccoli modelli
MQAh1minimomediocreFalcon-7B (storico)
GQAhh/g (g=2,4,8)medioquasi MHAstandard (Llama, Mistral, Gemma)
MLAhlatent compressominimocome MHADeepSeek V2/V3

Flash-Attention: il “come” più importante del “cosa”

Negli ultimi 3 anni, l’innovazione che ha cambiato di più la vita pratica del multi-head attention non è stata architetturale ma di kernel: FlashAttention (Dao et al., 2022, 2023, 2024).

L’idea: la matrice di attention n × n non si materializza mai in memoria globale (HBM). Si fa tutto a blocchi (tiling), in SRAM della GPU, ricomputando dove serve.

Risultato:

  • 2-4× più veloce in training
  • Memoria O(n) invece di O(n²) per il forward pass
  • Numericamente più stabile (calcola in modo “online” il softmax)

Senza FlashAttention, context >8k era impossibile su hardware single-GPU. Con FlashAttention-3 (2024), si tengono 128k context su una H100 con relativa naturalezza.

È una “rivoluzione invisibile”: stesso paper di Vaswani et al., stesso algoritmo, ma 10× più veloce in pratica. Quasi tutti i framework moderni la usano di default.


Glossario lampo

  • KV cache — memoria persistente di K, V durante la generation autoregressive.
  • Head dimd_k, dimensione di Q/K/V per ogni testa.
  • Concat-and-project — il pattern finale di MHA: si concatenano gli output delle teste e si proiettano via W_O.
  • Induction head — testa che fa pattern matching nel passato, base dell’in-context learning.

TEST hands-on

  1. Implementa multi-head attention in PyTorch in ~30 righe. Verifica le shape per d_model = 64, h = 8, n = 16. Output: (batch, n, d_model).
  2. Confronta i tuoi risultati con torch.nn.MultiheadAttention. Devono coincidere a meno di permutazione delle teste.
  3. Apri un transformer pretrainato (es. GPT-2 small via transformers). Estrai i pesi di una testa per un prompt corto e visualizza la matrice di attention come heatmap. Vedi se riesci a identificare pattern interpretabili (coreferenza, sintassi).
  4. Bonus: leggi transformer-circuits.pub/2022/in-context-learning-and-induction-heads/. È accessibile e cambia il modo in cui guardi i layer di un LLM.

Take-away

Multi-head attention non è “più attention”. È attention in sotto-spazi paralleli, ognuno libero di specializzarsi in un pattern diverso, senza pagare in compute. Le innovazioni del 2023-26 (GQA, MLA, FlashAttention) hanno tutte mirato a un solo problema: ridurre la KV cache senza perdere qualità. Da lì arrivano i context da 128k+ che diamo per scontati nel 2026.


➡️ Prossima puntata: Norm, attivazioni, residui — i mattoncini fra un blocco e l’altro.