Sliding Window Attention

L’attention à fenêtre glissante (SWA) limite chaque jeton à une fenêtre fixe de jetons voisins — les w jetons les plus récents dans un modèle de langage causal — plutôt qu’à toute la séquence comme dans l’auto-attention complète. Le calcul de l’attention croît alors de façon linéaire avec la longueur de la séquence plutôt que de façon quadratique, et un modèle génératif n’a besoin de conserver dans son cache KV que les clés et les valeurs des w derniers jetons. Longformer a fait de la fenêtre glissante le cœur de son mécanisme d’attention en 2020, puis le Mistral 7B de Mistral AI l’a intégrée à un grand modèle de langage (LLM) ouvert largement utilisé en 2023, avec une fenêtre de 4 096 jetons.

Comment la fenêtre s’élargit avec la profondeur#

Une couche de SWA ne voit que w jetons, mais les couches de transformeurs s’empilent. À la couche 2, un jeton lit des états cachés qui résument déjà chacun w jetons de la couche 1 ; l’information peut donc se propager sur environ w − 1 positions supplémentaires à chaque couche, tout comme des convolutions empilées élargissent le champ réceptif d’un réseau de neurones convolutif (CNN). Dans PyTorch, la fenêtre est un masque d’attention booléen, et multiplier ce masque par lui-même permet de suivre jusqu’où l’information se propage :

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

Pour 16 jetons, la fenêtre calcule 58 scores par couche au lieu de 136 ; la portée du dernier jeton augmente de 3 jetons par couche, passant de 4 à 7, 10 et 13. Le même masque peut être transmis sous la forme de attn_mask à scaled_dot_product_attention de PyTorch. Pour Mistral 7B, 32 couches avec une fenêtre de 4 096 jetons donnent une portée théorique de l’attention d’environ 131K jetons, même si l’information provenant de loin doit traverser de nombreux relais avant d’arriver.

Mistral 7B stocke également ses clés et ses valeurs dans un tampon circulaire de w emplacements : il écrit la position i dans l’emplacement i mod w et écrase l’entrée la plus ancienne. Pour une séquence de 32K jetons, cela a divisé par huit la mémoire du cache sans réduire la qualité ; la taille fixe du cache maintient la mémoire constante, quelle que soit la durée de la génération.

Entrelacement des couches locales et globales#

Les fenêtres strictement locales ne peuvent pas copier une information ancienne en une seule étape ; les modèles récents alternent donc des couches SWA et des couches globales complètes. Gemma 2 de Google les alternait une à une avec une fenêtre de 4 096 tokens, et Gemma 3 utilise cinq couches locales pour chaque couche globale avec une fenêtre de seulement 1 024 tokens, ce qui garde le cache KV petit dans une fenêtre de contexte pouvant atteindre 128K tokens. Seules les couches globales stockent les clés et les valeurs de l’ensemble du contexte, de sorte que cinq couches sur six ne contiennent chacune que 1 024 tokens.

Les transformeurs de vision utilisent le même mélange pour les images haute résolution, où chaque patch est un token. Le SAM 3 de Meta, disponible dans le package Ultralytics pour la segmentation d’images, applique l’attention à l’intérieur de fenêtres de patches non chevauchantes de 24 × 24 dans 28 des 32 blocs de son encodeur d’image, et globalement sur toute la grille de patches de 72 × 72 d’une image de 1 008 pixels dans les blocs 7, 15, 23 et 31.

SWA et notions apparentées#

• Swin Transformer : Divise l’image en fenêtres fixes et non chevauchantes, communes à tous les tokens de chaque fenêtre, puis décale la grille entre les couches afin que les fenêtres voisines échangent des informations.
• Attention éparse : La famille plus large de motifs fixes. SWA en est le membre le plus simple ; d’autres méthodes ajoutent des fenêtres dilatées, des tokens globaux ou des liens aléatoires.
• Puits d’attention : Les premiers tokens d’une séquence, qui absorbent le poids d’attention restant. Lorsqu’une fenêtre est appliquée à l’inférence à un modèle entraîné avec une attention complète, leur éviction interrompt la génération ; StreamingLLM les conserve donc dans le cache à côté de la fenêtre.
• Attention déformable : Chaque requête porte son attention sur quelques points d’échantillonnage dont elle apprend les décalages ; le voisinage dépend donc du contenu.
• Attention Flash : Un noyau GPU pour calculer l’attention exacte. Ses implémentations acceptent une taille de fenêtre ; les modèles SWA ignorent donc complètement les blocs situés hors de la fenêtre.