Attention Sinks

Un puits d'attention est un token qui reçoit une grande part du poids de l'attention dans de nombreuses têtes et couches, même si son contenu n'est pas pertinent pour la requête ; il s'agit le plus souvent du premier token de la séquence. Guangxuan Xiao et ses collègues du MIT, de Meta AI et de l'université Carnegie Mellon ont nommé cet effet dans Efficient Streaming Language Models with Attention Sinks, publié à ICLR 2024. Ils ont montré que les grands modèles de langage (LLMs), comme Llama 2, MPT, Falcon et Pythia, se dégradent lorsqu'un cache KV glissant évince les premiers tokens, et que leur en conserver quatre rétablit une génération stable.

Pourquoi les puits se forment#

L'auto-attention transforme les scores de chaque requête en poids à l'aide d'une softmax, et la somme de ces poids est toujours égale à un. Lorsqu'aucun token du contexte n'est pertinent pour une tête, le poids doit tout de même être attribué quelque part. Dans un modèle de langage causal, les premiers tokens sont visibles par tous les tokens ultérieurs ; l'entraînement en fait donc la destination par défaut du poids dont une tête n'a pas besoin. Dans l'analyse de StreamingLLM sur Llama 2, toutes les couches au-delà des deux premières concentraient l'attention sur les tokens initiaux dans toutes les têtes. L'ajout d'un logit de puits à la softmax montre cet effet :

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

Sans puits, la tête répartit tout son poids entre six valeurs non pertinentes et produit une sortie de norme 0,83. Avec un logit de puits égal à 3, le puits reçoit 77 % du poids et la sortie tombe à 0,19 ; la tête ne fait donc presque rien pour cette requête. Les modèles à poids ouverts gpt-oss d'OpenAI intègrent cette fonction grâce à un logit de puits appris pour chaque tête d'attention, qui offre à chacune une destination pour le poids excédentaire sans effet sur sa sortie.

Génération en continu avec StreamingLLM#

Un modèle qui génère indéfiniment finit par dépasser toute fenêtre de contexte et ne peut pas conserver chaque token dans son cache KV ; l'attention à fenêtre simple, qui ne conserve que les tokens les plus récents, échoue dès que les premiers tokens sont évincés. StreamingLLM conserve en permanence les clés et valeurs de quatre tokens initiaux, à côté d'une fenêtre glissante de tokens récents, et attribue les positions en fonction de leur emplacement dans le cache plutôt que dans le texte d'origine, en recalculant les intégrations positionnelles rotatives (RoPE) à chaque étape :

Cette méthode a maintenu la stabilité de la modélisation du langage sur plus de 4 millions de tokens et a été jusqu'à 22,2 fois plus rapide qu'une fenêtre glissante qui recalcule le cache pour chaque nouveau token. La fenêtre limite toujours ce dont le modèle peut se souvenir : les tokens évincés du milieu sont perdus à jamais, de sorte qu'un agent conversationnel ou un agent d'IA en continu ne se souvient que de ses échanges les plus récents. Les auteurs ont aussi préentraîné des modèles avec un seul token de substitution dédié et apprenable au début de chaque échantillon, qui a pris à lui seul le rôle de puits.

Puits d'attention et termes associés#

• Attention à fenêtre glissante : Limite chaque token à ses voisins récents. Appliquée à l'inférence d'un modèle entraîné avec une attention complète, elle échoue dès que les puits sortent de la fenêtre.
• Dégradation du contexte : Baisse de précision à mesure qu'un prompt s'allonge, même s'il tient encore dans la fenêtre de contexte.

Puits dans les modèles de vision#

Les transformeurs de vision (ViTs) présentent un comportement apparenté. Dans Vision Transformers Need Registers, des chercheurs de Meta et d'Inria ont découvert des tokens de norme inhabituellement élevée dans les zones d'arrière-plan peu informatives de ViTs supervisés et auto-supervisés, réaffectés à des calculs internes. L'ajout de quelques tokens de registre à l'entrée a absorbé ce rôle, supprimé les artefacts des cartes de caractéristiques et des cartes d'attention, et amélioré la prédiction dense.

SAM 2 de Meta, exécuté dans le paquet Ultralytics, gère sa mémoire vidéo selon un plan similaire, qui conserve les points d'ancrage et une fenêtre : pour chaque nouvelle trame, son attention mémoire lit les entrées de la banque mémoire correspondant aux trames indiquées par l'utilisateur et conservées pendant toute la vidéo, ainsi que celles des 6 trames les plus récentes, ce qui maintient constant le coût de l'attention par trame, quelle que soit la durée de la vidéo.