Une tête d’attention à la main

Jusqu’au chapitre 7, le modèle savait ajuster des fonctions sur des exemples individuels. Les courbes de loss descendaient, les MLP résolvaient XOR, les optimiseurs faisaient leur travail. Le problème restant : un token n’a aucun moyen de regarder les autres tokens. Le bigramme ne voyait que le mot précédent ; un MLP appliqué à un embedding fixe ne voit que ce vecteur.

L’attention est ce mécanisme, et elle porte tous les modèles de langage de pointe depuis 2017. Le papier d’origine la disait « tout ce dont vous avez besoin » — une exagération qui s’est révélée approximativement juste.

La mécanique est mécanique. Une fois que tu as exécuté chaque morceau et inspecté les formes, tu peux expliquer l’attention à un ami au comptoir. C’est exactement ce qu’on va faire : quatre cellules exécutables qui construisent une tête d’attention de l’intérieur, sur une phrase jouet de 5 tokens avec des embeddings en 4 dimensions, pour que les matrices tiennent à l’écran. Ensuite, tu sauvegarderas une fonction d’attention causale minimale en local.

Phrase : the cat sat on mat. Cinq tokens, embeddings 4D faits à la main.

1. Projeter en Q, K, V

Une seule tête d’attention a trois matrices apprises : W_Q, W_K, W_V, chacune de taille [d × d_head]. Multiplier l’entrée X (de forme [seq_len × d]) par chacune produit les queries, les keys et les values — trois vues de la même entrée, calculées chacune par un jeu de poids différent :

Intuition :

• Queries : “qu’est-ce que je cherche ?”
• Keys : “voilà ce que j’ai”
• Values : “voilà ce que j’apporterai si tu me choisis”

Écris la multiplication matricielle pour Q. (On réutilise la même routine pour K et V — une fois que tu l’as écrite une fois, c’est fini.)

// X is the [seq_len × d] matrix of token embeddings.
// W_Q is the [d × d] query-projection matrix.
// Compute Q = X · W_Q  — shape [seq_len × d].
// Return the result.
return X;

Le résultat est [seq_len × d_head], la même forme que l’entrée dans notre cas jouet. Chaque ligne est le vecteur query d’un token de la phrase.

2. Scorer chaque paire

On demande maintenant : à quel point chaque query s’intéresse à chaque key ? La réponse standard est un produit scalaire — grand quand deux vecteurs pointent dans la même direction, petit quand ils sont orthogonaux, négatif quand ils s’opposent.

S est une matrice [seq_len × seq_len]. S[i][j] signifie : « à quel point le token j est-il pertinent pour le token i ? » La ligne i est la ligne des scores vus depuis le token i.

// Q and K are both [seq_len × d_k].
// Compute the score matrix S = Q · Kᵀ — shape [seq_len × seq_len].
// S[i][j] = dot product of Q[i] (query i) and K[j] (key j).
// This says: "how relevant is token j to token i?".
return [];

La heatmap montre les scores bruts. La matrice n’est pas symétrique en général : i qui interroge j n’est pas la même question que j qui interroge i, parce qu’elles passent par des W_Q et W_K différents. C’est cette asymétrie qui permet à l’attention d’exprimer des relations orientées, du type « ce verbe est gouverné par ce sujet ».

3. Échelle et softmax

Les scores ne sont pas encore des probabilités : ils peuvent avoir n’importe quel signe, n’importe quelle magnitude. Deux transformations les convertissent en une distribution de probabilité, ligne par ligne :

1. Mettre à l’échelle chaque entrée par 1/√d_k. Sans ça, les produits scalaires de vecteurs en grande dimension grandiraient assez pour saturer le softmax et faire disparaître les gradients à l’entraînement.
2. Appliquer le softmax à chaque ligne. Chaque ligne somme alors à 1.

// S is the raw score matrix from the previous cell.
// 'd_k' is the head dimension.
// Two steps to turn scores into attention weights:
//   1. divide every entry by sqrt(d_k)  (stabilizes for large d_k)
//   2. apply softmax row-wise           (each row sums to 1)
// Return the [seq_len × seq_len] attention-weight matrix.
return S;

(Softmax par ligne, appliqué indépendamment à chacune.)

Regarde la ligne i dans la heatmap. Ses cellules te disent, pour le token i, quelle part de sa représentation mise à jour viendra de chaque autre token. Si le token 2 a une valeur forte en colonne 4, le modèle considère le token 4 comme très pertinent en tant que source pour la mise à jour du token 2. Les lignes somment à 1 : ce sont de vraies distributions de probabilité.

