Chapitre 2 : Auto-attention et attention multi-têtes

Introduction

En nous appuyant sur notre compréhension du mécanisme d'attention, ce chapitre explore comment l'attention est appliquée lorsqu'une séquence s'attend elle-même (auto-attention) et pourquoi l'utilisation de plusieurs têtes d'attention en parallèle améliore considérablement la capacité du modèle.

Auto-attention : Le cœur des Transformers

Qu'est-ce que l'auto-attention ?

Dans l'auto-attention, une séquence s'attend elle-même. Chaque position dans l'entrée peut s'attendre à chaque autre position dans la même entrée.

Entrée : "Le chat est assis sur le tapis"

Auto-attention : Chaque mot s'attend à TOUS les mots
├─ "Le" s'attend à ["Le", "chat", "est", "assis", "sur", "le", "tapis"]
├─ "chat" s'attend à ["Le", "chat", "est", "assis", "sur", "le", "tapis"]
├─ "est" s'attend à ["Le", "chat", "est", "assis", "sur", "le", "tapis"]
└─ ... et ainsi de suite

Pourquoi l'auto-attention ?

L'auto-attention permet :

1. Dépendances à longue portée : Connexions directes entre n'importe quelles deux positions
2. Traitement parallèle : Toutes les positions calculées simultanément
3. Contexte dynamique : Différents poids pour différentes entrées

Formulation mathématique

Pour l'auto-attention, Q, K, et V proviennent tous de la même entrée :

$$\text{SelfAttention}(X) = \text{softmax}\left(\frac{XW_Q(XW_K)^T}{\sqrt{d_k}}\right)XW_V$$

Implémentation

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    """
    Couche d'auto-attention où Q, K, V proviennent tous de la même entrée.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k or d_model
        
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k
    
    def forward(self, x, mask=None):
        """
        x : Tenseur d'entrée (batch, seq_len, d_model)
        mask : Masque d'attention optionnel
        """
        # Toutes les projections de la même entrée
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
        
        # Calculer les scores d'attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Comparaison des types d'attention

Auto-attention vs. Cross-attention

Type	Source de Q	Source de K, V	Cas d'utilisation
Auto-attention	Entrée X	Entrée X	Encodeur, Décodeur
Cross-attention	Décodeur	Encodeur	Modèles encodeur-décodeur

Exemple de cross-attention

class CrossAttention(nn.Module):
    """
    Cross-attention où Q provient d'une source, K et V d'une autre.
    Utilisé dans les architectures encodeur-décodeur.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k or d_model
        
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k
    
    def forward(self, query_input, key_value_input, mask=None):
        """
        query_input : Source pour Q (par exemple, état du décodeur)
        key_value_input : Source pour K, V (par exemple, sortie de l'encodeur)
        """
        Q = self.W_Q(query_input)
        K = self.W_K(key_value_input)
        V = self.W_V(key_value_input)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Attention multi-têtes

La motivation

Une seule tête d'attention ne peut capturer qu'un type de relation. L'attention multi-têtes exécute plusieurs opérations d'attention en parallèle, permettant au modèle de s'attendre conjointement à des informations provenant de différents sous-espaces de représentation.

Intuition

Différentes têtes peuvent apprendre différents motifs :

• Tête 1 : Relations syntaxiques (sujet-verbe)
• Tête 2 : Relations positionnelles (mots proches)
• Tête 3 : Relations sémantiques (synonymes)
• Tête 4 : Coréférence (pronoms vers noms)

Définition mathématique

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$$

Où chaque tête est :

$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

Dimensions des paramètres

Pour $h$ têtes avec une dimension de modèle $d_{model}$ :

• Chaque tête a une dimension $d_k = d_v = d_{model} / h$
• Le nombre total de paramètres reste à peu près le même que pour l'attention à une seule tête

# Exemple : d_model = 512, h = 8
d_model = 512
num_heads = 8
d_k = d_model // num_heads  # 64 par tête

Implémentation complète

