Développement de LLM, partie 8 — La self-attention entraînable

• Dans un Transformer de type GPT, décodeur uniquement, la self-attention entraînable calcule à quels tokens des entrées précédentes chaque token doit prêter attention afin de produire un vecteur de contexte
• Le cœur du mécanisme est la scaled dot product attention, qui utilise trois matrices apprenables Wq, Wk, Wv pour projeter les embeddings d’entrée dans les espaces query, key et value
• La matrice d’entrée X est transformée en Q=XWq, K=XWk, V=XWv, puis Ω=QKᵀ est divisé par √c avant d’appliquer une softmax par ligne pour obtenir les poids d’attention A
• Les vecteurs de contexte sont produits en une seule multiplication de matrices, C=AV, et l’ensemble du calcul peut être appliqué à tous les tokens avec 5 multiplications de matrices et une transposition
• Cette étape dépasse l’exemple jouet où l’on faisait directement le produit scalaire entre embeddings d’entrée, pour aller vers une attention entraînable implémentable avec PyTorch nn.Module et nn.Linear

---

La place de la self-attention dans le pipeline de traitement d’un LLM

• Un LLM basé sur un Transformer de type GPT, décodeur uniquement est conçu pour prédire le token suivant à partir des tokens vus jusque-là
• Le pipeline consiste à découper une chaîne en tokens, transformer chaque token en embedding de token, puis ajouter un embedding de position représentant l’information de position afin de former les embeddings d’entrée
• La self-attention génère, pour chaque embedding d’entrée, une liste de scores d’attention indiquant dans quelle mesure il faut prêter attention aux autres tokens
  • Dans la phrase d’exemple "the fat cat sat on the mat", lorsque l’on regarde "cat", "fat" peut être important
  • En revanche, lorsque l’on regarde "mat", l’importance de "fat" peut être relativement plus faible
• Après passage par softmax, les scores d’attention deviennent des poids d’attention dont la somme vaut 1, puis ces poids servent à calculer une somme pondérée des embeddings d’entrée pour produire un vecteur de contexte
• Le vecteur de contexte est traité comme un vecteur représentant le sens de chaque token dans le contexte de l’ensemble de l’entrée

L’objectif de la self-attention entraînable

• Jusqu’à l’étape précédente, on utilisait une self-attention jouet qui calculait directement le produit scalaire entre embeddings d’entrée
• L’objectif de cette étape est de construire un mécanisme d’attention entraînable capable de produire des scores d’attention à partir des vecteurs d’entrée
• La section 3.4 de Build a Large Language Model (from Scratch) de Sebastian Raschka l’implémente sous la forme de la scaled dot product attention
• L’accent est mis non pas sur la raison pour laquelle cette structure est efficace, mais sur les calculs qu’elle effectue

Les matrices Query, Key, Value et la projection dans des espaces distincts

• On note n la longueur de la séquence d’entrée, d la dimension des embeddings d’entrée et c la dimension des vecteurs de contexte
• La séquence d’embeddings d’entrée est représentée par x1, x2, x3, ... xn, et chaque embedding d’entrée est un vecteur de dimension d
• On définit trois matrices de poids apprenables
  • query weights matrix : Wq
  • key weights matrix : Wk
  • value weights matrix : Wv
• Chaque matrice est de taille d×c et projette un vecteur d’entrée de dimension d dans un espace de dimension c
• Le calcul qui envoie le vecteur d’entrée xm dans l’espace query est qm=xmWq
• Les espaces key et value projettent eux aussi les embeddings d’entrée dans des espaces c-dimensionnels distincts de la même manière

Voir les matrices comme des projections

• Les matrices peuvent être utilisées pour des transformations géométriques, par exemple faire tourner des points
• Une matrice carrée effectue une transformation dans le même espace dimensionnel, tandis qu’une matrice non carrée peut envoyer un vecteur vers un espace d’une autre dimension
• Par exemple, une matrice 3×2 peut transformer des vecteurs 3D en vecteurs 2D
• En graphisme 3D, la matrice de frustum qui transforme des points 3D en points 2D à l’écran est aussi un exemple de ce type de projection
• La self-attention projette les embeddings d’entrée dans trois espaces de projection différents — query, key et value — puis effectue les calculs à partir de ces vecteurs projetés
• Comme ces matrices de projection sont apprises pendant l’entraînement, cela introduit une indirection absente de l’attention par simple produit scalaire

Calcul des scores d’attention

