Les Mécanismes de l'Attention Causale

13 novembre 2024

Introduction

L’attention causale est le mécanisme à la base de la plupart des avancées en IA depuis 2017. Dans cet article, je vais détailler le calcul et, espérons-le, acquérir une meilleure intuition de son fonctionnement.

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

À un niveau élevé, 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$ est la dimension de l’embedding. Chaque ligne de cette matrice correspond à un token d’entrée, représenté comme un vecteur de dimension $d$ .

Alors pourquoi y a-t-il 3 entrées pour $SelfAttention$ ? C’est parce que, dans l’architecture Transformer, la séquence d’entrée est projetée par 3 couches linéaires différentes de taille $d \times d$ . 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}, W_{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 de $L \times d$ par $d \times L$ , résultant en une sortie de taille $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$ produit scalaire $d \times 1$ ), et il représente le produit scalaire entre $q_{i}$ et $k_{j}$ . Si nous nous rappelons de la formule

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

nous voyons 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º, ou lorsqu’ils pointent dans des directions opposées. Nous pouvons interpréter le produit scalaire comme une métrique de similarité, où les valeurs positives indiquent des vecteurs similaires, et les valeurs négatives indiquent le contraire.

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 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 de distribution 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 de taille $1 \times d$ . Visuellement, nous observons que SelfAttention combine sélectivement les tokens de valeur, pondérés par une distribution de probabilité générée par la manière dont les requêtes et les clés s’attendent mutuellement, c’est-à-dire ont un grand produit scalaire. Nous voyons également que le poids d’un token de sortie à l’indice $i$ ne dépend que des tokens d’entrée avec un indice $\leq i$ , en raison du masque causal que nous avons appliqué précédemment. Cela repose sur l’hypothèse causale, selon laquelle un token de sortie $O_{i}$ ne dépend pas des tokens futurs, ce qui est nécessaire lors de l’entraînement de modèles autoregressifs (c’est-à-dire la 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 i.i.d. Calculons la moyenne et la variance de $s = q \cdot k$ non divisé.

La moyenne 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}]

qui est $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”. Nous laissons 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})

←

Un bot expert pour 2048

Approximation Locale

→

Les Mécanismes de l'Attention Causale

Introduction

Annexe

Pourquoi diviser par d​ ?

Attention Multi-Têtes

Pourquoi diviser par $d$ ?