Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

1. Énoncé du problème

Les grands modèles de langage (LLM) rencontrent un goulot d'étranglement fondamental dans l'inférence à contexte long en raison de la complexité quadratique ($O(L^2)$) du mécanisme d'attention lors de l'étape de préremplissage (prefill). Bien que les optimisations matérielles comme FlashAttention réduisent la surcharge des entrées/sorties mémoire, elles ne modifient pas la complexité algorithmique.

Les méthodes d'accélération existantes souffrent de deux limitations principales :

• Sparsité structurée : Des méthodes comme Minference ou FlexPrefill appliquent une sparsité au niveau des blocs. Cela conserve souvent des tokens non pertinents s'ils sont regroupés avec des tokens importants dans le même bloc, conduisant à une sparsité sous-optimale.
• Éviction de tokens : Des méthodes comme FastKV ou GemFilter éliminent définitivement des tokens dans les premières couches. Cela suppose que l'importance des tokens est statique, ignorant la nature dynamique de la pertinence des tokens à travers les différentes couches et têtes d'attention. Une fois évictés, un token ne peut plus être réexaminé, ce qui entraîne une dégradation de la précision.

2. Méthodologie : Token Sparse Attention

Les auteurs proposent Token Sparse Attention, un mécanisme de sparsification au niveau des tokens, léger, dynamique et réversible. Contrairement aux méthodes basées sur l'éviction, il ne supprime pas définitivement les tokens mais compresse l'espace de calcul avant de restaurer la séquence.

Architecture de base

La méthode opère en deux étapes par tête d'attention :

1. Compression (Sélection QKV) :
   • Pour chaque tête d'attention $h$, un sous-ensemble d'indices de tokens $S_h$ est sélectionné sur la base de scores d'importance.
   • Les tenseurs de requête ($Q$), de clé ($K$) et de valeur ($V$) sont rassemblés pour former des tenseurs compressés $\hat{Q}, \hat{K}, \hat{V}$ avec une longueur de séquence réduite $L' \ll L$.
   • Caractéristique clé : La sélection est effectuée indépendamment par tête, permettant à différentes têtes de se concentrer sur différents sous-ensembles de tokens, répondant ainsi à l'hétérogénéité tête par tête.
   • Les tenseurs compressés restent denses et contigus, permettant l'utilisation de noyaux optimisés comme FlashAttention.
2. Décompression (Restauration de la sortie) :
   • La sortie d'attention $\hat{O}$ (forme $L' \times d$) est dispersée à nouveau dans un tenseur initialisé à zéro de la longueur de séquence originale $L \times d$ en utilisant les indices spécifiques à la tête.
   • Caractéristique clé : Les positions non sélectionnées sont mises à zéro, mais la connexion résiduelle de la couche précédente préserve l'information des tokens sautés. Cela permet au modèle de "revisiter" et de réévaluer l'importance des tokens dans les couches suivantes.

Couverture dynamique des tokens

Pour déterminer quels tokens conserver et combien, le papier introduit la Couverture dynamique des tokens :

• Notation : Une carte d'attention proxy légère est calculée en utilisant les requêtes récentes contre toutes les clés. Les scores d'importance des tokens sont dérivés en sommant les poids d'attention verticalement.
• Allocation du budget : Au lieu de conserver les $k$ meilleurs tokens, la méthode trie les tokens par importance croissante et élague la "queue" (les tokens les moins importants) jusqu'à ce que la masse cumulative des tokens élagués dépasse un seuil de couverture $\tau$. Cela traite l'élagage comme une régularisation structurelle contre le bruit d'attention.
• Sélection tête par tête : Après avoir déterminé le budget global, chaque tête sélectionne indépendamment ses $k$ meilleurs tokens sur la base de ses scores spécifiques.

Sélection de couches clairsemées

Les auteurs observent que l'application de la sparsité à toutes les couches dégrade les performances. Ils introduisent la Dérive de représentation inter-couches ($R_\ell$) pour mesurer la stabilité des représentations de tokens entre les couches.

• Les couches avec une faible dérive (représentations stables) sont sélectionnées pour la sparsification.
• Les couches avec une forte dérive (où les représentations changent significativement) sont conservées denses pour préserver le comportement du modèle.
• Cette sélection est une étape de prétraitement unique par modèle.

3. Contributions clés

1. Sparsification réversible au niveau des tokens : Une conception novatrice "Compresser et Décompresser" qui réduit le coût de calcul sans éviction définitive des tokens, permettant une réévaluation dynamique de l'importance des tokens à travers les couches.
2. Hétérogénéité tête par tête : La méthode prend en charge une sélection de tokens indépendante pour chaque tête d'attention, capturant le fait que différentes têtes priorisent différentes caractéristiques sémantiques.
3. Compatibilité transparente : L'approche est entièrement compatible avec les noyaux denses existants (FlashAttention) et clairsemés (Minference, FlexPrefill), agissant comme un accélérateur complémentaire plutôt que comme un remplacement.
4. Budgétisation dynamique : Une approche pilotée par les données pour déterminer les niveaux de sparsité basés sur les distributions d'importance des tokens et la stabilité couche par couche, plutôt que sur des ratios fixes.

4. Résultats expérimentaux

Les expériences ont été menées sur LLaMA-3.1-8B-Instruct et Mistral-Nemo-12B-Instruct en utilisant des benchmarks tels que RULER, InfiniteBench et Needle-in-a-Haystack.

• Compromis Précision-Latence :
  • Token Sparse Attention améliore constamment le compromis vitesse-précision.
  • À un contexte de 128K, il atteint jusqu'à $\times 3,23$ d'accélération de l'attention par rapport à FlashAttention avec moins de 1 % de dégradation de la précision.
  • Combiné avec FlexPrefill, il améliore l'accélération de $\times 2,44$ à $\times 2,76$ tout en maintenant une précision moyenne identique (87,27 %).
• Comparaison avec les méthodes d'éviction :
  • Il surpasse significativement les méthodes d'éviction de tokens (PyramidInfer, FastKV, GemFilter). Par exemple, à 128K, Token Sparse Attention atteint 86,84 % de précision contre 78,49 % pour PyramidInfer à des accélérations similaires.
• Évolutivité :
  • Les gains d'accélération augmentent avec la longueur du contexte (par exemple, gain minime à 4K, substantiel à 128K) car la sparsité de l'attention augmente avec la longueur de la séquence.
• Surcharge :
  • La surcharge liée à la notation et à la gestion des indices est minimale, représentant <11 % de la latence totale d'attention même à des niveaux de sparsité élevés.

5. Importance

Ce papier comble une lacune critique dans l'inférence des LLM à contexte long en s'éloignant de l'éviction irréversible de tokens vers une sparsification dynamique et entrelacée.

• Insight théorique : Il démontre que l'importance des tokens est hautement dynamique à travers les couches et les têtes, nécessitant des mécanismes de sélection flexibles plutôt qu'un élagage statique.
• Impact pratique : Il fournit une optimisation "plug-and-play" qui peut être superposée aux noyaux d'attention de pointe existants (FlashAttention, Minference, etc.) pour obtenir des accélérations massives dans les scénarios à contexte long sans réentraîner les modèles ni sacrifier une précision significative.
• Direction future : Il établit un nouveau point de conception pour une inférence évolutive, suggérant que la compression réversible est une stratégie supérieure pour gérer le goulot d'étranglement $O(L^2)$ par rapport à la suppression permanente de tokens.