Comment un bloc Transformer met-il à jour la représentation d’un token ?

English · 中文 · Français

First Principles of LLMs & RL · Article 07 · Article précédent

Note d’écriture : les idées et le matériau initial sont les miens. ChatGPT a aidé à la rédaction, à la révision, à la traduction, aux vérifications numériques et à la création des figures et animations originales.

À la fin de l’article précédent, nous avions deux vecteurs pour la position 2. Le premier était la représentation dont elle disposait en entrant dans la couche ; le second contenait ce qu’elle venait de lire aux positions visibles :

Il reste à écrire la ligne de code qui transmet un vecteur à la couche suivante. Lequel lui donner ?

On pourrait transmettre le second : les dimensions concordent, le programme s’exécuterait. Mais ce choix revient à remplacer la représentation de la position par le contenu qu’elle vient de lire. On pourrait aussi utiliser ce contenu pour mettre à jour la représentation existante. La différence entre ces deux lignes de code fournit un point de départ pour construire un bloc Transformer.

Si une position peut se lire elle-même, pourquoi rajouter son entrée ?

Reprenons le calcul du second vecteur. La position 2 pouvait lire la position 1 et se lire elle-même, mais ne pouvait pas accéder à la position future 3. Dans notre exemple, elle attribuait la moitié du poids à chacune des deux positions visibles :

Cette somme contient donc déjà une part d’information issue de la position 2. Encore faut-il regarder par quel chemin elle est passée. L’entrée a d’abord subi la projection Value. La valeur obtenue a ensuite été pondérée, puis mélangée à la contribution de la position 1. Dans cet exemple, la Value de la position 2 coïncide avec son entrée ; une autre projection pourrait les distinguer. Même si l’on conserve la projection, la contribution de la position 2 diminue dès que la tête accorde davantage de poids à la position 1.

D’où l’intérêt d’un autre chemin : une position devrait pouvoir consacrer son attention à ce qui se trouve ailleurs sans devoir en réserver une partie à la transmission de sa propre entrée.

Donnons à l’entrée un chemin direct jusqu’à l’addition. Prenons d’abord une seule tête dont la sortie a la même dimension que l’entrée, une projection de sortie égale à l’identité, et aucune normalisation. Notons u₂ la représentation obtenue après la mise à jour :

Ce chemin qui contourne la branche de lecture est la connexion résiduelle. Pourquoi « résiduelle » ? Pour passer de l’entrée à une sortie donnée, il suffit que la branche fournisse la différence entre les deux. La première coordonnée valait zéro et augmente d’une unité. La seconde valait déjà un ; la branche ne doit y ajouter qu’un demi. Avec d’autres paramètres, la correction peut aussi être négative.¹

Un test permet de saisir la différence sans attribuer de sens particulier aux coordonnées : annulons la sortie de la branche. Si le résultat de la lecture remplace l’entrée, la couche suivante reçoit zéro. Si on l’ajoute à l’entrée, elle retrouve celle-ci intacte. À l’inverse, une branche qui renverrait l’opposé de l’entrée pourrait l’annuler entièrement. Ce que la connexion résiduelle préserve, c’est un chemin direct pour l’entrée. La somme finale ne permet pas nécessairement d’en retrouver séparément tous les termes.

Les deux chemins partent de la même entrée et de la même sortie de branche. Que reste-t-il à la couche suivante lorsque cette dernière est nulle ?

Qu’impose une seule distribution de poids ?

Revenons au facteur un demi devant la première Value. Il s’applique à toutes les coordonnées du vecteur. Supposons que l’on veuille prendre davantage de la première coordonnée à la position 1 et davantage de la seconde coordonnée à la position 2. Une tête peut-elle régler ces deux contributions indépendamment ?

Notons a le poids de la position 1. Celui de la position 2 vaut alors 1 − a. Si les deux Values restent fixes, toute lecture s’écrit :

Augmentons légèrement a : la première coordonnée augmente, la seconde diminue. Si l’on représente les deux Values par des points dans le plan, la sortie ne peut que se déplacer sur le segment qui les relie. Pour obtenir 3/2 sur la première coordonnée, il faut fixer a à 3/4. Pour obtenir aussi 3/4 sur la seconde, il faudrait fixer le même a à 1/4. Choisir le poids avec une précision accrue ne résoudrait rien. Il nous manque un deuxième réglage indépendant.

Passons à deux têtes, qui renvoient chacune un scalaire. Pour isoler ce qui change, choisissons pour A une projection Value qui extrait la première coordonnée de l’exemple précédent, et pour B une projection qui extrait la seconde. Chaque tête calcule ses propres scores. À la position 2, donnons-leur les poids suivants, choisis à la main :