class MultiHeadAttention(nn.Module):
    """
    Mécanisme d'attention multi-têtes.
    """
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model doit être divisible par num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Projections linéaires pour Q, K, V (toutes les têtes combinées)
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        
        # Projection de sortie
        self.W_O = nn.Linear(d_model, d_model, bias=False)
        
        self.dropout = nn.Dropout(dropout)
    
    def split_heads(self, x, batch_size):
        """
        Diviser la dernière dimension en (num_heads, d_k).
        Transposer pour obtenir la forme (batch, num_heads, seq_len, d_k)
        """
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Projections linéaires
        Q = self.W_Q(query)
        K = self.W_K(key)
        V = self.W_V(value)
        
        # Diviser en têtes
        Q = self.split_heads(Q, batch_size)  # (batch, heads, seq, d_k)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)
        
        # Attention produit scalaire pondéré pour toutes les têtes en parallèle
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            # Étendre le masque pour la dimension des têtes
            if mask.dim() == 3:
                mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        
        # Appliquer l'attention aux valeurs
        context = torch.matmul(attention_weights, V)
        
        # Concaténer les têtes
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.d_model)
        
        # Projection linéaire finale
        output = self.W_O(context)
        
        return output, attention_weights

Utilisation de l'attention multi-têtes

# Exemple d'utilisation
d_model = 512
num_heads = 8
seq_len = 100
batch_size = 32

mha = MultiHeadAttention(d_model, num_heads)

# Entrée (même pour l'auto-attention)
x = torch.randn(batch_size, seq_len, d_model)

# Auto-attention
output, attn_weights = mha(x, x, x)
print(f"Forme de la sortie : {output.shape}")  # (32, 100, 512)
print(f"Forme de l'attention : {attn_weights.shape}")  # (32, 8, 100, 100)

Avantages de la parallélisation

Pourquoi les Transformers s'entraînent rapidement

Contrairement aux RNN, toutes les positions dans un Transformer peuvent être calculées simultanément :

# RNN : Séquentiel (lent)
for t in range(seq_len):
    h_t = rnn_cell(x[t], h_{t-1})

# Transformer : Parallèle (rapide)
output = multi_head_attention(x, x, x)  # Toutes les positions en même temps

Efficacité GPU

L'attention multi-têtes est particulièrement adaptée aux GPU :

# Toutes les têtes calculées dans une seule multiplication matricielle par lots
# Forme : (batch, heads, seq, d_k) @ (batch, heads, d_k, seq)
# → (batch, heads, seq, seq)
scores = torch.matmul(Q, K.transpose(-2, -1))  # Parallèle sur toutes les têtes !

Visualisation des multiples têtes

def visualize_multi_head_attention(attention_weights, tokens, num_heads_to_show=4):
    """
    Visualiser les motifs d'attention de plusieurs têtes.
    attention_weights : (batch, heads, seq, seq)
    """
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(1, num_heads_to_show, figsize=(4*num_heads_to_show, 4))
    
    for head_idx in range(num_heads_to_show):
        ax = axes[head_idx]
        attn = attention_weights[0, head_idx].detach().numpy()  # Premier élément du lot
        
        im = ax.imshow(attn, cmap='Blues')
        ax.set_xticks(range(len(tokens)))
        ax.set_yticks(range(len(tokens)))
        ax.set_xticklabels(tokens, rotation=45, ha='right')
        ax.set_yticklabels(tokens)
        ax.set_title(f'Tête {head_idx + 1}')
    
    plt.tight_layout()
    plt.show()

Motifs d'attention en pratique

La recherche a montré que différentes têtes apprennent des motifs spécialisés :

Type de motif	Ce qu'il capture	Exemple
Diagonal	Même position / identité	Le token s'attend à lui-même
Token précédent	Relations séquentielles	Motifs de modèle de langage
Premier token	Agrégation globale	Token CLS dans BERT
Syntaxique	Structure grammaticale	Accord sujet-verbe
Sémantique	Relations de sens	Synonymes, antonymes

Points clés à retenir

Résumé

1. L'auto-attention permet à chaque position de s'attendre à toutes les positions dans la même séquence
2. La cross-attention permet de s'attendre à une séquence différente (encodeur-décodeur)
3. L'attention multi-têtes exécute plusieurs opérations d'attention en parallèle
4. Différentes têtes apprennent différents types de relations
5. La parallélisation rend les Transformers efficaces sur les GPU

À venir

Dans le prochain chapitre, nous explorerons comment les Transformers gèrent l'information positionnelle grâce à l'encodage positionnel, et comment les tokens sont représentés grâce aux embeddings.

---

← Chapitre 1 : Mécanisme d'attention

Chapitre 3 : Encodage positionnel →