L’attention d’un Transformer, calculée à la main

Le mot « attention » invite à imaginer un lecteur qui choisit les mots importants. L'image est commode, mais elle cache le calcul. Dans un Transformer, une couche produit des combinaisons pondérées de vecteurs. Les coefficients de ces mélanges dépendent eux-mêmes des vecteurs présents dans la séquence. Il n'est pas nécessaire de prêter une intention au mécanisme pour le comprendre.

L'article Attention Is All You Need, publié par Vaswani et ses coauteurs en 2017, introduit une architecture fondée sur l'attention, sans récurrence ni convolution pour son modèle principal. Concentrons-nous sur un seul de ses blocs. Les nombres qui suivent sont un exemple pédagogique construit pour pouvoir refaire les opérations.

Une recherche, des étiquettes et du contenu

Chaque token est représenté par un vecteur. Trois transformations linéaires apprises construisent des requêtes $Q$, des clés $K$ et des valeurs $V$. Une requête est comparée aux clés ; les valeurs sont le contenu que l'on mélange ensuite. Dans une même séquence, ces trois objets viennent des mêmes représentations initiales, mais leurs matrices de projection diffèrent.

Pour une requête à deux dimensions, choisissons :

Les produits scalaires valent $1,0,1$. Le premier et le troisième token sont donc à égalité selon cette requête, même si leurs clés sont différentes. Le produit scalaire n'est pas une mesure magique de sens : il ne voit que les directions que l'apprentissage a placées dans cet espace.

Du score à un mélange

Avec une dimension de clé $d_k=2$, divisons les scores par $\sqrt{2}$. On obtient environ $0{,}707,0,0{,}707$. Le softmax transforme ces trois nombres en coefficients positifs dont la somme vaut un :

Pourquoi cette division ? Si les composantes des requêtes et des clés sont indépendantes, centrées et de variance un, la variance de leur produit scalaire croît comme $d_k$. Diviser par $\sqrt{d_k}$ stabilise cette échelle sous ces hypothèses. Sans cela, de grands scores peuvent rendre le softmax très concentré et réduire ses gradients utiles.

Prenons maintenant $v_1=(2,0)$, $v_2=(0,2)$ et $v_3=(2,2)$. La sortie vaut :

La sortie n'est donc pas le troisième token, ni son nom, ni une explication. C'est un nouveau vecteur, situé ici dans l'enveloppe convexe des trois valeurs.

Le même calcul pour toute la phrase

En empilant les requêtes dans une matrice, tous les produits scalaires deviennent une multiplication matricielle :

Le softmax s'applique ligne par ligne. Pour une séquence de $n$ tokens, $QK^\top$ comporte $n^2$ scores. Si le modèle doit prédire la suite sans regarder le futur, le masque $M$ interdit les positions futures en leur attribuant un score négatif infini avant le softmax. Un score interdit reçoit ainsi un poids nul.

Ce masque ne fournit pas l'ordre des mots. Sans information de position, permuter les lignes permute les résultats de façon correspondante. Les représentations positionnelles apportent une information que le mélange, à lui seul, ne possède pas.

Une implémentation assez petite pour être inspectée

Le code suivant calcule notre ligne sans bibliothèque. Soustraire le maximum avant l'exponentielle préserve le résultat tout en limitant le risque de débordement numérique.

from math import exp, sqrt
 
q = [1.0, 0.0]
keys = [[1, 0], [0, 1], [1, 1]]
values = [[2, 0], [0, 2], [2, 2]]
scores = [sum(a*b for a, b in zip(q, k)) / sqrt(2) for k in keys]
shift = max(scores)
weights = [exp(s - shift) for s in scores]
weights = [w / sum(weights) for w in weights]
output = [sum(w*v[d] for w, v in zip(weights, values)) for d in range(2)]
print(weights, output)

Deux expériences éclairent le mécanisme. Multiplier les scores par dix concentre les poids sur les deux maxima, mais ne départage pas leur égalité. Changer seulement les valeurs modifie la sortie sans modifier aucun coefficient d'attention. On voit ainsi que « où regarder » et « quoi transporter » sont deux calculs distincts.

Les multiples têtes répètent ce principe dans plusieurs sous-espaces appris, puis combinent leurs sorties. Elles ne correspondent pas nécessairement à des fonctions linguistiques simples. Enfin, visualiser les coefficients peut être instructif, mais un coefficient élevé n'est pas une preuve causale qu'un mot explique la décision finale : d'autres couches, connexions résiduelles et transformations interviennent. Le calcul local est précis ; l'interprétation globale demande une expérience supplémentaire.