Phase 07: Transformers Deep Dive

Une attention à plusieurs têtes

Type: Build

Languages: Python

Prerequisites: Phase 7 · 02 (Self-Attention from Scratch)

Time: ~75 minutes

Le problème

Une seule tête d'attention à soi compute une matrice d'attention. Cette matrice capture un type de relation généralement celle qui minimise la perte sur quoi que ce soit le signal d'entraînement. Si vos données ont un accord sujet-verbe, une co-référence, un discours à longue portée et un chunking syntaxique tous enracinés ensemble, une seule tête les épluche dans une seule distribution de max douce et perd la moitié du signal.

La correction du papier Vaswani 2017: exécuter plusieurs fonctions d'attention en parallèle, chacune avec ses propres projections Q, K, V, et concatener les sorties. Chaque tête fonctionne dans un sous-espace de dimension plus petit d_model / n_headsLes paramètres totaux restent les mêmes.

L'attention multi-tête est la norme par défaut de chaque transformateur dans les navires de 2026. Le seul argument est sur combien de têtes et si les touches et les valeurs partagent des projections (attention groupée, attention multi-quéret, attention latente multi-tête).

Le concept

!Multi-head attention splits, attends, concatenates

Split.Prenez .Xde forme (N, d_model)- Projet à Q, K, V de forme(N, d_model)- Ressoufflez-vous .(N, n_heads, d_head)où d_head = d_model / n_heads- Transposer à(n_heads, N, d_head)- Je suis désolé .

Attend in parallel.Exécutez des produits à l'échelle de points à l'intérieur de chaque tête.(N, d_head)Les têtes fonctionnent sur différents sous-espaces de l'embedding et ne parlent jamais pendant le calcul de l'attention lui-même.

Concatenate and project.Les têtes de pile retournent à (N, d_model)et multiplier par une matrice de sortie apprise W_ode forme (d_model, d_model)- Je suis là .W_oC'est là que les têtes se mélangent.

Why it works.Chaque tête peut se spécialiser sans rivaliser avec les autres pour le budget représentatif. Les études de sondage de 2019 à 2024 montrent des rôles distincts de tête: les têtes positionnelles, la tête qui assiste au jeton précédent, les têtes de copie, les têtes d'entité nommée, les têtes d'induction (qui sous-tendent l'apprentissage dans le contexte).

The 2026 lineage of variations:

GQA est la norme par défaut moderne car elle réduit la mémoire KV-cache d'un facteur de N/GMLA va plus loin en comprimant K/V dans un espace latent, puis en le projetant à l'heure de calcul coûte FLOPs, économise beaucoup plus de mémoire.

Faites-le

Étape 1: séparer les têtes de l'attention à tête unique que nous avons déjà

Prenez le SelfAttentionLeur résultat est le résultat de la leçon 02 et de l'envelopper avec une paire de fractionnement/concétation.code/main.pypour une mise en œuvre numpy; la logique est:

pythonCopydef split_heads(X, n_heads):
    n, d = X.shape
    d_head = d // n_heads
    return X.reshape(n, n_heads, d_head).transpose(1, 0, 2)  # (heads, n, d_head)

def combine_heads(H):
    h, n, d_head = H.shape
    return H.transpose(1, 0, 2).reshape(n, h * d_head)

Une remodèle et une transposition, pas de boucle, c'est exactement ce que fait PyTorch sousnn.MultiheadAttention- Je suis désolé .

Étape 2: exécuter des points-produit à l'échelle de l'attention par tête

Chaque tête obtient sa propre tranche de Q, K, V. L'attention devient un matmul en lots:

pythonCopydef mha_forward(X, W_q, W_k, W_v, W_o, n_heads):
    Q = X @ W_q
    K = X @ W_k
    V = X @ W_v
    Qh = split_heads(Q, n_heads)         # (heads, n, d_head)
    Kh = split_heads(K, n_heads)
    Vh = split_heads(V, n_heads)
    scores = Qh @ Kh.transpose(0, 2, 1) / np.sqrt(Qh.shape[-1])
    weights = softmax(scores, axis=-1)
    out = weights @ Vh                    # (heads, n, d_head)
    concat = combine_heads(out)
    return concat @ W_o, weights

Sur du matériel réel .Qh @ Kh.transpose(...)est unbmmLa GPU voit une seule forme de matmul .(heads, N, d_head) × (heads, d_head, N) -> (heads, N, N)- Ajouter des têtes est gratuit.

Étape 3: Variante de l'attention de la requête groupée

Seules les projections de clé et de valeur changent.n_headsles groupes; K et V obtiennent n_kv_heads < n_headsles groupes et sont répétés pour correspondre:

pythonCopydef gqa_project(X, W, n_kv_heads, n_heads):
    kv = split_heads(X @ W, n_kv_heads)       # (kv_heads, n, d_head)
    repeat = n_heads // n_kv_heads
    return np.repeat(kv, repeat, axis=0)      # (n_heads, n, d_head)

En conséquence , cela permet d' économiser la mémoire , car seulementn_kv_headsLes copies sont en direct dans le cache KV, pas n_headsLlama 3 70B utilise 64 têtes de requête avec 8 têtes KV un rétrécissement de cache 8x.

Étape 4: enquête sur ce que chaque tête a appris

Exécutez MHA sur une phrase courte avec 4 têtes.(N, N)Vous verrez différentes têtes choisir différentes structures même avec l'initialisation aléatoire qui est en partie le signal, en partie la symétrie rotative dans les sous-espaces.

Utilisez-le

Dans PyTorch, la version à une ligne:

pythonCopyimport torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)

