Multi-têtes et résidus

Au chapitre 8, tu as construit une tête d’attention. Elle calculait un motif d’attention — une matrice de qui écoute qui — et réarrangeait les représentations des tokens en conséquence. C’est un seul motif de routage par couche.

Ce n’est pas assez. Un vrai texte contient plusieurs relations simultanées : accord sujet/verbe, résolution de pronoms, rattachement des modificateurs, similarité sémantique, biais de position. N’importe laquelle peut être la bonne chose à regarder selon le token. Un modèle avec une seule tête par couche doit choisir.

La correction est celle qu’on devine : lancer plusieurs têtes en parallèle. Chaque tête a ses propres W_Q, W_K, W_V, donc chacune finit avec un motif d’attention différent. On concatène leurs sorties, on projette le résultat, on passe à la suite. C’est l’attention multi-têtes.

Ce chapitre introduit aussi les connexions résiduelles et la LayerNorm — les deux pièces de plomberie qui permettent d’empiler réellement des couches d’attention sans que le réseau s’effondre à l’entraînement. Elles ressemblent à des détails ajoutés après coup ; elles sont porteuses. En local, tu ajouteras ces opérations qui préservent la forme aux pièces que tu as déjà écrites.

Même phrase jouet qu’au chapitre 8 (the cat sat on mat), mais avec d_model = 8 et H = 4 têtes, donc d_head = 2 pour chaque tête.

1. Combiner plusieurs têtes

Tu as déjà une tête d’attention qui fonctionne. Le chapitre a pré-calculé 4 têtes différentes (des graines aléatoires différentes pour W_Q, W_K, W_V) et te donne la sortie de chacune sous forme de matrice [seq_len × d_head]. La cellule les combine.

La recette standard :

1. Concaténer la sortie de chaque tête sur l’axe des features. La ligne de chaque token devient [head₀, head₁, head₂, head₃], de longueur H × d_head = d_model.
2. Projeter la matrice concaténée à travers un W_O appris, de forme [d_model × d_model].

La projection de sortie laisse le modèle décider comment mélanger les têtes. Sans elle, les sorties des têtes seraient simplement collées bout à bout, sans aucune chance d’interagir.

// 'headOutputs' is an array of 4 head outputs, each [seq_len × d_head].
// 'W_O' is the output projection matrix, shape [d_model × d_model].
// The script defines matmul(A, B).
//
// 1. Concatenate every head's output along the feature axis.
//    Each token's row becomes [head0_out..., head1_out..., ...].
// 2. Project the concatenated matrix through W_O.
// 3. Return the [seq_len × d_model] result.
return [];

Le résultat a la même forme que l’entrée — [seq_len × d_model] — mais la nouvelle représentation de chaque token reflète désormais quatre motifs d’attention différents, mélangés.

2. Inspecter ce que les têtes apprennent vraiment

On a affirmé que les têtes voient des choses différentes. Vérifions. Les quatre matrices d’attention pré-calculées sont affichées ci-dessous en heatmaps. Certaines sont nettes (quelques cellules dominent chaque ligne) ; d’autres sont diffuses (la masse est répartie uniformément sur la ligne).

Une façon standard de quantifier la « concentration » est l’entropie : H = -Σ p log p. Une entropie basse veut dire que la tête se concentre sur peu de tokens. L’entropie maximale pour une ligne de 5 tokens vaut log(5) ≈ 1,61.

Calcule l’entropie moyenne par tête.

// 'headAttentions' is an array of 4 attention matrices, each [seq_len × seq_len].
// Compute the average Shannon entropy of each head's attention rows:
//   row_entropy = -Σ p_i * log(p_i)
//   head_entropy = mean(row_entropy over rows)
//
// Low entropy → the head focuses on a few tokens (concentrated pattern).
// High entropy → it spreads attention evenly (no clear preference).
// Return an array of 4 entropies, one per head.
return [];

Tu devrais voir des différences nettes entre les têtes. Certaines finissent presque uniformes (entropie haute — la tête ne préfère rien en particulier). D’autres ont des pics marqués (entropie basse — la tête a décidé de se concentrer). Dans un vrai transformer entraîné, on trouve des têtes spécialisées : têtes de copie, têtes d’induction, têtes qui suivent les noms propres, têtes qui suivent la syntaxe. On a nommé des motifs qu’on a trouvés, pas des motifs qu’on avait demandés.

3. Résiduel + LayerNorm

Si on branchait la sous-couche d’attention directement dans la sous-couche d’attention suivante, l’entraînement s’effondrerait. Deux raisons, toutes deux pénibles :

