Le routage partagé entre couches accélère l'attention sparse de 7x sans perte de qualité

L'attention reste le goulot d'étranglement. Chaque token supplémentaire de contexte dans un transformateur signifie davantage de clés et de valeurs à stocker et à parcourir, et le décodage d'une longue chaîne de pensée peut ralentir même sur du matériel haut de gamme. L'attention sparse tente de résoudre ce problème en ignorant les parties non pertinentes du cache, mais cette correction coûte souvent en précision (sparsity par blocs) ou en vitesse (sparsity par tokens, où la sélection des bons tokens à chaque couche est un goulot d'étranglement en soi).

Un nouvel article, CLSA: You Only Index Once, soutient que les deux problèmes ont une cause racine unique : chaque couche exécute son propre routage de manière indépendante. Les auteurs proposent de partager l'indice de routage entre les couches dans une architecture partageant le KV-cache, de sorte que la coûteuse sélection top-k s'exécute une seule fois et que le résultat soit réutilisé. Les expériences rapportent un décodage jusqu'à 7,6x plus rapide et un débit global amélioré de 17,1x à 128K tokens, avec une précision équivalente à l'attention totale sur les benchmarks standards.

C'est un cas rare où la correction architecturale améliore à la fois l'efficacité et la qualité du modèle simultanément, plutôt que de sacrifier l'une pour l'autre.

Comment CLSA fonctionne

La méthode s'appuie sur des architectures d'attention hybrides qui partagent les caches KV entre les couches croisées du décodeur, comme la famille YOCO. Dans ces modèles, le cache KV est déjà réutilisé, mais chaque couche croisée du décodeur exécute encore son propre routage top-k sur l'ensemble du cache. CLSA ajoute une couche indexeur partagée : elle calcule les scores d'attention une fois, sélectionne les tokens top-k, et transmet ce même indice à chaque couche suivante.

Ce passage unique remplace N passages de routage, où N est le nombre de couches croisées du décodeur. Le routage devient ainsi un coût fixe plutôt que linéaire en fonction du nombre de couches, ce qui devient d'autant plus important que les modèles s'approfondissent et que les contextes s'allongent.

L'astuce clé est de conserver la sparsité au niveau des tokens, soit la sélection fine qui garantit une haute précision, tout en amortissant les frais généraux qui la rendent normalement lente. Les méthodes de sparsité par blocs ignorent des pages mémoire entières pour gagner en vitesse mais manquent des tokens pertinents à l'intérieur des blocs ignorés. CLSA évite cela en conservant une granularité par token et en la payant exactement une fois.

D'où vient l'accélération

L'article mesure trois étapes d'inférence : le pré-remplissage (traitement de la requête), le stockage du cache KV (empreinte mémoire des clés et valeurs mises en cache), et le décodage autorégressif (génération de tokens un par un). Toutes trois s'améliorent simultanément :

Ces chiffres proviennent d'un modèle avec une fenêtre de contexte de 128K tokens. La seule accélération du décodage, 7,6x, signifie qu'une génération qui prend 8 secondes avec une attention totale se termine en un peu plus d'une seconde avec CLSA. Les gains de débit sont encore plus importants car le modèle peut traiter plus de requêtes par lot avec une empreinte mémoire réduite.

La précision reste stable

Sur le benchmark long contexte RULER (moyenné sur cinq modèles de requête et plusieurs longueurs), CLSA correspond à l'attention totale à moins de 0,5 point près. Sur les tâches à contexte court, MMLU, ARC-Challenge, HellaSwag, la différence est inférieure à 0,3 point. C'est essentiellement du bruit. Sur GSM8K (raisonnement mathématique), l'écart s'élargit à environ 1,2 point, ce que les auteurs attribuent au fait que de petites sélections de tokens peuvent manquer des nombres critiques dans les chaînes arithmétiques.

Néanmoins, le compromis est inhabituel : une méthode qui offre un gain de vitesse de 7x tout en perdant moins de 1,5 point sur n'importe quel benchmark est rare parmi les propositions d'attention sparse. Les méthodes de sparsité par blocs perdent généralement 3 à 5 points sur les mêmes tâches pour une vitesse comparable.

La stabilité de la précision vient du fait que l'indice de routage est calculé une fois en utilisant les scores d'attention de la couche indexeur, qui voit le même état de requête et de cache que toutes les autres couches. L'ensemble top-k pour une couche est un proxy solide pour l'ensemble top-k de toutes les couches dans une architecture à cache partagé.

Le coût : limites de compatibilité

CLSA n'est pas un remplacement direct pour n'importe quel transformateur. Il nécessite un backbone partageant le cache KV (YOCO, ou des modèles construits avec des couches croisées de décodeur) pour fonctionner. Un transformateur standard décodeur uniquement qui ne partage pas les caches entre les couches ne bénéficierait pas du routage partagé, car le cache KV de chaque couche est différent et l'indice pointerait vers des entrées obsolètes.

Les praticiens utilisant le décodage spéculatif pour le raisonnement sur longs contextes pourraient trouver CLSA complémentaire : le modèle brouillon pourrait utiliser le même routage partagé pour la rapidité, tandis que le modèle cible exécute une attention totale pour la vérification. L'article n'explore pas cette combinaison, mais les architectures sont compatibles.

L'autre limitation est le nombre fixe de sélections top-k. L'article utilise k=512 tokens sur un cache de 128K, soit environ 0,4% de densité. Pour les tâches nécessitant une attention très dense, par exemple l'analyse de chaque token d'un contrat financier pour un contrôle de conformité, une sparsité aussi agressive pourrait manquer des détails même si l'indice de routage est précis. Les auteurs suggèrent que la couche indexeur peut ajuster k dynamiquement, mais cela est laissé pour des travaux futurs.

Ce que cela signifie pour l'inférence sur longs contextes

L'affirmation la plus forte de l'article est que le compromis efficacité-précision n'est pas inévitable ; c'est un artefact du routage indépendant par couche. Un seul passage de routage partagé résout les deux aspects : la précision reste élevée car la sélection est au niveau des tokens, et la vitesse s'améliore car le coût du routage est amorti entre les couches.

À 128K de contexte, le gain de débit de 17x signifie qu'un seul GPU peut servir environ 17x plus de requêtes simultanées sur de longs contextes qu'avec une attention totale. Pour les fournisseurs d'inférence cloud, cela réduit directement le coût par token. Pour les modèles à poids ouverts sur appareil, cela réduit l'écart de mémoire et de latence entre l'exécution de requêtes courtes et longues.

Il reste des questions ouvertes : comment CLSA se comporte au-delà de 128K (l'article teste jusqu'à ce point) ; si la sélection dynamique de k peut combler l'écart de <1,5 point sur les tâches de raisonnement arithmétique ; et comment la méthode se scale sur des modèles de 70B+ paramètres où l'indice de routage lui-même devient une grande structure de données. Mais en tant que solution au goulot d'étranglement spécifique qui rend l'inférence sur longs contextes coûteuse, elle aborde le bon goulot d'étranglement de manière principlée.