4. Mécanismes d'Attention

Reading time: 13 minutes

Mécanismes d'Attention et Auto-Attention dans les Réseaux de Neurones

Les mécanismes d'attention permettent aux réseaux de neurones de se concentrer sur des parties spécifiques de l'entrée lors de la génération de chaque partie de la sortie. Ils attribuent des poids différents à différentes entrées, aidant le modèle à décider quelles entrées sont les plus pertinentes pour la tâche à accomplir. Cela est crucial dans des tâches comme la traduction automatique, où comprendre le contexte de l'ensemble de la phrase est nécessaire pour une traduction précise.

Comprendre les Mécanismes d'Attention

Dans les modèles traditionnels de séquence à séquence utilisés pour la traduction linguistique, le modèle encode une séquence d'entrée en un vecteur de contexte de taille fixe. Cependant, cette approche a des difficultés avec les longues phrases car le vecteur de contexte de taille fixe peut ne pas capturer toutes les informations nécessaires. Les mécanismes d'attention répondent à cette limitation en permettant au modèle de considérer tous les tokens d'entrée lors de la génération de chaque token de sortie.

Exemple : Traduction Automatique

Considérons la traduction de la phrase allemande "Kannst du mir helfen diesen Satz zu übersetzen" en anglais. Une traduction mot à mot ne produirait pas une phrase anglaise grammaticalement correcte en raison des différences dans les structures grammaticales entre les langues. Un mécanisme d'attention permet au modèle de se concentrer sur les parties pertinentes de la phrase d'entrée lors de la génération de chaque mot de la phrase de sortie, conduisant à une traduction plus précise et cohérente.

Introduction à l'Auto-Attention

L'auto-attention, ou intra-attention, est un mécanisme où l'attention est appliquée au sein d'une seule séquence pour calculer une représentation de cette séquence. Elle permet à chaque token de la séquence de prêter attention à tous les autres tokens, aidant le modèle à capturer les dépendances entre les tokens, quelle que soit leur distance dans la séquence.

Concepts Clés

• Tokens : Éléments individuels de la séquence d'entrée (par exemple, mots dans une phrase).
• Embeddings : Représentations vectorielles des tokens, capturant des informations sémantiques.
• Poids d'Attention : Valeurs qui déterminent l'importance de chaque token par rapport aux autres.

Calcul des Poids d'Attention : Un Exemple Étape par Étape

Considérons la phrase "Hello shiny sun!" et représentons chaque mot avec un embedding en 3 dimensions :

• Hello : [0.34, 0.22, 0.54]
• shiny : [0.53, 0.34, 0.98]
• sun : [0.29, 0.54, 0.93]

Notre objectif est de calculer le vecteur de contexte pour le mot "shiny" en utilisant l'auto-attention.

Étape 1 : Calculer les Scores d'Attention

Pour chaque mot de la phrase, calculez le score d'attention par rapport à "shiny" en calculant le produit scalaire de leurs embeddings.

Score d'Attention entre "Hello" et "shiny"

Score d'Attention entre "shiny" et "shiny"

Score d'Attention entre "sun" et "shiny"

Étape 2 : Normaliser les Scores d'Attention pour Obtenir les Poids d'Attention

Appliquez la fonction softmax aux scores d'attention pour les convertir en poids d'attention qui s'additionnent à 1.

Calcul des exponentielles :

Calcul de la somme :

Calcul des poids d'attention :

Étape 3 : Calculer le Vecteur de Contexte

Le vecteur de contexte est calculé comme la somme pondérée des embeddings de tous les mots, en utilisant les poids d'attention.

Calcul de chaque composant :

• Embedding Pondéré de "Hello" :

• Embedding Pondéré de "shiny" :

• Embedding Pondéré de "sun" :

Somme des embeddings pondérés :

vecteur de contexte=[0.0779+0.2156+0.1057, 0.0504+0.1382+0.1972, 0.1237+0.3983+0.3390]=[0.3992,0.3858,0.8610]

Ce vecteur de contexte représente l'embedding enrichi pour le mot "shiny", incorporant des informations de tous les mots de la phrase.

Résumé du Processus

