Les Mécanismes de l'Auto-Attention Causale

13 novembre 2024

L’auto-attention causale est le mécanisme à l’origine de la plupart des avancées en IA depuis 2017. Dans cet article, je vais détailler le calcul étape par étape pour, je l’espère, acquérir une meilleure intuition de son fonctionnement.

SelfAttention (Q, K, V) = softmax (mask (\frac{Q K ^{T}}{d})) V

À un haut niveau, cette fonction prend une séquence et la transforme en une autre. Une séquence est une liste d’embeddings de tokens, un tenseur de forme $L \times d$ , où $L$ est la longueur de la séquence d’entrée et $d$ la dimension de l’embedding. Chaque ligne de cette matrice correspond à un token d’entrée, représenté par un vecteur de dimension $d$ .

Mais alors, pourquoi y a-t-il 3 entrées à $SelfAttention$ ? C’est parce que, dans l’architecture Transformer, la séquence d’entrée est projetée par 3 couches linéaires $d \times d$ différentes. Si $X$ est la séquence d’entrée,

Q = X W_{Q}, K = X W_{K}, V = X W_{V}

où $W_{Q}, W_{K}, V$ sont de taille $d \times d$ . Ainsi, $Q, K, V$ sont simplement des représentations différentes de la même séquence d’entrée.

Calculons $SelfAttention$ étape par étape. D’abord, nous effectuons $Q K^{T}$ , qui est un produit scalaire $L \times d$ par $d \times L$ , donnant un résultat $L \times L$ . Que fait cette opération ?

Q K^{T} = q_{1} q_{2} ⋮ q_{L} [k_{1}^{T} k_{2}^{T} \dots k_{L}^{T}] = q_{1} k_{1}^{T} q_{2} k_{1}^{T} ⋮ q_{L} k_{1}^{T} q_{1} k_{2}^{T} q_{2} k_{2}^{T} ⋮ q_{L} k_{2}^{T} \dots \dots ⋱ \dots q_{1} k_{L}^{T} q_{2} k_{L}^{T} ⋮ q_{L} k_{L}^{T}

Le résultat de $q_{i} k_{j}^{T}$ est un scalaire ( $1 \times d$ point $d \times 1$ ), et c’est le produit scalaire vectoriel entre $q_{i}$ et $k_{j}$ . Si l’on se souvient de la formule

a \cdot b = ∥ a ∥∥ b ∥ cos θ

on voit que le produit scalaire est positif lorsque $θ$ , l’angle entre $a$ et $b$ , est proche de 0º, et négatif lorsque l’angle est de 180º, c’est-à-dire lorsqu’ils pointent dans des directions opposées. On peut interpréter le produit scalaire comme une mesure de similarité, où les valeurs positives indiquent des vecteurs similaires et les valeurs négatives indiquent l’opposé.

Ainsi, notre matrice finale $L \times L$ est remplie de scores de similarité entre chaque paire de tokens $q$ et $k$ . Le résultat est divisé par $d$ pour éviter que la variance n’explose pour de grandes dimensions d’embedding. Voir l’annexe pour plus de détails.

L’étape suivante consiste à appliquer la fonction $mask$ , qui définit toutes les valeurs qui ne sont pas dans la section triangulaire inférieure de la matrice d’entrée à $- \infty$ .

mask (\frac{1}{d} Q K^{T}) = \frac{1}{d} q_{1} k_{1}^{T} q_{2} k_{1}^{T} q_{3} k_{1}^{T} ⋮ q_{L} k_{1}^{T} - \infty q_{2} k_{2}^{T} q_{3} k_{2}^{T} ⋮ q_{L} k_{2}^{T} - \infty - \infty q_{3} k_{3}^{T} ⋮ q_{L} k_{3}^{T} \dots \dots \dots ⋱ \dots - \infty - \infty - \infty ⋮ q_{L} k_{L}^{T}

À cela, nous appliquons $softmax$ , qui convertit chaque ligne de valeurs de la matrice en une distribution de probabilité. La fonction est définie comme une application de $R^{L} \to R^{L}$ , où le $i$ ème élément de sortie est donné par

softmax (x)_{i} = \frac{e ^{x_{i}}}{\sum _{j = 1}^{L} e ^{x_{j}}} pour i = 1, 2, \dots, L

Deux choses à noter ici :

La somme de tous les éléments de sortie est $1$ , comme attendu pour une distribution de probabilité.
Si un élément d’entrée $x_{i}$ est $- \infty$ , alors $softmax (x)_{i} = 0$ .

Après avoir appliqué la fonction $softmax$ aux scores de similarité masqués, nous obtenons :

S = softmax (mask (\frac{1}{d} Q K^{T})) = S_{1, 1} S_{2, 1} S_{3, 1} ⋮ S_{L, 1} 0 S_{2, 2} S_{3, 2} ⋮ S_{L, 2} 00 S_{3, 3} ⋮ S_{L, 3} \dots \dots \dots ⋱ \dots 000 ⋮ S_{L, L}

