Les mathématiques de l'attention

Ce que vous saurez faire

Écrivez l’attention d’un transformeur avec quelques produits matriciels. Voyez ensuite le rôle de chaque partie de la formule : la mise à l’échelle, le masque, les têtes et le coût. Chaque étape comprend un petit exemple à vérifier à la main, une vérification dans PyTorch et une illustration.

Carnet JupyterConstruire l'attention à partir des formules, dans un carnetColab · Kaggle · GitHub · en anglais

Le carnet (notebook) Jupyter écrit attention() et l’attention à têtes multiples à partir des formules, puis les compare à PyTorch. Il mesure la variance de q · k pour plusieurs tailles de clés et chronomètre la matrice des scores sur votre machine. Il ne nécessite aucun téléchargement de données ni aucun entraînement. Environ 30 minutes.

Ouvrir dans ColabOuvrir dans KaggleVoir sur GitHub

Le carnet est en anglais. Colab et Kaggle ouvrent votre propre copie. Il faut un compte Google ou Kaggle gratuit. Vous êtes bloqué ? Consultez le carnet de solution.

Avant de commencerCe qu’il faut savoir et avoir sous la mainEnviron 55 minutes

Ce que vous devez savoir

• Les requêtes, les clés, les valeurs, les scores, softmax et la somme pondérée, dans L’attention étape par étape. Cette page part de son exemple détaillé.
• L’attention à têtes multiples, dans Les blocs du transformeur, et les deux masques, dans Position et masquage.
• Le produit scalaire, dans Les opérations sur les vecteurs du cours de mathématiques, et sa section La multiplication matricielle : plusieurs produits scalaires à la fois.
• La variance et l’écart-type, dans Décrire un jeu de données, et la dérivée comme pente, dans Fonctions et variation.

Ce qu’il vous faut sur votre ordinateur

• Rien pour la page. Le carnet s’ouvre dans Colab ou Kaggle avec PyTorch déjà installé.

Testé avec Python 3.12 et PyTorch 2.14.1 sur le processeur (macOS).

Cette leçon est facultative : elle n’est pas requise pour obtenir le certificat.

L’idée, en un paragraphe

Vous avez vu l’idée dans L’attention étape par étape. Cette page présente les mathématiques. Chaque mot pose une question (sa requête) et la compare à ce que chaque mot offre (sa clé). La comparaison est un produit scalaire, divisé par une racine carrée. Softmax transforme les scores en poids dont la somme vaut 1. Le nouveau vecteur du mot est la moyenne pondérée de ce que les mots transmettent (leurs valeurs). Dans l’exemple détaillé, « damaged on arrival » (endommagé à l’arrivée), chaque vecteur contient deux nombres. Le mot « damaged » lit surtout « arrival », avec un poids de 0,728.

La question à laquelle cette page répond

La leçon sur l’attention vous a donné une recette. Cette page demande pourquoi la recette contient exactement ces parties. Pourquoi un produit scalaire? Pourquoi diviser par √d_k, et non par d_k ou par rien? Pourquoi un masque ajoute-t-il moins l’infini, au lieu de fixer un poids à 0? Pourquoi plusieurs petites têtes plutôt qu’une grande? Et pourquoi l’attention devient-elle si coûteuse pour les longs textes? Toute la recette tient sur une ligne, tirée de l’article de 2017 Attention Is All You Need :

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}} + M\right) V$$

Ici, Q, K et V sont les matrices des requêtes, des clés et des valeurs, avec une ligne par jeton (token). d_k est le nombre de nombres dans une clé (la taille de la clé). M est le masque, une matrice de zéros et de moins l’infini. La leçon sur l’attention utilisait d pour la taille de la clé. À partir d’ici, d désigne le nombre de nombres par jeton et d_k, le nombre par clé. Ces deux nombres diffèrent quand il y a plusieurs têtes. La suite de cette page décompose la formule étape par étape, avec la forme inscrite à côté de chaque matrice.

Étape 1 : d’une requête à toutes les requêtes

La formule. Commencez une étape avant la leçon sur l’attention, avec les vecteurs des jetons eux-mêmes. Placez-les dans une matrice X de n lignes (une par jeton) et de d colonnes (les nombres par jeton). Trois matrices de poids apprises transforment X en requêtes, clés et valeurs :

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V \qquad (n, d) \times (d, d_k) \to (n, d_k)$$

En mots : chaque ligne de Q est la requête d’un jeton, obtenue à partir de son vecteur avec la même matrice de poids. Un seul produit matriciel crée les requêtes de tous les jetons à la fois.

Un petit exemple à vérifier. Reprenez les trois mots, avec maintenant d = 4 nombres chacun, et trois matrices de poids fictives de 4 × 2. Elles sont choisies pour que Q, K et V correspondent exactement aux requêtes, clés et valeurs de la leçon sur l’attention :

from_x.py

import math
import torch

X = torch.tensor([[1.0, 0.0, 0.0, 1.0],      # damaged
                  [0.0, 1.0, 0.0, 0.0],      # on
                  [0.0, 0.0, 1.0, 0.0]])     # arrival