• Si l’on considère une entrée donnée xm, le score d’attention vis-à-vis d’une autre entrée xp est défini comme le produit scalaire entre la projection query et la projection key
• Les formules sont les suivantes
  • qm=xmWq
  • kp=xpWk
  • ωm,p=qm·kp
• On pourrait répéter ce calcul pour toutes les entrées avec des boucles, mais une multiplication de matrices permet de tout calculer d’un seul coup
• Si l’on note X la matrice contenant l’ensemble des embeddings d’entrée, alors X est de taille n×d
• La matrice key est calculée d’un seul coup avec K=XWk
  • Le résultat K est de taille n×c
  • Chaque ligne est le vecteur obtenu en projetant l’embedding d’entrée correspondant dans l’espace key
• La matrice query se calcule de la même manière avec Q=XWq
• Les produits scalaires entre toutes les queries et toutes les keys s’obtiennent avec QKᵀ
  • Q est de taille n×c
  • Kᵀ est de taille c×n
  • Le résultat Ω est de taille n×n
• Ωm,p est le score d’attention indiquant à quel point il faut prêter attention à xp pour produire le vecteur de contexte de xm

Mise à l’échelle et normalisation softmax

• Comme dans l’exemple précédent, les scores d’attention sont convertis en poids dont la somme vaut 1 via une softmax
• La softmax accentue les grandes valeurs et réduit les petites, tout en ajustant l’ensemble pour que leur somme soit égale à 1
• Dans un LLM réel, d et c peuvent se compter en milliers, si bien qu’une softmax appliquée seule peut produire de petits gradients
• Dans ce cas, la softmax peut se comporter « comme une step function »
  • On peut l’interpréter comme une situation où la plus grande valeur domine et où toutes les autres deviennent très petites
• Pour atténuer cela, on divise les scores d’attention par la racine carrée de la dimension de l’espace de projection c avant d’appliquer la softmax
• En notation matricielle, cela donne
  • A=softmax(Ω/√c, axis=1)
• axis=1 est une notation de style PyTorch qui signifie que la softmax est appliquée ligne par ligne
• Le résultat A est la matrice des scores d’attention normalisés, c’est-à-dire la matrice des poids d’attention

Génération des vecteurs de contexte

• La projection dans l’espace value se calcule avec V=XWv
• A est une matrice de poids d’attention de taille n×n
  • Am,p est le poids d’attention à appliquer à l’entrée p lors de la construction du vecteur de contexte de xm
• V est de taille n×c, et chaque ligne est le vecteur obtenu en projetant un embedding d’entrée dans l’espace value
• La matrice des vecteurs de contexte se calcule avec C=AV
  • Le résultat C est de taille n×c
  • La m-ième ligne de C est le vecteur de contexte associé à l’entrée xm
• Ce calcul effectue, en une seule multiplication de matrices, l’opération qui consiste à multiplier les vecteurs value par les poids d’attention puis à les additionner pour chaque token

Résumé du calcul global

• La matrice d’entrée X contient les embeddings d’entrée de la séquence de tokens et est de taille n×d
• Trois matrices apprenables projettent l’entrée dans les espaces query, key et value
  • Q=XWq
  • K=XWk
  • V=XWv
• Les scores d’attention sont calculés à partir du produit scalaire entre query et key
  • Ω=QKᵀ
• Après mise à l’échelle des scores, une softmax par ligne produit les poids d’attention
  • A=softmax(Ω/√c, axis=1)
• La projection value est multipliée par les poids d’attention pour produire les vecteurs de contexte
  • C=AV
• L’ensemble du mécanisme de self-attention peut produire les vecteurs de contexte de tous les tokens d’entrée avec 5 multiplications de matrices et une transposition

Implémentation PyTorch et prochaine étape

• La section 3.4 du livre implémente les calculs ci-dessus en code PyTorch et crée une sous-classe simple de nn.Module effectuant les mêmes opérations matricielles
• La première version utilise de simples objets nn.Parameter pour les trois matrices de poids
• La deuxième version utilise nn.Linear pour un entraînement plus efficace
• Deux sujets seront ensuite abordés
  • causal self-attention : une méthode où, lorsqu’on traite un token donné, on ne prête pas attention aux tokens suivants
  • multi-head attention : annoncée comme moins compliquée qu’on pourrait le croire au premier abord
• Le traitement par lots reste une question distincte
  • Même avec une seule séquence d’entrée, on utilise déjà une matrice de scores d’attention
  • Pour traiter plusieurs séquences d’entrée en parallèle, des tenseurs d’ordre supérieur aux matrices peuvent être nécessaires
• L’article suivant est Writing an LLM from scratch, part 9 -- causal attention