Où les entrées $S_{i, j}$ sont définies comme :

S_{i, j} = \frac{e ^{mask (\frac{Q K ^{T}}{d})_{i, j}}}{\sum _{k = 1}^{L} e ^{mask (\frac{Q K ^{T}}{d})_{i, k}}}

La matrice résultante $S$ a des lignes qui sont des distributions de probabilité de longueur $L$ . La dernière étape consiste à mapper notre matrice de valeurs $V$ par ces distributions de probabilité pour obtenir notre nouvelle séquence.

SelfAttention (Q, K, V) = SV = S_{1, 1} S_{2, 1} S_{3, 1} ⋮ S_{L, 1} 0 S_{2, 2} S_{3, 2} ⋮ S_{L, 2} 00 S_{3, 3} ⋮ S_{L, 3} \dots \dots \dots ⋱ \dots 000 ⋮ S_{L, L} V_{1} V_{2} V_{3} ⋮ V_{L} = S_{1, 1} V_{1} S_{2, 1} V_{1} + S_{2, 2} V_{2} S_{3, 1} V_{1} + S_{3, 2} V_{2} + S_{3, 3} V_{3} ⋮ S_{L, 1} V_{1} + S_{L, 2} V_{2} + \dots + S_{L, L} V_{L}

Notez que $S_{i, j}$ est un scalaire et $V_{k}$ est un vecteur d’embedding $1 \times d$ . Visuellement, on observe que SelfAttention combine sélectivement les tokens de Valeur, pondérés par une distribution de probabilité générée par la façon dont les requêtes et les clés s’attendent mutuellement, c’est-à-dire ont un grand produit scalaire. On voit aussi que le poids d’un token de sortie à l’indice $i$ ne dépend que des tokens d’entrée d’indice $\leq i$ , en raison du masque causal appliqué précédemment. Ceci est basé sur l’hypothèse causale, qu’un token de sortie $O_{i}$ ne dépend pas des tokens futurs, ce qui est requis lors de l’entraînement de modèles autorégressifs (c’est-à-dire de prédiction du token suivant).

J’espère que vous avez trouvé cela utile !

Annexe

Pourquoi diviser par $d$ ?

Nous faisons cela pour éviter que la variance n’explose lorsque $d$ augmente.

Supposons que $q_{i}, k_{i} \sim N (μ = 0, σ^{2} = 1)$ et soient i.i.d. Calculons l’espérance et la variance du produit scalaire non normalisé $s = q \cdot k$ .

L’espérance est trivialement nulle :

E [s] = E [i = 1 \sum d q_{i} k_{i}] = i = 1 \sum d E [q_{i} k_{i}] = i = 1 \sum d E [q_{i}] E [k_{i}] = 0

Et la variance est :

Var (s) = E [s^{2}] - (E [s])^{2} = E [s^{2}] = d

car

E [s^{2}] = E [i = 1 \sum d j = 1 \sum d q_{i} k_{i} q_{j} k_{j}] = i = 1 \sum d j = 1 \sum d E [q_{i} k_{i} q_{j} k_{j}]

ce qui vaut $0$ pour $i \neq = j$ (puisque $q_{i}, q_{j}$ et $k_{i}, k_{j}$ sont i.i.d). Pour $i = j$ ,

i = 1 \sum d E [q_{i}^{2} k_{i}^{2}] = i = 1 \sum d E [q_{i}^{2}] E [k_{i}^{2}] = i = 1 \sum d 1 \cdot 1 = d

puisque $E [q_{i}^{2}] = E [k_{i}^{2}] = σ^{2} = 1$ .

Ainsi, si nous divisons par $1/ d$ , notre nouvelle variance est

Var (\frac{s}{d}) = \frac{1}{d} Var (s) = 1

comme souhaité.

Attention Multi-Têtes

La plupart des systèmes modernes utilisent l’attention multi-têtes, qui calcule $SelfAttention$ en parallèle sur plusieurs “têtes”. On pose généralement $d_{k} = d_{v} = d_{model} / H$ , où $H$ est le nombre de têtes.

Q_{h} K_{h} V_{h} = X W_{h}^{Q} = X W_{h}^{K} = X W_{h}^{V} W_{h}^{Q} \in R^{d_{model} \times d_{k}} W_{h}^{K} \in R^{d_{model} \times d_{k}} W_{h}^{V} \in R^{d_{model} \times d_{v}}

head_{h} = SelfAttention (Q_{h}, K_{h}, V_{h}) = softmax (mask (\frac{Q _{h} K _{h}^{T}}{d _{k}})) V_{h}

MultiHead (Q, K, V) = Concat (head_{1}, head_{2}, \dots, head_{H})

✦ Aucune IA n'a été utilisée dans la conception, la recherche, la rédaction ou l'édition de cet article.

Les Mécanismes de l'Auto-Attention Causale

Annexe

Pourquoi diviser par d​ ?

Attention Multi-Têtes

Pourquoi diviser par $d$ ?