Qu'est-ce que le masque de padding de clé par rapport au masque d'attention dans PyTorch?

Dans PyTorch, lors de l'utilisation de transformateurs et d'attention, il est courant d'utiliser des masques pour traiter correctement des séquences de différentes longueurs et empêcher l'attention sur des positions indésirables.

• Masque de padding clé (Key Padding Mask) : c'est un masque qui indique quelles positions dans la séquence sont du padding (remplissage), afin que le modèle ne les prenne pas en compte lors du calcul de l'attention. Il s'agit généralement d'un tenseur booléen où True indique une position de padding. Il est appliqué aux clés (keys) et valeurs (values) dans le mécanisme d'attention pour ignorer les tokens de padding.

• Masque d'attention (Attention Mask) : un masque plus général qui contrôle quelles positions peuvent influencer la position courante lors du calcul de l'attention. Par exemple, dans les tâches de génération automatique de texte, un masque empêche de regarder en avant (tokens futurs) pour éviter la fuite d'informations. Ce masque peut être une matrice triangulaire, où certaines connexions sont interdites.

Exemple:

import torch
from torch.nn import MultiheadAttention

# Supposons que nous ayons une séquence avec padding
key_padding_mask = torch.tensor([[False, False, True, True]])  # True pour padding

# Masque d'attention pour empêcher l'attention aux tokens futurs
attn_mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()  # Masque triangulaire supérieure

mha = MultiheadAttention(embed_dim=8, num_heads=2)
query = torch.rand(4, 1, 8)  # (longueur de séquence, lot, dimension d'embedding)
key = value = query

output, attn_weights = mha(query, key, value, attn_mask=attn_mask, key_padding_mask=key_padding_mask)

Ainsi, le masque de padding clé est utilisé pour ignorer les tokens de padding, et le masque d'attention pour contrôler la structure de l'attention (par exemple, empêcher l'attention aux positions futures).