W_Q = torch.tensor([[0.5, 0.5], [0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
W_K = torch.tensor([[1.0, 0.0], [0.0, 0.5], [0.0, 2.0], [0.0, 0.0]])
W_V = torch.tensor([[1.0, 0.0], [0.0, 0.0], [0.0, 1.0], [1.0, 0.0]])
Q, K, V = X @ W_Q, X @ W_K, X @ W_V          # (3, 4) @ (4, 2) -> (3, 2)
print("X", tuple(X.shape), "-> Q, K, V", tuple(Q.shape))
print(Q)
print(K)
print(V)

La sortie :

X (3, 4) -> Q, K, V (3, 2)
tensor([[0.5000, 1.5000],
        [0.0000, 0.0000],
        [1.0000, 0.0000]])
tensor([[1.0000, 0.0000],
        [0.0000, 0.5000],
        [0.0000, 2.0000]])
tensor([[2., 0.],
        [0., 0.],
        [0., 1.]])

Vérifiez la ligne 1 à la main. Le vecteur de « damaged » est (1, 0, 0, 1). Il sélectionne donc les lignes 1 et 4 de W_Q et les additionne : (0,5, 0,5) + (0, 1) = (0,5, 1,5). C’est la requête de « damaged » dans la leçon sur l’attention.

Tous les scores à la fois. Le score de la requête i avec la clé j est un produit scalaire. Sous forme matricielle, c’est la ligne i de Q multipliée par la colonne j de Kᵀ, les clés transposées. Un seul produit donne donc tous les scores :

$$S = Q K^\top, \qquad S_{ij} = q_i \cdot k_j = \sum_{t=1}^{d_k} q_{it}\, k_{jt} \qquad (n, d_k) \times (d_k, n) \to (n, n)$$

scores.py

S = Q @ K.T                                  # (3, 2) @ (2, 3) -> (3, 3)
print(S)
print((S / math.sqrt(2)).round(decimals=3))

La sortie :

tensor([[0.5000, 0.7500, 3.0000],
        [0.0000, 0.0000, 0.0000],
        [1.0000, 0.0000, 0.0000]])
tensor([[0.3540, 0.5300, 2.1210],
        [0.0000, 0.0000, 0.0000],
        [0.7070, 0.0000, 0.0000]])

La ligne 1 contient les trois scores de « damaged » : 0,5, 0,75 et 3,0. La leçon sur l’attention calculait ces mêmes nombres un à un. Pourquoi un produit scalaire? Il est grand quand la requête et la clé pointent dans la même direction, et il ne coûte que d_k multiplications. Un produit matriciel calcule les n × n scores en une seule opération rapide.

Parcourez toute la formule ci-dessous. Chaque matrice affiche sa forme, et la ligne de « damaged » est indiquée. Faites d’abord votre prédiction.

Explorer

Parcourir l'attention sous forme matricielle

L’entrée X et les matrices de poids

Étape

1 sur 7

Forme du résultat

(3, 4)

Poids de « damaged »

pas encore

X (3, 4)

	1	2	3	4
damaged	1	0	0	1
on	0	1	0	0
arrival	0	0	1	0

W_Q (4, 2)

	1	2
0,5	0,5
0	0
1	0
0	1

W_K (4, 2)

	1	2
1	0
0	0,5
0	2
0	0

W_V (4, 2)

	1	2
1	0
0	0
0	1
1	0

Chaque ligne de X est un jeton avec d = 4 nombres : n = 3 jetons, donc X est (3, 4). Les trois matrices de poids sont (4, 2) : elles transforment 4 nombres en d_k = 2. Dans un vrai modèle, elles sont apprises; ici, elles sont inventées.

Toute la formule sur 3 jetons, avec d = 4 et d_k = 2. Les matrices sont fictives, pour que Q, K et V correspondent à l'exemple détaillé de la leçon sur l'attention. Tous les autres nombres sont calculés à partir de ces matrices et vérifiés avec PyTorch.

Description textuelle de l’explorateur

L’explorateur présente la formule d’attention en sept étapes. Il affiche une matrice pour chaque entrée et résultat, sa forme et la ligne de « damaged ». Étape 1 : X a la forme (3, 4), avec une ligne par jeton : damaged (1, 0, 0, 1), on (0, 1, 0, 0), arrival (0, 0, 1, 0); W_Q, W_K et W_V ont la forme (4, 2). Étape 2 : Q = X W_Q, K = X W_K et V = X W_V ont la forme (3, 2). Q contient les lignes (0,5, 1,5), (0, 0) et (1, 0); K contient (1, 0), (0, 0,5) et (0, 2); V contient (2, 0), (0, 0) et (0, 1). Étape 3 : S = Q Kᵀ a la forme (3, 3), avec les lignes (0,5, 0,75, 3), (0, 0, 0) et (1, 0, 0). Étape 4 : S divisé par √2 contient les lignes (0,354, 0,53, 2,121), (0, 0, 0) et (0,707, 0, 0). Étape 5 : on ajoute le masque M. Sans masque, M ne contient que des zéros. Avec le masque causal, M contient −∞ au-dessus de la diagonale. Étape 6 : softmax, ligne par ligne. Sans masque, les poids sont (0,124, 0,148, 0,728), (0,333, 0,333, 0,333) et (0,503, 0,248, 0,248). Avec le masque causal, ils sont (1, 0, 0), (0,5, 0,5, 0) et (0,503, 0,248, 0,248). La somme de chaque ligne vaut 1. Étape 7 : O = A V a la forme (3, 2). Sans masque, ses lignes sont (0,248, 0,728), (0,667, 0,333) et (1,007, 0,248). Avec le masque causal, elles sont (2, 0), (1, 0) et (1,007, 0,248). Commandes : un bouton pour chaque étape, Back (précédent) et Next (suivant), et un choix entre aucun masque et le masque causal.

Le tableau rassemble les formes. Lisez-le comme une liste de vérification : les dimensions intérieures de chaque produit doivent correspondre, et le résultat conserve les dimensions extérieures.

Étape	Formule	Exemple	En général
Entrée	X	(3, 4)	(n, d)
Requêtes, clés, valeurs	X W_Q, X W_K, X W_V	(3, 4) × (4, 2) → (3, 2)	(n, d_k)
Scores	Q Kᵀ	(3, 2) × (2, 3) → (3, 3)	(n, n)
Mise à l’échelle et masquage	S / √d_k + M	(3, 3)	(n, n)
Poids	softmax, ligne par ligne	(3, 3), la somme de chaque ligne vaut 1	(n, n)
Sortie	A V	(3, 3) × (3, 2) → (3, 2)	(n, d_v)

Ici, d_v est la taille d’une valeur. Dans un transformeur, d_v = d_k, donc la sortie a la forme de Q.

Étape 2 : pourquoi diviser par √d_k

La leçon sur l’attention disait que les grands scores rendent softmax extrême. Vous pouvez maintenant mesurer leur grandeur et déduire le facteur exact.

La question. Avant l’entraînement, les nombres des requêtes et des clés sont à peu près aléatoires, autour de 0. Leur dispersion est d’environ 1. Quelle est alors la grandeur d’un score typique q · k, et comment augmente-t-elle avec la taille de clé d_k?

La démonstration, en trois lignes. Supposons que chaque nombre q_t et k_t est indépendant, avec une moyenne de 0 et une variance de 1. Le produit de deux de ces nombres indépendants a aussi une moyenne de 0. Sa variance est E[q_t²] E[k_t²] = 1 × 1 = 1. Le score additionne d_k de ces produits, et les variances des termes indépendants s’additionnent :

$$\operatorname{Var}(q \cdot k) = \operatorname{Var}\Big(\sum_{t=1}^{d_k} q_t k_t\Big) = \sum_{t=1}^{d_k} \operatorname{Var}(q_t k_t) = \sum_{t=1}^{d_k} \mathbb{E}[q_t^2]\,\mathbb{E}[k_t^2] = d_k$$

L’écart-type d’un score est donc √d_k. Diviser chaque score par √d_k donne la variance d_k / (√d_k)² = 1, quelle que soit la taille de clé. Voilà la raison de la racine carrée. Diviser par d_k réduirait trop les scores (variance 1 / d_k). Sans division, ils augmentent avec d_k.

Une vérification en code. Tirez 20 000 requêtes et clés aléatoires pour cinq tailles de clés, puis mesurez :

variance.py

torch.manual_seed(0)
for d_k in (4, 16, 64, 256, 1024):
    q = torch.randn(20000, d_k, dtype=torch.float64)       # 20,000 random queries ...
    k = torch.randn(20000, d_k, dtype=torch.float64)       # ... and keys, every number N(0, 1)
    dots = (q * k).sum(1)                                   # 20,000 dot products q . k
    Qs = torch.randn(2000, d_k, dtype=torch.float64)        # 2,000 queries, 32 keys each
    Ks = torch.randn(2000, 32, d_k, dtype=torch.float64)
    raw = torch.einsum("nd,nkd->nk", Qs, Ks)                # (2000, 32) scores
    top_raw = torch.softmax(raw, 1).max(1).values.mean()
    top_scaled = torch.softmax(raw / math.sqrt(d_k), 1).max(1).values.mean()
    print(f"d_k {d_k:5}  Var(q.k) {dots.var():8.2f}  after / sqrt(d_k) {(dots / math.sqrt(d_k)).var():.3f}"
          f"  largest weight: unscaled {top_raw:.3f}, scaled {top_scaled:.3f}")

La sortie :

d_k     4  Var(q.k)     4.03  after / sqrt(d_k) 1.006  largest weight: unscaled 0.335, scaled 0.152
d_k    16  Var(q.k)    15.84  after / sqrt(d_k) 0.990  largest weight: unscaled 0.637, scaled 0.165
d_k    64  Var(q.k)    63.69  after / sqrt(d_k) 0.995  largest weight: unscaled 0.822, scaled 0.163
d_k   256  Var(q.k)   256.00  after / sqrt(d_k) 1.000  largest weight: unscaled 0.903, scaled 0.164
d_k  1024  Var(q.k)  1024.99  after / sqrt(d_k) 1.001  largest weight: unscaled 0.953, scaled 0.164

torch.randn tire des nombres dans la loi normale de moyenne 0 et de variance 1, notée N(0, 1). einsum calcule les 32 produits scalaires de chaque requête avec ses propres clés. La variance mesurée vaut 4,03, 15,84, 63,69, 256,00 et 1 024,99 : d_k, comme le prédit la démonstration. Après la division, elle vaut 1 chaque fois. Le graphique montre les mêmes nombres sur des axes logarithmiques, où « variance = d_k » est une droite.

• Var(q · k), sans mise à l'échelle
• Var(q · k / √d_k), avec mise à l'échelle

La variance mesurée de q · k pour des requêtes et des clés aléatoires (20 000 paires par taille de clé). Sans mise à l'échelle, elle vaut d_k; après division par √d_k, elle vaut 1 pour chaque taille de clé.

Description textuelle du graphique

Un graphique en courbes avec des axes logarithmiques. L’axe des x représente la taille de clé d_k : 4, 16, 64, 256 et 1 024. L’axe des y représente la variance du score. La courbe sans mise à l’échelle monte en ligne droite : 4,03, 15,84, 63,69, 256,00 et 1 024,99. La variance vaut donc d_k. La courbe avec mise à l’échelle, q · k divisé par √d_k, reste horizontale à 1 : 1,006, 0,990, 0,995, 1,000 et 1,001.

L’effet des grands scores sur softmax. Le code ci-dessus a aussi transmis les 32 scores de chaque requête à softmax. Sans mise à l’échelle, le plus grand poids passe de 0,335 pour d_k = 4 à 0,953 pour d_k = 1 024. Pour d_k = 1 024, 73 % des lignes donnent plus de 0,99 à une seule clé. Softmax devient presque un choix strict d’une seule clé : il est saturé. Avec la mise à l’échelle, le plus grand poids reste près de 0,16 pour chaque taille de clé. Déplacez le curseur et comparez les deux rangées de barres.

Explorer

Changer la taille de clé et observer softmax

Var(q · k)

4,03 (écart type 2,01)

Plus grand poids, sans → avec mise à l’échelle

0,335 → 0,152

Lignes avec un poids au-dessus de 0,99

0,1 % → 0,0 %

Taille du gradient du softmax

0,292 → 0,234

À d_k = 4, les scores q · k ont une variance de 4,03, donc un score typique se situe à environ ±2,01 de 0. Sans mise à l’échelle, le softmax répartit encore son poids (plus grand poids 0,335 en moyenne). Divisée par √4 = 2, la variance est de 1,01 et le plus grand poids de 0,152, à peu près comme pour chaque d_k.

Sans mise à l’échelle : softmax(q · k) max. 0,164

1

Avec mise à l’échelle : softmax(q · k / √4) max. 0,084

1

Une requête aléatoire et 32 clés aléatoires à d_k = 4, chaque nombre tiré de N(0, 1). Les clés sont triées par score. Les valeurs ci-dessus donnent les moyennes sur 2 000 requêtes de ce type.

Une requête aléatoire et 32 clés aléatoires par taille de clé, avec chaque nombre tiré de N(0, 1). L'affichage donne les moyennes sur 2 000 de ces requêtes, issues de la même exécution que le code ci-dessus. La grandeur du gradient provient du code de gradient ci-dessous.

Description textuelle de l’explorateur

L’explorateur a un curseur pour la taille de clé d_k : 4, 16, 64, 256 ou 1 024. Il montre deux rangées de 32 barres : les poids softmax d’une requête aléatoire sur 32 clés aléatoires. Les poids sont affichés sans et avec mise à l’échelle, du plus grand au plus petit. Pour d_k = 4, les deux rangées sont étalées : le plus grand poids sans mise à l’échelle vaut 0,164, et celui avec mise à l’échelle vaut 0,084. Pour d_k = 64, la rangée sans mise à l’échelle a une barre à 0,797 et une à 0,200. Les autres sont près de 0; la rangée avec mise à l’échelle descend encore doucement à partir de 0,149. Pour d_k = 1 024, la rangée sans mise à l’échelle est une seule barre à 1,000. La rangée avec mise à l’échelle commence à 0,243. L’affichage montre les moyennes mesurées sur 2 000 requêtes. La variance de q · k passe de 4,03 à 1 024,99. Le plus grand poids, sans et avec mise à l’échelle, vaut 0,335 et 0,152 pour d_k = 4. Il vaut 0,953 et 0,164 pour d_k = 1 024. La proportion de lignes ayant un poids supérieur à 0,99 vaut 0,1 % et 0 % pour d_k = 4. Elle vaut 73,2 % et 0 % pour d_k = 1 024. La grandeur du gradient softmax vaut 0,292 et 0,234 pour d_k = 4, et 0,057 et 0,240 pour d_k = 1 024.

Pourquoi un softmax saturé nuit à l’apprentissage. L’entraînement modifie un score en suivant le gradient de la perte, comme dans La descente de gradient. Le gradient passe par softmax, et la dérivée du poids softmax p_i par rapport au score s_j est :

$$\frac{\partial p_i}{\partial s_j} = p_i\,(\delta_{ij} - p_j), \qquad \delta_{ij} = 1 \text{ if } i = j, \text{ else } 0$$

En mots : chaque élément est un poids multiplié par un moins un poids, ou l’opposé du produit de deux poids. Quand un poids est près de 1 et tous les autres près de 0, chaque élément est près de 0. Par exemple, une ligne avec les poids 0,99 et 0,01 donne p(1 − p) = 0,99 × 0,01 = 0,0099. Presque aucun gradient ne revient aux scores, donc les requêtes et les clés apprennent très lentement. Le code mesure la grandeur de cette dérivée, en moyenne sur 2 000 lignes de 32 clés :

gradient.py

def softmax_gradient_size(scores):
    """Mean size (Frobenius norm) of d softmax / d scores = diag(p) - p p^T, over the rows."""
    p = torch.softmax(scores, dim=-1)
    J = torch.diag_embed(p) - p.unsqueeze(-1) * p.unsqueeze(-2)
    return J.norm(dim=(-2, -1)).mean().item()

torch.manual_seed(1)
for d_k in (4, 16, 64, 256, 1024):
    raw = torch.einsum("nd,nkd->nk", torch.randn(2000, d_k, dtype=torch.float64),
                       torch.randn(2000, 32, d_k, dtype=torch.float64))
    print(f"d_k {d_k:5}  gradient size: unscaled {softmax_gradient_size(raw):.4f},"
          f" scaled {softmax_gradient_size(raw / math.sqrt(d_k)):.4f}")

La sortie :

d_k     4  gradient size: unscaled 0.2918, scaled 0.2337
d_k    16  gradient size: unscaled 0.2875, scaled 0.2388
d_k    64  gradient size: unscaled 0.1948, scaled 0.2395
d_k   256  gradient size: unscaled 0.1129, scaled 0.2405
d_k  1024  gradient size: unscaled 0.0570, scaled 0.2403

J contient toutes les dérivées ci-dessus pour une ligne. Sa norme de Frobenius est la racine carrée de la somme des carrés de ses éléments. C’est un nombre qui indique « combien de gradient passe ». Avec la mise à l’échelle, la grandeur reste près de 0,24 pour chaque taille de clé. Sans mise à l’échelle, elle tombe à 0,057 pour d_k = 1 024, soit environ 4 fois moins. Quand d_k est petit, les scores sans mise à l’échelle ne sont pas encore extrêmes. La division compte donc peu dans ce cas. Le tout petit transformeur de ce cours utilise d_k = 16 et divise par 4.

Étape 3 : softmax ligne par ligne, puis la somme pondérée

La formule. Écrivez Z = S / √d_k + M pour les scores mis à l’échelle et masqués. Softmax traite chaque ligne de Z séparément, donc la somme de chaque ligne de la matrice des poids A vaut 1. Chaque ligne de sortie est ensuite une somme pondérée des lignes de valeurs :

$$A_{ij} = \frac{e^{Z_{ij}}}{\sum_{l=1}^{n} e^{Z_{il}}}, \qquad o_i = \sum_{j=1}^{n} A_{ij}\, v_j \quad\Longleftrightarrow\quad O = A V$$

En mots : la ligne i de A indique combien le jeton i lit chaque jeton j. La ligne i de O mélange les valeurs avec ces poids. Ici, o_i et v_j sont des lignes de O et V. En code, toute la formule tient en quelques lignes. Voici la fonction lisible utilisée dans la suite de la leçon :

attention.py

def attention(Q, K, V, mask=None):
    """softmax(Q K^T / sqrt(d_k) + mask) V.  Q: (..., n, d_k), K: (..., m, d_k), V: (..., m, d_v)."""
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)       # (..., n, m)
    if mask is not None:
        scores = scores + mask                               # -inf where a query may not look
    weights = torch.softmax(scores, dim=-1)                  # every row sums to 1
    return weights @ V, weights                              # (..., n, d_v), (..., n, m)

out, weights = attention(Q, K, V)
print(weights.round(decimals=3))
print(out.round(decimals=3))

La sortie :

tensor([[0.1240, 0.1480, 0.7280],
        [0.3330, 0.3330, 0.3330],
        [0.5030, 0.2480, 0.2480]])
tensor([[0.2480, 0.7280],
        [0.6670, 0.3330],
        [1.0070, 0.2480]])

Le ... dans les formes signifie « toutes les dimensions qui précèdent celles-ci » : un lot, des têtes ou les deux. K.transpose(-2, -1) échange les deux dernières dimensions : c’est donc Kᵀ pour chaque élément d’un lot. dim=-1 fait travailler softmax le long de chaque ligne. Les clés peuvent venir d’une autre séquence de m jetons, ce que la prochaine leçon utilise pour la traduction. La ligne 1 correspond au calcul à la main de la leçon sur l’attention : poids 0,124, 0,148 et 0,728, sortie (0,248, 0,728).

Étape 4 : ajouter les masques avant softmax

La formule. Un masque est une matrice M de la même forme que les scores. Elle contient 0 là où une requête peut lire une clé, et −∞ là où elle ne le peut pas. On l’ajoute aux scores mis à l’échelle, avant softmax :

$$A = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right), \qquad M_{ij} = \begin{cases} 0 & \text{token } i \text{ may read token } j \\ -\infty & \text{otherwise} \end{cases}$$

En mots : une clé cachée reçoit le score −∞, et e^−∞ = 0, donc son poids vaut exactement 0. Les clés visibles se partagent tout le poids de 1. Deux masques sont courants. Le masque causal contient −∞ au-dessus de la diagonale, donc un jeton ne lit jamais les jetons qui le suivent. Le masque de remplissage contient −∞ dans les colonnes des jetons de remplissage, donc aucun jeton ne lit les espaces de remplissage. Position et masquage présentait les deux sur le tout petit transformeur; les voici sur les trois mots :

masks.py

causal = torch.triu(torch.full((3, 3), float("-inf")), diagonal=1)   # -inf above the diagonal
print(causal)
out, weights = attention(Q, K, V, mask=causal)
print(weights.round(decimals=3))
padding = torch.tensor([0.0, 0.0, float("-inf")])     # pretend "arrival" is padding: hide column 3
out, weights = attention(Q, K, V, mask=padding)       # the same row is added to every row of scores
print(weights.round(decimals=3))

La sortie :

tensor([[0., -inf, -inf],
        [0., 0., -inf],
        [0., 0., 0.]])
tensor([[1.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.0000],
        [0.5030, 0.2480, 0.2480]])
tensor([[0.4560, 0.5440, 0.0000],
        [0.5000, 0.5000, 0.0000],
        [0.6700, 0.3300, 0.0000]])

torch.triu(…, diagonal=1) conserve la partie d’une matrice au-dessus de la diagonale et fixe le reste à 0. Le masque de remplissage est une seule ligne; PyTorch l’ajoute à chaque ligne des scores. C’est la diffusion, présentée dans Tenseurs et formes. La figure montre les deux masques et les poids qu’ils donnent.

Les masques sont des matrices de 0 et de −∞ ajoutées aux scores mis à l'échelle. Après softmax, chaque cellule −∞ a un poids de 0 (orange), et la somme des poids visibles de chaque ligne vaut 1.

Description textuelle du graphique

Deux rangées de deux grilles, chacune de 3 sur 3, pour les jetons damaged, on et arrival. Les lignes représentent le jeton qui lit; les colonnes, le jeton qu’il lit. Première rangée : le masque causal M contient 0 sur et sous la diagonale, et −∞ au-dessus. Après son ajout et softmax, les poids sont : damaged 1, 0, 0; on 0,5, 0,5, 0; arrival 0,503, 0,248, 0,248. Deuxième rangée : le masque de remplissage, où « arrival » sert de remplissage, contient −∞ dans toute la troisième colonne. Les poids sont : damaged 0,456, 0,544, 0; on 0,5, 0,5, 0; arrival 0,67, 0,33, 0. Chaque cellule cachée a un poids d’exactement 0, et la somme des poids visibles de chaque ligne vaut 1.

Pourquoi −∞, et non 0? Ajouter 0 à un score ne change rien. Et un score de 0 n’est pas « caché » : softmax le transforme en e^0 = 1 avant de diviser par le total. Essayez les deux idées sur la ligne de « damaged » :

why_inf.py

scaled = torch.tensor([0.354, 0.530, 2.121])          # row "damaged", already divided by sqrt(2)
print("add -inf:      ", torch.softmax(scaled + torch.tensor([0.0, float("-inf"), float("-inf")]), 0).round(decimals=3))
print("set scores to 0:", torch.softmax(torch.tensor([0.354, 0.0, 0.0]), 0).round(decimals=3))

La sortie :

add -inf:       tensor([1., 0., 0.])
set scores to 0: tensor([0.4160, 0.2920, 0.2920])

Avec −∞, « damaged » ne lit que lui-même. Quand les scores sont fixés à 0, il donne encore 0,584 de son poids aux deux mots « cachés ». Une troisième idée, fixer les poids à 0 après softmax, échoue aussi. La somme de la ligne ne vaudrait plus 1, à moins de diviser à nouveau. Ajouter −∞ avant softmax fait les deux tâches en une étape. Cela conserve aussi la forme de la matrice, donc une seule opération rapide traite chaque ligne.

Étape 5 : l’attention multitête

La formule. Une tête donne à chaque jeton une ligne de poids. Elle ne peut donc suivre qu’un type de relation à la fois. L’attention multitête exécute h petites têtes en parallèle. Chaque tête i a ses propres projections, de taille d × d_k avec d_k = d / h. Les sorties des têtes sont réunies (concaténées), puis une autre matrice W_O les mélange :

$$\text{head}_i = \text{Attention}(X W_Q^{(i)},\, X W_K^{(i)},\, X W_V^{(i)}), \qquad \text{MultiHead}(X) = [\text{head}_1, \ldots, \text{head}_h]\, W_O$$

En mots : h attentions indépendantes, chacune sur une tranche de d_k nombres, puis une couche linéaire sur le résultat réuni. En pratique, les h matrices de requêtes des têtes se trouvent côte à côte dans une matrice W_Q de taille d × d. Un seul produit calcule donc toutes les requêtes. Un changement de forme sépare ensuite les d colonnes en h têtes de d_k nombres. Le schéma montre chaque forme, avec les nombres du tout petit transformeur pour le billet T-65052 (30 jetons, d = 64, h = 4).

Attention multitête : séparer en h têtes, calculer l'attention dans chaque tête, réunir les têtes et les mélanger avec W_O. La forme à l'entrée et à la sortie est (n, d).

Description textuelle du diagramme

Les données circulent de gauche à droite. X a la forme (n, d); pour le billet T-65052 dans le tout petit transformeur, c’est (30, 64). Il est multiplié par trois matrices de poids W_Q, W_K et W_V, chacune de forme (d, d). Cela donne Q, K et V, chacun de forme (n, d). Leurs d colonnes sont séparées en h têtes de d_k = d / h nombres. Chacun devient donc (h, n, d_k), soit (4, 30, 16) pour le billet. Chaque tête, de la tête 1 à la tête h, calcule sa propre attention : softmax de Q multiplié par K transposé, divisé par la racine carrée de d_k, plus le masque M, puis multiplié par V. Chaque tête a une matrice de poids (n, n) et une sortie (n, d_k). Les h sorties sont concaténées : h fois d_k colonnes donnent d colonnes, donc le résultat est (n, d). Une dernière matrice W_O, de forme (d, d), mélange les têtes. La sortie est (n, d), la même forme que X.

Le nombre de paramètres. Comptez les nombres appris. W_Q, W_K, W_V et W_O ont chacun la taille d × d et un biais de d nombres :

$$4 \cdot d^2 + 4 \cdot d \quad\Rightarrow\quad d = 64:\; 16{,}384 + 256 = 16{,}640, \qquad d = 512:\; 1{,}050{,}624$$

Le nombre de têtes n’apparaît pas dans la formule. Plus de têtes divisent les mêmes d colonnes en tranches plus nombreuses et plus petites. C’est pourquoi Les blocs du transformeur donnait 16 640 paramètres pour 1, 4 et 8 têtes. L’article de 2017 utilisait d = 512 et h = 8. Chacune de ses couches d’attention a donc 1 050 624 paramètres si chaque matrice a un biais, comme dans PyTorch. (Les formules de l’article ne montrent aucun biais; sans eux, le nombre est 4d² = 1 048 576.) Voici la couche multitête, construite à partir de attention() :

multihead.py

class MultiHead(torch.nn.Module):
    """h heads of size d_k = d / h. Q = X W_Q + b_Q (and K, V alike); heads side by side; output = concat W_O + b_O."""

    def __init__(self, d, h):
        super().__init__()
        self.d, self.h, self.d_k = d, h, d // h
        self.W_Q, self.W_K, self.W_V, self.W_O = (torch.nn.Linear(d, d) for _ in range(4))

    def split(self, X):                                      # (B, n, d) -> (B, h, n, d_k)
        B, n, _ = X.shape
        return X.view(B, n, self.h, self.d_k).transpose(1, 2)

    def forward(self, Xq, Xkv, mask=None):
        Q, K, V = self.split(self.W_Q(Xq)), self.split(self.W_K(Xkv)), self.split(self.W_V(Xkv))
        O, weights = attention(Q, K, V, mask)                # (B, h, n, d_k), (B, h, n, m)
        B, _, n, _ = O.shape
        O = O.transpose(1, 2).reshape(B, n, self.d)          # concatenate the heads
        return self.W_O(O), weights

torch.manual_seed(0)
Xb = torch.randn(3, 7, 64)                                   # 3 sequences, 7 tokens, d = 64
q, k, v = (torch.randn(3, 7, 16) for _ in range(3))          # for the one-head check below
mine = MultiHead(64, 4)
print("parameters:", sum(p.numel() for p in mine.parameters()),
      "| d = 512:", sum(p.numel() for p in MultiHead(512, 8).parameters()))
out, weights = mine(Xb, Xb)
print("out", tuple(out.shape), "weights", tuple(weights.shape))

La sortie :

parameters: 16640 | d = 512: 1050624
out (3, 7, 64) weights (3, 4, 7, 7)

view lit les mêmes nombres avec une nouvelle forme, (B, n, h, d_k). transpose(1, 2) place les têtes avant les jetons. Ensuite, attention() travaille sur chaque tête de chaque séquence à la fois, grâce aux ... dans ses formes. Après l’attention, les transformations inverses réunissent les têtes : (B, h, n, d_k) redevient (B, n, d). Les poids ont un carré n × n par tête : (3, 4, 7, 7).

Pourquoi plusieurs têtes? Chaque tête a ses propres W_Q et W_K. Chacune apprend donc sa propre idée d’une « bonne correspondance ». Une tête peut suivre les mots liés au paiement pendant qu’une autre suit la demande. Leurs poids sont des lignes softmax distinctes. Une relation ne retire donc pas de poids à une autre. Le coût est un d_k plus petit par tête : chaque tête compare les jetons sur 16 nombres plutôt que 64. Le vrai billet à la fin de cette page montre quatre têtes qui ont appris quatre motifs différents.

Étape 6 : le coût de n × n

La formule. Comptez le travail pour une tête. Q Kᵀ contient n × n éléments, chacun étant un produit scalaire de d_k nombres. A V demande encore n × n × d_k multiplications. Il faut aussi stocker la matrice des scores elle-même, soit n × n nombres :

$$\text{time} \propto n^2 \cdot d_k, \qquad \text{memory for the scores} = n^2 \times 4 \text{ bytes (float32)}$$

En mots : doublez les jetons, et le travail et la mémoire des scores quadruplent. Tout le reste d’un bloc, les projections et la couche à propagation avant, augmente seulement en proportion de n.

memory.py

for n in (128, 1024, 8192):
    cells = n * n                                            # one score per (query, key) pair
    print(f"n {n:5}: {cells:>11,} scores = {cells * 4 / 1e6:7.2f} MB in float32 (one head)")

La sortie :

n   128:      16,384 scores =    0.07 MB in float32 (one head)
n  1024:   1,048,576 scores =    4.19 MB in float32 (one head)
n  8192:  67,108,864 scores =  268.44 MB in float32 (one head)

À 8 192 jetons, la matrice des scores d’une tête occupe 268 Mo, pour une séquence dans une couche. La sortie softmax est une deuxième matrice de même taille. En comparaison, Q, K et V ensemble n’occupent que 6,3 Mo. Les graphiques montrent la mémoire et le temps mesurés sur un processeur partagé, pour une tête avec d_k = 64. Les valeurs sont les médianes de 5 passages. Les deux axes sont logarithmiques, donc la croissance en n² forme une droite : 4 fois plus chaque fois que n double.

• Matrice des scores
• Q, K et V ensemble

Mémoire pour une tête en float32. La matrice des scores n × n croît selon n²; les requêtes, les clés et les valeurs croissent seulement selon n.

Description textuelle du graphique

Un graphique linéaire avec des axes logarithmiques. L’axe des x représente le nombre de jetons n, de 128 à 8 192. La matrice des scores demande 0,07 Mo à 128 jetons, 1,05 Mo à 512, 4,19 Mo à 1 024, 16,78 Mo à 2 048, 67,11 Mo à 4 096 et 268,44 Mo à 8 192 : une droite à forte pente, avec 4 fois plus chaque fois que n double. Q, K et V ensemble doublent seulement chaque fois que n double, de 0,098 Mo à 6,3 Mo.

• attention() lisible
• Fonction fusionnée de PyTorch
• Croissance en n² à partir de n = 1 024

Temps pour une tête, lot de 1, d_k = 64, sur un processeur partagé (2 fils d'exécution, médiane de 5 passages) : à titre indicatif seulement. La ligne pointillée représente une croissance pure en n² à partir du temps de la fonction lisible à 1 024 jetons.

Description textuelle du graphique

Un graphique linéaire avec des axes logarithmiques. L’axe des x représente le nombre de jetons n, de 128 à 8 192; l’axe des y représente le temps en millisecondes. La fonction attention() lisible prend 0,035 ms à 128 jetons, 0,225 ms à 512, 0,849 ms à 1 024, 3,54 ms à 2 048, 14,8 ms à 4 096 et 60,0 ms à 8 192. La fonction fusionnée scaled_dot_product_attention de PyTorch prend 0,027, 0,144, 0,437, 1,70, 6,86 et 27,0 ms. Une ligne pointillée montre une croissance pure en n² à partir de 0,849 ms à 1 024 jetons. Au-dessus de 1 024 jetons, les deux lignes mesurées suivent sa pente. Sous environ 512 jetons, les coûts fixes comptent davantage et les lignes sont plus plates.

n	Scores	Matrice des scores	attention() lisible	Fonction fusionnée de PyTorch
128	16 384	0,07 Mo	0,035 ms	0,027 ms
1 024	1 048 576	4,19 Mo	0,849 ms	0,437 ms
8 192	67 108 864	268,44 Mo	60,0 ms	27,0 ms

De 1 024 à 8 192 jetons (8 fois plus), le temps est multiplié par 71. C’est proche du facteur 64 de n². La fonction scaled_dot_product_attention de PyTorch est environ deux fois plus rapide, car elle fusionne les étapes en une opération optimisée. Elle effectue tout de même le travail en n². Ces temps proviennent d’une seule machine partagée. Sur la vôtre, les nombres seront différents, mais pas la pente.

C’est pourquoi les modèles à long contexte ont besoin de techniques particulières. Certains noyaux ne stockent jamais toute la matrice des scores à la fois. Certaines couches limitent l’attention à une fenêtre de jetons voisins. Pendant la génération, le cache KV présenté dans Les coûts de l’inférence conserve les clés et les valeurs des jetons précédents. Elles ne sont donc pas calculées à nouveau.

Étape 7 : la vérification avec PyTorch

Une formule sur papier n’est utile que si le code la calcule. Le code ci-dessous compare donc les fonctions lisibles de cette page avec celles de PyTorch. attention() est comparée à F.scaled_dot_product_attention, et MultiHead à nn.MultiheadAttention, avec les mêmes poids copiés dans cette dernière :

check.py

F = torch.nn.functional
ref = torch.nn.MultiheadAttention(64, 4, batch_first=True)
with torch.no_grad():                                        # copy our weights into PyTorch's layer
    ref.in_proj_weight.copy_(torch.cat([mine.W_Q.weight, mine.W_K.weight, mine.W_V.weight]))
    ref.in_proj_bias.copy_(torch.cat([mine.W_Q.bias, mine.W_K.bias, mine.W_V.bias]))
    ref.out_proj.weight.copy_(mine.W_O.weight)
    ref.out_proj.bias.copy_(mine.W_O.bias)
    causal = torch.triu(torch.full((7, 7), float("-inf")), diagonal=1)
    future = torch.triu(torch.ones(7, 7, dtype=torch.bool), diagonal=1)        # PyTorch's layer wants True/False
    is_pad = torch.arange(7)[None, :] >= torch.tensor([7, 5, 3])[:, None]     # lengths 7, 5 and 3
    padding = torch.zeros(3, 1, 1, 7).masked_fill(is_pad[:, None, None, :], float("-inf"))
    one_head = [attention(q, k, v)[0] - F.scaled_dot_product_attention(q, k, v),
                attention(q, k, v, causal)[0] - F.scaled_dot_product_attention(q, k, v, is_causal=True),
                attention(q, k, v, padding[:, 0])[0] - F.scaled_dot_product_attention(q, k, v, attn_mask=padding[:, 0])]
    print("one head vs scaled_dot_product_attention:", f"{max(d.abs().max() for d in one_head):.1e}")
    biggest = 0.0
    for my_mask, theirs in [(None, {}), (causal, {"attn_mask": future}), (padding, {"key_padding_mask": is_pad}),
                            (causal + padding, {"attn_mask": future, "key_padding_mask": is_pad})]:
        o1, w1 = mine(Xb, Xb, my_mask)
        o2, w2 = ref(Xb, Xb, Xb, average_attn_weights=False, **theirs)
        biggest = max(biggest, (o1 - o2).abs().max().item(), (w1 - w2).abs().max().item())
    print("multi-head vs nn.MultiheadAttention:", f"{biggest:.1e}")

La sortie :

one head vs scaled_dot_product_attention: 2.4e-07
multi-head vs nn.MultiheadAttention: 1.8e-07

nn.MultiheadAttention stocke W_Q, W_K et W_V empilées dans une matrice, in_proj_weight, de forme (3d, d). Le code réunit donc les nôtres avec torch.cat. La vérification couvre l’absence de masque, le masque causal, le remplissage (longueurs de séquence 7, 5 et 3) et les deux masques ensemble. Elle compare les sorties et les poids. Le plus grand écart est de 2,4 × 10⁻⁷. C’est un arrondi : float32 conserve environ 7 chiffres significatifs, et les deux codes additionnent les mêmes nombres dans un ordre différent. En float64, la même vérification donne 4,4 × 10⁻¹⁶. Les formules de cette page correspondent donc à ce que PyTorch calcule.

Étape 8 : un vrai billet

La dernière vérification utilise le tout petit transformeur entraîné du cours et le billet de validation T-65052 : « Ce n’est pas un problème de livraison et je garde le souffleur à feuilles. Ma carte a été débitée deux fois pour la même commande. Pouvez-vous supprimer le débit en double? » L’équipe de Grace est celle des paiements, et le modèle prédit les paiements. Le code copie les poids entraînés du bloc 2 dans MultiHead et recalcule les poids d’attention :

ticket.py

from ticketnet import TinyTransformer, Vocab, load_split

DATA = data_dir()
train_rows, valid_rows = (load_split(DATA / f"{s}.csv") for s in ("train", "valid"))
vocab = Vocab([r["text"] for r in train_rows])
tf = TinyTransformer(len(vocab))
tf.load_state_dict(torch.load(WEIGHTS, weights_only=True))   # the course's trained model
tf.eval()
row = next(r for r in valid_rows if r["ticket_id"] == "T-65052")
ids = torch.tensor([vocab.encode(row["text"])])
tokens = [vocab.itos[i] for i in ids[0].tolist()]
with torch.no_grad():
    h = tf.embed(ids) + tf.position(torch.arange(ids.shape[1]))
    h = tf.blocks.layers[0](h)                               # block 1 as a whole
    block = tf.blocks.layers[1]
    x = block.norm1(h)                                       # block 2 normalises before attention
    pt = block.self_attn                                     # PyTorch's layer, trained
    ours = MultiHead(64, 4)
    for lin, W, b in zip((ours.W_Q, ours.W_K, ours.W_V), pt.in_proj_weight.chunk(3), pt.in_proj_bias.chunk(3)):
        lin.weight.copy_(W)
        lin.bias.copy_(b)
    ours.W_O.weight.copy_(pt.out_proj.weight)
    ours.W_O.bias.copy_(pt.out_proj.bias)
    _, w_ours = ours(x, x)
    _, w_pt = pt(x, x, x, average_attn_weights=False)
print("weights", tuple(w_ours.shape), "max difference", f"{(w_ours - w_pt).abs().max():.1e}")
q = tokens.index("not")
for head in range(4):
    top = w_ours[0, head, q].topk(2)
    print(f"head {head + 1}: 'not' reads", [(tokens[i], round(v.item(), 2)) for v, i in zip(top.values, top.indices)])

La sortie :

weights (1, 4, 30, 30) max difference 2.1e-07
head 1: 'not' reads [('twice', 0.5), ("'", 0.07)]
head 2: 'not' reads [('charge', 0.55), ('the', 0.09)]
head 3: 'not' reads [('remove', 0.5), ('order', 0.1)]
head 4: 'not' reads [('remove', 0.48), ('charge', 0.15)]

data_dir() et WEIGHTS pointent vers les données de billets du cours et le modèle enregistré. Les blocs du transformeur l’entraîne et l’enregistre. chunk(3) redécoupe la matrice empilée de PyTorch en W_Q, W_K et W_V. Notre formule reproduit les poids du modèle entraîné à 2,1 × 10⁻⁷ près. Sur les deux blocs, pour les sorties et les poids, le plus grand écart est de 9,5 × 10⁻⁷. Les cartes thermiques montrent les quatre têtes du bloc 2.

Les vrais poids d'attention du tout petit transformeur, bloc 2, billet T-65052 : une carte thermique de 30 × 30 par tête. Les lignes représentent les jetons qui lisent, les colonnes les jetons qu'ils lisent. Une couleur plus forte indique un poids plus grand (échelle de racine carrée, relative au plus grand poids de la tête). La ligne de « not » (pas) est encadrée.

Description textuelle du graphique

Quatre cartes thermiques, une pour chaque tête du bloc 2, chacune de 30 lignes sur 30 colonnes pour les 30 jetons du billet T-65052. La somme de chaque ligne est 1. Les têtes montrent des bandes verticales : la plupart des jetons lisent les mêmes quelques mots. La tête 1 a une forte bande sur « twice » (deux fois), qui reçoit en moyenne 0,355 du poids sur les 30 jetons lecteurs. La tête 2 a sa bande sur « charge » (débit), avec 0,208 en moyenne. La tête 3 répartit davantage son poids : « remove » (supprimer) reçoit 0,118 et « you » (vous), 0,116. La tête 4 a sa bande sur « remove » (0,274). La ligne encadrée de « not » lit twice à 0,50 dans la tête 1, charge à 0,55 dans la tête 2, remove à 0,50 dans la tête 3 et remove à 0,48 dans la tête 4. Aucune tête n’envoie « not » vers « delivery » (livraison), le mot juste après.

Les têtes ont appris des motifs différents : la tête 1 recueille « twice », la tête 2 « charge », la tête 4 « remove », et la tête 3 répartit son attention. Ce sont les « relations différentes » de l’étape 5, dans un vrai modèle. Pour chaque billet, bloc et tête, ouvrez l’explorateur dans L’attention étape par étape. Gardez son avertissement en tête : les poids ne sont pas des explications. Ils montrent où chaque jeton a pris son information, pas pourquoi le modèle a choisi les paiements.

Ce que cela signifie pour Larkfield, et la suite

Pour le routeur de Larkfield, les mathématiques expliquent trois faits pratiques. Un billet de 120 jetons donne à chaque tête une matrice des scores de 120 × 120, soit 14 400 nombres. C’est petit, donc l’attention coûte peu pour les billets courts. Le d_k = 16 du tout petit transformeur est assez petit pour que la mise à l’échelle compte moins que dans les grands modèles. Mais la conserver ne coûte rien. Les quatre têtes de chaque bloc sont aussi moins coûteuses qu’elles semblent. Elles coûtent les mêmes 16 640 paramètres qu’une tête de 64 nombres.

La formule est restée presque identique depuis 2017. Les changements concernent surtout ce qui l’entoure. Les modèles plus récents intègrent les positions au produit scalaire. Plusieurs têtes de requêtes partagent un plus petit nombre de têtes de clés et de valeurs. Les noyaux fusionnés évitent de stocker la matrice n × n. La prochaine leçon reconstruit le transformeur original à partir des éléments de cette page.

Exercice guidé : les formes d’un bloc pour un long billet

1. Le tout petit transformeur lit un billet de n = 120 jetons, avec d = 64 et h = 4. Écrivez d_k.
2. Écrivez la forme de Q après W_Q, puis après split (avec une dimension de lot de 1).
3. Écrivez la forme des poids d’attention d’un bloc et le nombre de scores qu’ils contiennent.
4. Calculez la mémoire de ces scores en float32, en ko.
5. Écrivez la forme de la sortie après la réunion des têtes et l’application de W_O.
6. Écrivez le nombre de paramètres de cette couche d’attention à partir de la formule.

Afficher les réponses attendues

1. d_k = d / h = 64 / 4 = 16.
2. Q = X W_Q est (1, 120, 64). Après split : (1, 4, 120, 16), quatre têtes de 16 nombres.
3. Un carré de 120 × 120 par tête : (1, 4, 120, 120), donc 4 × 14 400 = 57 600 scores.
4. 57 600 × 4 octets = 230 400 octets = 230,4 ko par bloc. La sortie softmax est une deuxième matrice de cette taille.
5. Après réunion : (1, 120, 4 × 16) = (1, 120, 64); après W_O : (1, 120, 64), la forme de l’entrée.
6. 4 × 64² + 4 × 64 = 16 640, le même nombre pour n’importe quel nombre de têtes.

Le carnet (notebook) Jupyter exécute MultiHead(64, 4) sur une entrée (1, 120, 64) et vérifie chaque forme.

À vous de jouer : des plongements plus grands

La division par √d_k suppose que les nombres des requêtes et des clés ont une variance d’environ 1. Multipliez chaque vecteur de jeton de l’exemple détaillé par 4. Avant de calculer, faites une prédiction : qu’arrive-t-il aux scores et aux poids de « damaged » (endommagé)?

Afficher la réponse

Q et K deviennent tous deux 4 fois plus grands. Chaque score devient donc 4 × 4 = 16 fois plus grand. La ligne « damaged » devient 0,354 × 16 = 5,66, 8,48 et 33,9 après la division. Softmax sature :

your_turn.py

out, weights = attention(4 * X @ W_Q, 4 * X @ W_K, 4 * X @ W_V)      # every embedding 4 times larger
print(weights.round(decimals=3))

La sortie :

tensor([[0.0000, 0.0000, 1.0000],
        [0.3330, 0.3330, 0.3330],
        [1.0000, 0.0000, 0.0000]])

« damaged » donne maintenant tout son poids à « arrival » (arrivée), et « arrival » tout son poids à « damaged ». Les poids sont presque à encodage unitaire : un poids vaut presque 1 et les autres presque 0. Les gradients qui traversent softmax sont donc proches de 0. Diviser par √d_k corrige le nombre de termes dans un produit scalaire, pas leur taille. C’est une raison pour laquelle le tout petit transformeur normalise chaque jeton (norm1) avant l’attention, comme dans Les blocs du transformeur. La ligne « on » (à) a la requête (0, 0), qui reste à 0 quand elle est multipliée par 4. Ses poids restent donc égaux.

Erreur courante

Un masque True/False qui signifie le contraire

Symptôme : Tomás remplace le masque −∞ par un masque True/False (vrai/faux), parce que PyTorch accepte les deux. Il utilise True (vrai) pour « masqué », comme le fait nn.MultiheadAttention, et le passe à F.scaled_dot_product_attention. Il n’y a aucune erreur, mais les sorties sont incorrectes :

pitfall.py

future = torch.triu(torch.ones(3, 3, dtype=torch.bool), diagonal=1)   # True above the diagonal
print(future)
wrong = F.scaled_dot_product_attention(Q, K, V, attn_mask=future)     # True means "may read" here!
right = F.scaled_dot_product_attention(Q, K, V, attn_mask=~future)
print(wrong)
print(right)

La sortie :

tensor([[False,  True,  True],
        [False, False,  True],
        [False, False, False]])
tensor([[0.0000, 0.8308],
        [0.0000, 1.0000],
        [0.0000, 0.0000]])
tensor([[2.0000, 0.0000],
        [1.0000, 0.0000],
        [1.0070, 0.2483]])

Cause : les deux fonctions interprètent un masque True/False de façons opposées. Dans nn.MultiheadAttention, True signifie « ne peut pas lire ». Dans F.scaled_dot_product_attention, True signifie « peut lire ». Le mauvais masque permet donc à chaque jeton de lire seulement le futur. La dernière ligne n’a aucune clé visible et renvoie des zéros plutôt qu’une erreur.

Correction : raisonnez avec les mathématiques : le masque s’ajoute aux scores, avec 0 pour « peut lire » et −∞ pour « ne peut pas lire ». Un masque à virgule flottante de 0 et de −∞ signifie la même chose dans les deux fonctions. Si vous utilisez un masque True/False, lisez la documentation de cette fonction. Testez ensuite le masque sur un tout petit exemple dont vous connaissez la réponse, comme le fait check.py.

Vérification des connaissances

Vérifiez votre compréhension

Pourquoi l'attention divise-t-elle les scores par la racine carrée de d_k, plutôt que par d_k ou par rien?

Afficher la réponse

Prenez une requête et une clé dont les nombres sont indépendants, avec une variance de 1. Leur produit scalaire additionne d_k produits de variance 1. Sa variance est donc d_k et sa dispersion, √d_k. Diviser par √d_k ramène la variance à 1, quelle que soit la taille des clés. Sans cette division, les scores augmentent avec d_k. La fonction softmax sature (poids maximal moyen de 0,953 à d_k = 1 024). Son gradient devient petit (0,057 au lieu de 0,240). Diviser par d_k réduirait la variance à 1 / d_k, et chaque ligne de poids deviendrait presque uniforme.

Vérifiez votre compréhension

Prédisez : une ligne de scores mis à l'échelle contient 2,0, 1,0 et 0,5. Le masque causal cache la dernière clé. Quels sont les poids? Quels seraient-ils si vous remplaciez plutôt le score caché par 0?

Afficher la réponse

Avec −∞, le troisième poids vaut exactement 0, et softmax répartit le poids entre 2,0 et 1,0. e^2 = 7,389 et e^1 = 2,718, pour un total de 10,107. Les poids sont donc 0,731, 0,269 et 0. Si le score caché vaut 0, il compte encore. e^0 = 1, le total est 11,107, et les poids sont 0,665, 0,245 et 0,090. La clé « cachée » reçoit encore 9 % du poids. Voilà pourquoi un masque ajoute −∞ avant softmax.

Vérifiez votre compréhension

Amira veut appliquer l'attention du tout petit transformeur à des fils de soutien de 8 192 jetons plutôt que 120. Comment la taille de la matrice de scores d'une tête change-t-elle? Pourquoi davantage de têtes ne réduirait-il pas le coût?

Afficher la réponse

La matrice de scores grandit selon n² : de 120 × 120 = 14 400 nombres à 8 192 × 8 192 = 67 108 864. Cela représente 268 Mo en float32 pour une tête, une séquence et une couche. La sortie de softmax double ce volume. Le temps augmente à peu près aussi vite (60 ms par tête sur le processeur partagé utilisé pour cette page). Ajouter des têtes n’ajoute pas de paramètres (4d² + 4d, quel que soit h). Mais chaque tête possède sa propre matrice n × n. La mémoire augmente donc aussi avec h, et chaque tête traite moins de nombres (d_k = d / h). Pour les longs textes, Amira a besoin d’un modèle conçu pour eux. Il utilise des noyaux qui ne stockent pas la matrice complète, des fenêtres d’attention et le cache KV pendant la génération.

Récapitulatif

• Forme matricielle : Q = X W_Q, K = X W_K, V = X W_V, (n, d) → (n, d_k); les scores S = Q Kᵀ ont la forme (n, n); A = softmax(S / √d_k + M), ligne par ligne; la sortie est A V, de forme (n, d_k).
• Mise à l’échelle : pour des nombres aléatoires de variance unitaire, Var(q · k) = d_k (valeurs mesurées : 4,03 … 1 024,99). Diviser par √d_k ramène cette variance à 1. Sans cette division, softmax sature (poids maximal de 0,335 → 0,953) et son gradient diminue vers 0.
• Masques : une matrice de 0 et de −∞ ajoutée avant softmax. Un score de −∞ donne un poids d’exactement 0, et les poids visibles totalisent encore 1. Un score de 0 compterait encore.
• Attention multitête : h têtes de d_k = d / h, réunies et combinées par W_O; 4d² + 4d paramètres, quel que soit h (16 640 à d = 64, 1 050 624 à d = 512). Les têtes apprennent des relations différentes, comme le montrent les quatre têtes de T-65052.
• Coût : un temps en O(n² · d_k) et une mémoire en n² par tête, donc O(n² · d) pour l’ensemble des h têtes : 268 Mo et 60 ms pour une tête à 8 192 jetons sur un processeur partagé. Les longs contextes exigent d’autres techniques.
• Vérification : le code lisible donne les mêmes résultats que PyTorch à 2,4 × 10⁻⁷ près en float32, avec et sans masques.

Ensuite, le transformeur original de 2017, reconstruit à partir de ces éléments et suivi étape par étape.

Leçon suivante : Le transformeur original, étape par étape