Les deux nombres voulus sont désormais disponibles en même temps. A et B n’ont pas à s’accorder : chacune a ses scores et son softmax. Nous avons changé la règle de lecture ; les poids du tableau n’ont pas été mesurés sur le modèle de l’article 06. Extraire une seule coordonnée est également un choix pédagogique : une projection Value apprise peut combiner plusieurs coordonnées. Pour nos deux Values fixes, les têtes supplémentaires permettent de prélever des caractéristiques différentes dans des proportions différentes selon la position. En changeant aussi la projection Value, une tête unique pourrait fort bien produire ces deux mêmes nombres sur cette entrée particulière. La comparaison porte sur le mode de lecture.

Une fois les lectures terminées, il reste à décider quoi écrire. La concaténation place simplement les sorties côte à côte. Gardons la convention de l’article 06 : le vecteur d’une position s’écrit en colonne ; sa transposée forme une ligne dans la matrice de la séquence. Notons s₂ le vecteur concaténé et choisissons une petite matrice de projection :

La première colonne écrit la différence entre les sorties des deux têtes ; la seconde écrit la sortie de B. La mise à jour que fournit la branche d’attention, puis la représentation obtenue après l’addition, valent donc :

Les têtes déterminent quelle part prendre à chaque position ; la projection de sortie combine ce qu’elles rapportent ; l’addition résiduelle inscrit cette combinaison dans la représentation courante. Avec quatre têtes produisant deux nombres chacune, par exemple, la concaténation donne huit nombres. La projection de sortie les ramène à la dimension de l’entrée, ce qui rend possible l’addition coordonnée par coordonnée. C’est la construction par concaténation puis projection du Transformer d’origine.²

Les poids sont choisis pour cet exemple. Suivre la lecture de chaque tête, la concaténation et la projection, puis repérer l’addition résiduelle distincte.

Que peut faire une position de l’information qu’elle vient de recevoir ?

Après cette mise à jour, la position 2 contient deux nombres : 3/4 et 7/4. Donnons-lui une petite tâche : si la seconde coordonnée dépasse la première, ajouter l’excédent à la première ; sinon, laisser la représentation telle quelle.

Faut-il consulter à nouveau le contexte ? Les deux nombres nécessaires sont déjà présents. Notons p et q les coordonnées de cette position. Une seule unité cachée suffit pour calculer la différence, ramener les valeurs négatives à zéro et envoyer le résultat vers la première coordonnée de sortie :

Dans notre exemple, la différence vaut exactement un. Après addition de la mise à jour, on obtient :

Voilà une opération que peut effectuer un réseau feed-forward appliqué à chaque position (FFN). La première transformation linéaire extrait la seconde coordonnée moins la première ; ReLU laisse cette caractéristique agir seulement si la différence est positive ; la seconde transformation linéaire choisit où écrire le résultat. Sans ReLU, les deux transformations linéaires pourraient se ramener à une seule et l’on perdrait cette réponse limitée à un côté du seuil. Rien ne dit que deux coordonnées d’un modèle entraîné se prêtent à cette interprétation. Nous avons choisi les paramètres pour pouvoir suivre le calcul à la main.

Comparons deux positions dans la séquence. L’attention peut faire dépendre le résultat de la position 2 du contenu de la position 1. Le FFN reçoit le vecteur mis à jour à la position 2 et transforme ce seul vecteur. En l’appliquant à la position 1, on utilise les mêmes paramètres ; le résultat peut différer parce que l’entrée diffère. Travailler position par position ne signifie donc pas traiter un mot isolé : l’attention a déjà acheminé du contexte jusqu’ici.²

Le Transformer d’origine utilise un FFN avec ReLU. Dans mon implémentation, j’utilise SwiGLU : à chaque position, une première projection fournit des caractéristiques candidates ; une deuxième, suivie d’une activation, les module avant la projection de retour à la dimension du modèle. Le calcul interne à la position change. Il n’y a toujours pas de lecture à la position suivante.³

Quel chemin passe par la normalisation ?

Pour suivre les additions, nous avons jusqu’ici laissé la normalisation de côté. Or, si l’on empile ces couches, chacune ajoute sa mise à jour. Les vecteurs que recevront les branches suivantes n’auront pas forcément la même amplitude qu’à l’entrée.

Cela affecte déjà la lecture. Prenons une attention par produit scalaire avec des projections sans biais. Si l’on double tous les vecteurs d’entrée sans rien changer d’autre, les Queries et les Keys doublent également. Leurs produits scalaires sont donc multipliés par quatre. Un léger avantage pour une position peut alors se traduire, après softmax, par une lecture beaucoup plus concentrée. Dans l’article 06, la division par √dₖ répondait à un problème d’échelle lié à la dimension du produit scalaire. Ici la dimension n’a pas changé ; c’est l’amplitude des entrées qui a varié.