1. Calculer les Scores d'Attention : Utilisez le produit scalaire entre l'embedding du mot cible et les embeddings de tous les mots de la séquence.
2. Normaliser les Scores pour Obtenir les Poids d'Attention : Appliquez la fonction softmax aux scores d'attention pour obtenir des poids qui s'additionnent à 1.
3. Calculer le Vecteur de Contexte : Multipliez l'embedding de chaque mot par son poids d'attention et additionnez les résultats.

Auto-Attention avec Poids Entraînables

En pratique, les mécanismes d'auto-attention utilisent des poids entraînables pour apprendre les meilleures représentations pour les requêtes, les clés et les valeurs. Cela implique l'introduction de trois matrices de poids :

La requête est les données à utiliser comme auparavant, tandis que les matrices de clés et de valeurs sont simplement des matrices aléatoires entraînables.

Étape 1 : Calculer les Requêtes, Clés et Valeurs

Chaque token aura sa propre matrice de requête, de clé et de valeur en multipliant ses valeurs de dimension par les matrices définies :

Ces matrices transforment les embeddings originaux en un nouvel espace adapté au calcul de l'attention.

Exemple

En supposant :

• Dimension d'entrée din=3 (taille de l'embedding)
• Dimension de sortie dout=2 (dimension souhaitée pour les requêtes, clés et valeurs)

Initialisez les matrices de poids :

import torch.nn as nn

d_in = 3
d_out = 2

W_query = nn.Parameter(torch.rand(d_in, d_out))
W_key = nn.Parameter(torch.rand(d_in, d_out))
W_value = nn.Parameter(torch.rand(d_in, d_out))

Calculer les requêtes, les clés et les valeurs :

queries = torch.matmul(inputs, W_query)
keys = torch.matmul(inputs, W_key)
values = torch.matmul(inputs, W_value)

Étape 2 : Calculer l'attention par produit scalaire mis à l'échelle

Calculer les scores d'attention

Semblable à l'exemple précédent, mais cette fois, au lieu d'utiliser les valeurs des dimensions des tokens, nous utilisons la matrice clé du token (déjà calculée en utilisant les dimensions) : . Donc, pour chaque requête qi​ et clé kj​ :

Mettre à l'échelle les scores

Pour éviter que les produits scalaires ne deviennent trop grands, mettez-les à l'échelle par la racine carrée de la dimension clé dk​ :

Appliquer Softmax pour obtenir les poids d'attention : Comme dans l'exemple initial, normalisez toutes les valeurs pour qu'elles s'additionnent à 1.

Étape 3 : Calculer les vecteurs de contexte

Comme dans l'exemple initial, il suffit de sommer toutes les matrices de valeurs en multipliant chacune par son poids d'attention :

Exemple de code

En prenant un exemple de https://github.com/rasbt/LLMs-from-scratch/blob/main/ch03/01_main-chapter-code/ch03.ipynb, vous pouvez consulter cette classe qui implémente la fonctionnalité d'auto-attention dont nous avons parlé :

import torch

inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your     (x^1)
[0.55, 0.87, 0.66], # journey  (x^2)
[0.57, 0.85, 0.64], # starts   (x^3)
[0.22, 0.58, 0.33], # with     (x^4)
[0.77, 0.25, 0.10], # one      (x^5)
[0.05, 0.80, 0.55]] # step     (x^6)
)

import torch.nn as nn
class SelfAttention_v2(nn.Module):

def __init__(self, d_in, d_out, qkv_bias=False):
super().__init__()
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

def forward(self, x):
keys = self.W_key(x)
queries = self.W_query(x)
values = self.W_value(x)

attn_scores = queries @ keys.T
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)

context_vec = attn_weights @ values
return context_vec

d_in=3
d_out=2
torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

Attention Causale : Masquer les Mots Futurs

Pour les LLMs, nous voulons que le modèle ne considère que les tokens qui apparaissent avant la position actuelle afin de prédire le prochain token. L'attention causale, également connue sous le nom de masquage d'attention, y parvient en modifiant le mécanisme d'attention pour empêcher l'accès aux tokens futurs.

Application d'un Masque d'Attention Causale

Pour mettre en œuvre l'attention causale, nous appliquons un masque aux scores d'attention avant l'opération softmax afin que les scores restants s'additionnent toujours à 1. Ce masque fixe les scores d'attention des tokens futurs à moins l'infini, garantissant qu'après le softmax, leurs poids d'attention sont nuls.

Étapes

