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 le mécanisme qui laisse le contexte circuler. Depuis 2017, elle porte tous les modèles de langage modernes. L'attention est le mécanisme qui permet à un token de regarder d'autres tokens.

On va la construire de l’intérieur : quatre cellules exécutables sur une phrase jouet de 5 tokens avec des embeddings 4D. Ensuite, tu sauvegarderas une attention causale minimale localement.

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

1. Projeter en Q, K, V

Une tête d’attention a trois matrices apprises : W_Q, W_K, W_V. Multiplier l’entrée X par chacune produit les queries, keys et values :

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 de matrice pour Q. La même routine sert pour K et V.

Le résultat est [seq_len × d_head]. Chaque ligne est la query d’un token.

2. Scorer chaque paire

On demande maintenant : à quel point chaque query s’intéresse à chaque key ? La réponse standard est un produit scalaire.

S est une matrice [seq_len × seq_len]. S[i][j] signifie : “à quel point le token j est pertinent pour le token i ?”

La heatmap montre les scores bruts. La matrice n’est pas forcément symétrique : i qui regarde j n’est pas la même question que j qui regarde i, car les projections Q et K sont différentes.

3. Échelle et softmax

Les scores ne sont pas encore des probabilités. Deux transformations les convertissent ligne par ligne :

1. Diviser par √d_k.
2. Appliquer softmax à chaque ligne.

Chaque ligne de la heatmap dit, pour un token donné, quelle fraction de sa représentation mise à jour viendra de chaque autre token. Les lignes somment à 1 : ce sont de vraies distributions.

4. Mélanger les values

Dernière étape : la sortie de chaque token est une somme pondérée des value vectors.

Sous forme matricielle : output = A · V.

C’est toute l’attention simple tête : cinq lignes de maths, quatre opérations matricielles. Empile ça, entraîne sur beaucoup de texte, et tu obtiens la famille GPT.

Pourquoi ça marche

Une tête peut apprendre différents motifs :

• Copie : chaque token regarde le précédent.
• Lookup : chaque the regarde le nom qui suit.
• Accord : chaque verbe regarde son sujet.
• Résumé : chaque token moyenne toute la séquence.

La descente de gradient découvre les motifs nécessaires. On ne les code pas à la main. L'attention permet au modèle d'apprendre différents motifs de relation entre les tokens.

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

• [1] calcule les trois vues apprises de x.
• [2] compare la query du token i à toutes les keys. j <= i est le masque causal.
• [3] transforme les scores en distribution.
• [4] construit la moyenne pondérée des values.

Le masque causal est crucial : le token i ne peut lire que 0..i. Sinon, pendant l’entraînement next-token, il pourrait regarder la réponse.

Pour aller plus loin

• The Illustrated Transformer.
• Karpathy, “Let’s build GPT from scratch”.
• Step by Token, chapitre 4.
• 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.