Résumé technique : SOCKET - Estimateur de noyau de collision souple pour l'attention parcimonieuse

Énoncé du problème

Les grands modèles de langage (LLM) rencontrent une limitation fondamentale d'évolutivité lors de l'inférence sur des contextes longs. Alors que la phase de préremplissage traite la séquence d'entrée complète, la phase de décodage itérative exige que chaque token nouvellement généré porte attention à tous les tokens précédemment mis en cache. À mesure que la longueur du contexte augmente (de 2 000 à plusieurs millions de tokens), ce mécanisme d'attention dense devient de plus en plus limité par la mémoire, dépassant souvent la capacité des GPU et nécessitant des transferts coûteux entre le CPU et le GPU.

Les méthodes d'attention parcimonieuse existantes tentent d'atténuer ce problème en restreignant le calcul à un sous-ensemble de tokens. Cependant, beaucoup reposent sur des heuristiques dépendantes des données (par exemple, le clustering k-moyennes) qui entraînent des coûts élevés de temps jusqu'au premier token (TTFT) et manquent de robustesse face aux décalages de distribution. D'autres utilisent le hachage sensible à la localité (LSH) traditionnel, qui attribue des signaux de collision binaires. L'article soutient que cette formulation « dure » est mal adaptée à la stabilité du classement ; les collisions binaires ne peuvent exprimer une similarité partielle, conduisant à des classements grossiers et bruyants qui dégradent la qualité de la récupération, en particulier sous des budgets mémoire stricts.

Méthodologie : SOCKET

Les auteurs introduisent SOCKET (Estimateur de noyau de collision souple), un mécanisme d'attention parcimonieuse agnostique aux données qui reformule le LSH d'un générateur de candidats binaire en un noyau de notation principé et conscient de la similarité.

Mécanisme central : LSH souple

Contrairement au LSH traditionnel, qui attribue une requête à un seul bac discret par table de hachage, SOCKET emploie un LSH souple :

1. Précalcul (Préremplissage) : Les clés sont hachées dans $L$ tables de hachage indépendantes à l'aide de projections aléatoires. Chaque clé se voit attribuer un seul identifiant de bac par table, et les normes des vecteurs de valeur sont mises en cache.
2. Attribution souple (Décodage) : Lorsqu'une nouvelle requête arrive, elle n'est pas attribuée à un seul bac. Au lieu de cela, elle induit une distribution de probabilité sur tous les bacs de chaque table de hachage. Cela est réalisé en calculant un score de hachage souple basé sur la projection de la requête sur des hyperplans aléatoires, passé à travers une fonction softmax contrôlée par une température.
3. Notation : Pour chaque clé, un score continu est calculé en agrégeant la masse de probabilité attribuée au bac spécifique de la clé à travers toutes les $L$ tables. Ce score est pondéré par la norme du vecteur de valeur de la clé.
4. Sélection : Les $k$ clés principales sont sélectionnées de manière déterministe sur la base de ces scores souples, et une attention exacte est calculée sur ce sous-ensemble.

Fondement théorique

L'article fournit des garanties théoriques montrant que l'attention calculée à l'aide de scores LSH souples approxime un noyau étroitement lié à l'attention softmax (spécifiquement, un noyau angulaire).

• Stabilité : L'analyse théorique (Théorème 3) démontre que le bacisation souple introduit un biais contrôlable ($\epsilon_\tau$) qui s'annule lorsque la température $\tau \to 0$, tout en préservant la régularité.
• Corrélation : Le Lemme 4 et les résultats empiriques montrent que les scores souples préservent mieux la structure directionnelle de la similarité requête-clé projetée que les scores durs. Le LSH dur repose sur la magnitude $\ell_1$ et rejette l'information directionnelle à l'intérieur des orthants, conduisant à des classements instables pour un $L$ fini. Le LSH souple maintient une relation linéaire avec les coordonnées projetées dans le régime de faible signal, résultant en une corrélation plus élevée avec la similarité réelle et une variance d'estimateur plus faible.

Contributions clés