1. Calculer les Scores d'Attention : Comme auparavant.
2. Appliquer le Masque : Utiliser une matrice triangulaire supérieure remplie de moins l'infini au-dessus de la diagonale.

mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1) * float('-inf')
masked_scores = attention_scores + mask

3. Appliquer Softmax : Calculer les poids d'attention en utilisant les scores masqués.

attention_weights = torch.softmax(masked_scores, dim=-1)

Masquage des Poids d'Attention Supplémentaires avec Dropout

Pour prévenir le surapprentissage, nous pouvons appliquer dropout aux poids d'attention après l'opération softmax. Le dropout met aléatoirement à zéro certains des poids d'attention pendant l'entraînement.

dropout = nn.Dropout(p=0.5)
attention_weights = dropout(attention_weights)

Un abandon régulier est d'environ 10-20%.

Code Exemple

Exemple de code provenant de https://github.com/rasbt/LLMs-from-scratch/blob/main/ch03/01_main-chapter-code/ch03.ipynb:

import torch
import torch.nn as nn

inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your     (x^1)
[0.55, 0.87, 0.66], # journey  (x^2)
[0.57, 0.85, 0.64], # starts   (x^3)
[0.22, 0.58, 0.33], # with     (x^4)
[0.77, 0.25, 0.10], # one      (x^5)
[0.05, 0.80, 0.55]] # step     (x^6)
)

batch = torch.stack((inputs, inputs), dim=0)
print(batch.shape)

class CausalAttention(nn.Module):

def __init__(self, d_in, d_out, context_length,
dropout, qkv_bias=False):
super().__init__()
self.d_out = d_out
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.dropout = nn.Dropout(dropout)
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # New

def forward(self, x):
b, num_tokens, d_in = x.shape
# b is the num of batches
# num_tokens is the number of tokens per batch
# d_in is the dimensions er token

keys = self.W_key(x) # This generates the keys of the tokens
queries = self.W_query(x)
values = self.W_value(x)

attn_scores = queries @ keys.transpose(1, 2) # Moves the third dimension to the second one and the second one to the third one to be able to multiply
attn_scores.masked_fill_(  # New, _ ops are in-place
self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)  # `:num_tokens` to account for cases where the number of tokens in the batch is smaller than the supported context_size
attn_weights = torch.softmax(
attn_scores / keys.shape[-1]**0.5, dim=-1
)
attn_weights = self.dropout(attn_weights)

context_vec = attn_weights @ values
return context_vec

torch.manual_seed(123)

context_length = batch.shape[1]
d_in = 3
d_out = 2
ca = CausalAttention(d_in, d_out, context_length, 0.0)

context_vecs = ca(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

Étendre l'attention à tête unique à l'attention à plusieurs têtes

L'attention à plusieurs têtes consiste en termes pratiques à exécuter plusieurs instances de la fonction d'auto-attention, chacune avec ses propres poids, de sorte que différents vecteurs finaux soient calculés.

Exemple de code

Il pourrait être possible de réutiliser le code précédent et d'ajouter simplement un wrapper qui le lance plusieurs fois, mais voici une version plus optimisée de https://github.com/rasbt/LLMs-from-scratch/blob/main/ch03/01_main-chapter-code/ch03.ipynb qui traite toutes les têtes en même temps (réduisant le nombre de boucles for coûteuses). Comme vous pouvez le voir dans le code, les dimensions de chaque token sont divisées en différentes dimensions selon le nombre de têtes. De cette façon, si un token a 8 dimensions et que nous voulons utiliser 3 têtes, les dimensions seront divisées en 2 tableaux de 4 dimensions et chaque tête utilisera l'un d'eux :

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert (d_out % num_heads == 0), \
"d_out must be divisible by num_heads"

self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
self.dropout = nn.Dropout(dropout)
self.register_buffer(
"mask",
torch.triu(torch.ones(context_length, context_length),
diagonal=1)
)

def forward(self, x):
b, num_tokens, d_in = x.shape
# b is the num of batches
# num_tokens is the number of tokens per batch
# d_in is the dimensions er token

keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)

# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)

# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)

# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)

# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection

return context_vec

torch.manual_seed(123)

batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)

context_vecs = mha(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

Pour une implémentation compacte et efficace, vous pourriez utiliser la classe torch.nn.MultiheadAttention dans PyTorch.

Références

• https://www.manning.com/books/build-a-large-language-model-from-scratch