Apprendre Comment Implémenter l'Attention Multi-Tête | Construction des Composants de Transformer

Glissez pour afficher le menu


              12345678910111213141516171819202122232425262728293031323334353637383940
            
import numpy as np

# Toy input: 3 tokens, embedding size 6
x = np.array([
    [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
    [0.6, 0.5, 0.4, 0.3, 0.2, 0.1],
    [0.2, 0.1, 0.4, 0.3, 0.6, 0.5],
])

# Multi-head attention parameters
num_heads = 2
embed_size = x.shape[1]
head_dim = embed_size // num_heads

# Random weights for queries, keys, and values for each head
np.random.seed(42)
W_q = np.random.randn(num_heads, embed_size, head_dim)
W_k = np.random.randn(num_heads, embed_size, head_dim)
W_v = np.random.randn(num_heads, embed_size, head_dim)

attention_outputs = []

for head in range(num_heads):
    # Linear projections
    Q = x @ W_q[head]
    K = x @ W_k[head]
    V = x @ W_v[head]
    
    # Scaled dot-product attention
    scores = Q @ K.T / np.sqrt(head_dim)
    attention_weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    head_output = attention_weights @ V
    attention_outputs.append(head_output)

# Concatenate outputs from all heads
multi_head_output = np.concatenate(attention_outputs, axis=-1)

print("Output from head 0:\n", attention_outputs[0])
print("Output from head 1:\n", attention_outputs[1])
print("Concatenated multi-head output:\n", multi_head_output)

L'attention multi-tête fonctionne en appliquant le même mécanisme d'auto-attention indépendamment sur plusieurs têtes en parallèle. Chaque tête utilise son propre ensemble de matrices de poids pour les requêtes, clés et valeurs (W_q, W_k et W_v), projetant l'entrée dans des sous-espaces de dimension inférieure. L'entrée x dans cet exemple contient trois jetons, chacun avec une taille d'embedding de 6, et est divisée en deux têtes.

Pour chaque tête, le code répète les étapes standard de l'auto-attention : il calcule les matrices de requête (Q), de clé (K) et de valeur (V), calcule les scores d'attention en prenant le produit scalaire de Q et K.T (avec une mise à l'échelle appropriée), et utilise une fonction softmax pour obtenir les poids d'attention. Ces poids sont ensuite utilisés pour combiner les vecteurs de valeur, produisant la sortie de chaque tête.

Il n'y a pas de nouveau mécanisme ici, l'attention multi-tête est simplement le processus d'auto-attention effectué en parallèle pour chaque tête. Après le traitement de toutes les têtes, leurs sorties sont concaténées selon la dernière dimension. Cette approche permet au modèle de capturer des relations variées dans les données, car chaque tête peut se concentrer sur différents aspects de l'entrée. Les résultats affichés illustrent comment l'exécution du même mécanisme d'auto-attention plusieurs fois et la combinaison des résultats fournissent une représentation plus riche pour chaque jeton.

Tout était clair ?

Merci pour vos commentaires !

Section 2. Chapitre 2

Demandez à l'IA

Posez n'importe quelle question ou essayez l'une des questions suggérées pour commencer notre discussion

Section 2. Chapitre 2