1. LSH souple comme classeur stable : Les auteurs proposent un mécanisme de LSH souple agnostique aux données qui sert de fonction de classement nettement plus efficace que le LSH dur traditionnel. Il remplace les signaux de collision binaires par une agrégation probabiliste et consciente de la similarité, permettant une sélection stable des $k$ premiers sans vote ad hoc.
2. Efficacité via des noyaux personnalisés : Un noyau de notation CUDA personnalisé et un backend Flash Decode Triton sont implémentés. Cela permet à SOCKET d'atteindre un débit jusqu'à 1,5 fois supérieur à celui de FlashAttention lors du décodage sur des contextes longs, tout en maintenant ou dépassant la précision des méthodes parcimonieuses précédentes.
3. Conception agnostique aux données : Contrairement aux approches basées sur le clustering, SOCKET ne nécessite aucun réentraînement, aucune calibration ni aucun prétraitement hors ligne. Cela permet un TTFT plus rapide et une robustesse face aux décalages de distribution.
4. Garanties théoriques : L'article établit des bornes d'erreur de bout en bout pour l'estimateur d'attention par comptage souple, décomposant l'erreur en variance d'échantillonnage, erreur d'approximation de table finie et biais de bacisation souple.

Résultats expérimentaux

Les auteurs ont évalué SOCKET sur les modèles Llama-3.1-8B-Instruct, Llama-3.2-1B et Qwen3 (4B, 8B, 30B) à travers deux benchmarks : LongBench (tâches réelles) et RULER (tâches de diagnostic synthétiques).

• Précision : Sur LongBench, SOCKET a constamment surpassé des lignes de base comme Quest, PQCache, Double Sparsity et MagicPig. Par exemple, sur Llama-3.1-8B à une parcimonie de 33×, SOCKET a obtenu un score moyen de 47,83, surpassant la ligne de base la plus forte (Quest) d'environ 0,84 point. Sur RULER-32K, il a égalé ou dépassé les méthodes précédentes, en particulier à des niveaux de parcimonie élevés (par exemple, 50×).
• Efficacité : SOCKET a démontré des gains de débit significatifs sur les GPU NVIDIA A100 et H200. Sur le H200, il a surpassé FlashAttention-2 et FlashAttention-3 à mesure que les longueurs de contexte augmentaient, réalisant un accélération de 1,58× par rapport à FlashAttention-2 à 100K tokens et 1,58× par rapport à FlashAttention-3 à 140K tokens (sous une parcimonie de 33×).
• Surcharge mémoire : La méthode introduit une surcharge mémoire modeste d'environ 15 % au-delà du cache KV standard (stockage des indices de table de hachage et des normes de valeur), ce qui est nettement inférieur à la surcharge requise par le LSH dur pour atteindre une précision comparable (ce qui nécessiterait beaucoup plus de tables de hachage).

Importance et affirmations

L'article affirme que SOCKET adresse le goulot d'étranglement critique de l'inférence des LLM sur des contextes longs en offrant une alternative principée et agnostique aux données aux méthodes d'attention parcimonieuse existantes. En passant d'un hachage binaire et dur à une notation probabiliste et souple, SOCKET réalise un compromis supérieur entre précision et efficacité.

Les auteurs soulignent que leur approche :

• Élimine le besoin de prétraitement dépendant des données, réduisant ainsi le TTFT et améliorant la flexibilité de déploiement.
• Fournit une stabilité théorique, garantissant que le classement des tokens reste cohérent même avec un nombre limité de tables de hachage, une faiblesse du LSH traditionnel.
• Offre des gains pratiques au niveau du système, prouvant que les améliorations algorithmiques dans les noyaux de notation peuvent se traduire directement par des améliorations significatives du débit sur les architectures GPU modernes sans sacrifier les performances du modèle.

Ce travail positionne le LSH souple non pas simplement comme un outil de récupération, mais comme un noyau de notation robuste capable de remplacer l'attention dense lors de la phase de décodage, permettant une mise à l'échelle efficace des LLM vers des contextes de plusieurs millions de tokens.