4. Mélanger les values

Dernière étape : la sortie de chaque token est la somme des vecteurs value pondérée par l’attention. Le token i prend une moyenne pondérée des values de tous les tokens, avec les poids de sa ligne d’attention A[i].

Sous forme matricielle, output = A · V, de forme [seq_len × d]. La même forme que l’entrée X — la tête réarrange la représentation de chaque token en y tirant des morceaux des autres.

// A is the [seq_len × seq_len] attention-weights matrix.
// V is the [seq_len × d] value matrix.
// Output = A · V  — shape [seq_len × d].
// Each output row is a weighted sum of value rows, weighted by attention.
return [];

C’est tout le calcul d’une attention à une tête. Cinq lignes de maths, quatre opérations matricielles. Empile la même machinerie quelques dizaines de fois (c’est ce que font les chapitres suivants), entraîne-la sur un milliard de tokens de texte, et tu obtiens GPT.

Pourquoi ça marche

Le mécanisme n’est pas profond, mais ses implications le sont. Une seule tête d’attention peut implémenter, selon les poids qu’elle a appris :

• Une copie — chaque token regarde le précédent, la sortie est la value du token d’avant. (Utile pour la répétition.)
• Un lookup — pour chaque « the », regarder le nom qui suit. (Fréquent en modélisation du langage.)
• Un accord — pour chaque verbe, regarder son sujet. (Dépendance à longue portée.)
• Un résumé — chaque token regarde à peu près également tous les autres, ce qui moyenne la séquence. (Utile dans la dernière couche.)

La procédure d’entraînement (descente de gradient sur la loss next-token) découvre quels motifs le réseau a besoin d’avoir. On ne les lui dit jamais. Un LLM moderne a des dizaines de ces têtes par couche et des dizaines de couches ; chaque tête se spécialise pendant l’entraînement.

5. Ajouter l’attention causale localement

Crée llm/attention.py :

"""Readable attention helpers before the PyTorch version."""
from __future__ import annotations
 
import math
 
 
Vector = list[float]
Matrix = list[Vector]
 
 
def dot(a: Vector, b: Vector) -> float:
    return sum(x * y for x, y in zip(a, b))
 
 
def softmax(values: Vector) -> Vector:
    m = max(values)
    exps = [math.exp(v - m) for v in values]
    total = sum(exps)
    return [v / total for v in exps]
 
 
def matmul(x: Matrix, w: Matrix) -> Matrix:
    columns = list(zip(*w))
    return [[dot(row, list(col)) for col in columns] for row in x]
 
 
def causal_attention(x: Matrix, wq: Matrix, wk: Matrix, wv: Matrix) -> Matrix:
    # [1]
    q = matmul(x, wq)
    k = matmul(x, wk)
    v = matmul(x, wv)
    scale = math.sqrt(len(k[0]))
 
    out: Matrix = []
    for i, query in enumerate(q):
        # [2]
        scores = [
            dot(query, key) / scale if j <= i else -1e9
            for j, key in enumerate(k)
        ]
        # [3]
        weights = softmax(scores)
        # [4]
        out.append([
            sum(weight * value[d] for weight, value in zip(weights, v))
            for d in range(len(v[0]))
        ])
    return out

Lis-le comme quatre passages sur la même séquence :

• [1] q, k et v sont trois vues apprises de x. Mêmes tokens, questions différentes.
• [2] scores compare la query du token i à chaque key. j <= i est le masque causal : les tokens passés et courant sont visibles, les futurs reçoivent -1e9.
• [3] softmax(scores) fabrique une distribution de probabilité pour le token i.
• [4] construit une moyenne pondérée des vecteurs value. C’est la nouvelle représentation du token i.

L’ajout important, c’est le masque causal : le token i ne peut lire que les tokens 0..i. Sans lui, le modèle pourrait tricher pendant l’entraînement next-token en regardant la réponse.

Pour aller plus loin

• The Illustrated Transformer, par Jay Alammar — l’explication visuelle de l’attention la plus reprise sur internet.
• Karpathy, « Let’s build GPT from scratch » — deux heures d’attention codée en direct.
• Step by Token, chapitre 4 traite l’attention sous l’angle de la compréhension.
• La référence mathématique de chaque cellule vit dans components/chapter/ch08/_shared.ts.

Prochaine étape : multi-têtes et résidus — une tête ne suffit pas, et il faut une connexion pour en empiler beaucoup.