GQA à partir de PyTorch 2.5+:

pythonCopyfrom torch.nn.functional import scaled_dot_product_attention

# scaled_dot_product_attention auto-dispatches Flash Attention on CUDA.
# For GQA, pass Q of shape (B, n_heads, N, d_head) and K,V of shape
# (B, n_kv_heads, N, d_head). PyTorch handles the repeat.
out = scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

How many heads?Règles générales des modèles de production en 2026:

d_headIl est l'unité de la quantité qu'une tête peut "voir".sqrt(d_head)Si vous dépassez 256 points, vous perdez le bénéfice "beaucoup de petits spécialistes".

La faire partir

Regardez !outputs/skill-mha-configurator.md. La compétence recommande le nombre de têtes, le nombre de têtes kv et la stratégie de projection pour un nouveau transformateur, compte tenu du budget des paramètres, de la longueur des séquences et de l'objectif de déploiement.

Exercices

1. Easy.Prenez le MHA de code/main.pyet le changement n_headsde 1 à 16 avec d_model=64Une fois que vous avez terminé, vous pouvez faire une copie synthétique.
2. Medium.MQA (un seul KV partagé sur tous les têtes de requête). Mesurer combien de paramètres du compte de gouttes par rapport à MHA complet. Comptez combien la taille de la caisse KV se réduit à l'inférence pour N = 2048.
3. Hard.Mettez en œuvre une version minuscule de l'attention latente multi-tête: comprimant K, V à un rang-rLaissez-le dans le cache KV, décomprimez-le à l'heure de l'attention.rla mémoire cache passe-t-elle en dessous de 1/8 de la MHA complète tandis que la qualité reste à l'intérieur de 1 bit de la validation ?

Les termes clés

Pour en savoir plus

• Vaswani et al. (2017). Attention Is All You Need §3.2.2 la spécification originale de plusieurs têtes.
• Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need le document MQA.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints comment convertir l'AMM en GQA après la formation.
• DeepSeek-AI (2024). DeepSeek-V2 Technical Report MLA et pourquoi il bat MHA/GQA sur la mémoire cache.
• Olsson et al. (2022). In-context Learning and Induction Heads Regardez mécaniquement ce que font réellement les têtes.

This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.

Browse the complete course catalog or open this lesson on GitHub.