• Gradient qui disparaît. Chaque sous-couche comprime un peu le signal du gradient. Empiles-en 12 et le gradient tout en bas est microscopique.
• Dérive des représentations. Chaque sous-couche transforme les activations dans une géométrie différente. Empiles-en beaucoup et les magnitudes explosent ou s’écrasent à zéro.

La connexion résiduelle règle le premier problème : au lieu de output = sublayer(input), on écrit output = input + sublayer(input). Le gradient dispose maintenant d’un chemin propre pour revenir en arrière à travers l’addition, quoi que fasse la sous-couche. (C’est l’astuce qui a rendu les ResNets praticables en 2015, et elle est aujourd’hui dans tous les modèles profonds.)

La LayerNorm règle le second : après l’addition résiduelle, on normalise la ligne de chaque token pour qu’elle ait une moyenne de 0 et un écart-type de 1. La couche suivante voit des activations d’échelle connue, quoi qu’il se soit passé avant.

Ensemble, la recette par sous-couche est :

Exécute-la. Le chapitre donne à la cellule l’input d’origine (X) et sublayerOutput (la sortie multi-têtes de la cellule 1), plus un petit eps pour la stabilité numérique dans la division par l’écart-type.

// 'input' is the original X (the input to the attention sublayer).
// 'sublayerOutput' is what came out of the attention sublayer.
// 'eps' is a small constant (1e-5) for numerical stability.
//
// 1. Residual: add input + sublayerOutput element-wise.
// 2. LayerNorm: for each row, subtract its mean and divide by its
//    standard deviation (with the eps added inside the sqrt).
//
// Return the normalized [seq_len × d_model] matrix. Each row should
// have mean ≈ 0 and std ≈ 1.
return sublayerOutput;

Regarde les statistiques de ligne. La moyenne de chaque ligne doit valoir 0 (à la précision des flottants près) et son écart-type doit être très proche de 1. C’est l’invariant que la LayerNorm offre à la sous-couche suivante : chaque token, à chaque couche, arrive à la même échelle.

Pourquoi ce chapitre compte

Nous avons maintenant toutes les pièces d’un bloc transformer :

• L’attention (chapitre 8) : aller chercher de l’information à d’autres positions.
• Le multi-têtes (ce chapitre) : faire tourner plusieurs routes d’attention en parallèle.
• Résiduel + LayerNorm (ce chapitre) : la connectivité et la normalisation qui permettent d’empiler beaucoup de blocs.

Le chapitre suivant les assemble en un bloc transformer complet, puis en empile plusieurs pour former l’architecture réelle.

4. Ajouter les helpers résiduels et de normalisation

Ajoute ces helpers à llm/nn.py :

import math
 
 
def add(a: Matrix, b: Matrix) -> Matrix:
    return [
        [x + y for x, y in zip(row_a, row_b)]
        for row_a, row_b in zip(a, b)
    ]
 
 
def layer_norm(x: Matrix, eps: float = 1e-5) -> Matrix:
    out: Matrix = []
    for row in x:
        mean = sum(row) / len(row)
        var = sum((value - mean) ** 2 for value in row) / len(row)
        denom = math.sqrt(var + eps)
        out.append([(value - mean) / denom for value in row])
    return out

Ces helpers sont petits parce que leur rôle est structurel :

• add est la connexion résiduelle. Elle garde l’ancienne représentation et y ajoute le changement proposé par la sous-couche.
• layer_norm travaille ligne par ligne, donc chaque token est normalisé indépendamment des autres.
• mean recentre les features d’un token autour de zéro.
• var mesure à quel point ces features sont dispersées.
• Diviser par sqrt(var + eps) donne à la couche suivante une échelle prévisible. eps évite la division par zéro.

Tu as maintenant l’invariant dont dépend un transformer : chaque sous-couche accepte une matrice et retourne une matrice de la même forme, pour que le flux résiduel puisse continuer à circuler.

Pour aller plus loin

• Vaswani et al., « Attention Is All You Need » (2017). Le papier original du transformer. L’architecture domine toujours, neuf ans plus tard.
• transformer-circuits.pub d’Anthropic — les écrits du labo d’interprétabilité mécanistique sur ce que font réellement les têtes.
• He et al., « Deep Residual Learning » (2015) — le papier qui a mis les connexions résiduelles sur la carte.

Prochaine étape : le bloc transformer complet — combiner tout ce qu’on a construit dans l’unité qu’on empile réellement.