Une solution consiste à laisser circuler la représentation dans le chemin résiduel tout en présentant une version normalisée à chaque branche. Prenons RMSNorm. Pour le vecteur z d’une position, de dimension d, notons γ les coefficients ajustables propres à chaque coordonnée et ε un petit nombre positif qui évite un dénominateur nul :

Le dénominateur est calculé sur les coordonnées de cette position, sans moyenne entre tokens. Lorsque ε est négligeable devant la moyenne des carrés, une multiplication de tout le vecteur par un facteur positif multiplie approximativement le numérateur et le dénominateur par le même facteur. L’augmentation de la norme du vecteur ne suffit donc plus à rendre les scores plus extrêmes. Contrairement à LayerNorm, RMSNorm ne soustrait pas la moyenne.⁴

Écrivons maintenant le bloc entier. X contient la séquence d’entrée : une ligne pour chacune des n positions, et d_model coordonnées par ligne. U désigne la séquence après l’écriture de l’attention, et Y la sortie du bloc. N₁ et N₂ sont deux RMSNorm aux paramètres distincts, appliquées séparément à chaque ligne ; MHA comprend l’attention causale multi-têtes et sa projection de sortie. Omettons la dimension de batch et le dropout :

Regardons les deux additions. La première ajoute la mise à jour à X lui-même. La seconde l’ajoute à U, qui contient déjà la mise à jour de l’attention. La normalisation se trouve à l’entrée de chaque branche ; le chemin direct la contourne. C’est l’agencement Pre-Norm utilisé ici. Les calculs à la main isolaient certaines opérations ; dans le bloc complet, les branches reçoivent des entrées normalisées, ce qui peut changer les résultats numériques. En laissant de côté la gestion du cache, les deux écritures dans l’implémentation se lisent ainsi :

u = x + attention(attention_norm(x))
y = u + feed_forward(feed_forward_norm(u))

Dans le bloc complet, les entrées des branches sont normalisées, contrairement aux calculs locaux précédents. Le FFN agit séparément à chaque position ; la représentation mise à jour alimente les projections de la couche suivante.⁵

Reprenons le test du début. Si les deux branches renvoient une mise à jour nulle pour leur entrée respective, la première ligne restitue l’entrée et la seconde aussi. Oublions le premier x + : le résultat tombe à zéro, bien que tous les tenseurs aient toujours les bonnes dimensions. Nous testons ce bloc précis, avant toute normalisation finale distincte du modèle. Dans l’architecture Post-LN du Transformer d’origine, la normalisation suit l’addition résiduelle. Même si la sous-couche renvoie zéro, l’entrée passe encore par LayerNorm : la conclusion sur l’identité ne se transpose pas telle quelle.²⁶

Retrouvons la rétropropagation de l’article 03. Aplatissons toute l’entrée d’une sous-couche résiduelle en un vecteur colonne x. Appelons F sa branche et y sa sortie. Si g est le gradient colonne de la perte par rapport à la sortie, la règle de dérivation en chaîne donne, là où F est différentiable :

J_F désigne la matrice jacobienne de la branche. Le premier terme g à droite n’a pas besoin de traverser la branche ; il provient de l’addition directe du calcul avant. L’autre terme peut le renforcer ou l’annuler. Cette écriture nous dit par quel chemin le gradient peut passer, sans garantir qu’il reste suffisamment grand.¹

D’où vient la Query de la couche suivante ?

Reprenons I like et They like. Dans l’exemple précédent, like a le même identifiant de token et occupe la même position. À la première couche, il peut donc partir de la même représentation et émettre la même requête. Le mot qui le précède étant différent, cette requête peut néanmoins rapporter des contenus différents.

Après un bloc, l’identifiant de like n’a pas changé. En revanche, la ligne transmise à la couche suivante a reçu une mise à jour de l’attention, puis une autre du FFN à partir des caractéristiques désormais disponibles à cette position. La couche suivante calcule sa Query à partir de cette ligne mise à jour, avec ses propres paramètres de normalisation et de projection. Une différence dans le préfixe peut désormais influer sur le résultat de la première lecture, mais aussi sur les indices qui guideront la seconde. La projection peut rester insensible à cette différence. L’architecture rend la dépendance possible sans imposer d’interprétation grammaticale.

Voilà ce que l’empilement de deux couches apporte : la seconde lecture peut prendre pour point de départ une représentation qui contient déjà de l’information rapportée par la première.

À la toute dernière couche, il reste pourtant un pas à franchir. La dernière position livre encore un vecteur de représentation, alors que nous attendons un mot. Comment passer de ce vecteur à une distribution sur le vocabulaire ? Une fois le token choisi et ajouté à la séquence, quels calculs faudra-t-il reprendre, et lesquels pourra-t-on conserver ? C’est là que commence la génération.