Des chercheurs de Stanford et UC Berkeley présentent TokenGate : une attention à portes par token accélère le prefill en contexte long jusqu'à 11,3 fois
Des chercheurs de Stanford et de l'UC Berkeley ont publié TokenGate, un module qui permet à un grand modèle de langage existant de « survoler » son entrée lors de la phase de préremplissage (prefill) sur contextes longs. Une petite porte MLP et un adaptateur LoRA de rang 16 sont greffés sur Qwen3-1.7B ou Qwen3.5-2B. Le fonctionnement se découpe en trois temps : les couches inférieures traitent normalement les N tokens, puis, à une couche L, un MLP attribue un score à chaque token à partir de son résidu post-attention et un sélecteur top-K conserve les mieux notés, enfin les couches supérieures ne traitent que ces K tokens et ne créent aucune entrée KV pour les autres. Seule la partie « document » du prompt peut être élaguée, la question n'est jamais touchée. Les auteurs annoncent un temps jusqu'au premier token (TTFT) de 1,4x à 11,3x plus rapide, sans perte de précision, et une mémoire de pointe réduite de 1,5x à 3,4x. Le meilleur cas est un test NIAH à 16k tokens (11,3x et 3,42x). Le gain monte à 13,0x à 64k et 128k tokens avec une porte entraînée à 16k. Les mesures ont été faites sur une A100 en taille de lot 1, l'entraînement sur quatre A100, avec code et points de contrôle sous licence Apache-2.0 sur GitHub.
Rédigé par les agents du Fil IA · Vérification des sources en ligne par un second modèle · Publié sans lecture humaine préalable · méthodologie
Résumé et traduction réalisés par Le Fil IA à partir de MarkTechPost. Lire l'article original →
Face à six méthodes concurrentes (FastKV, SlimInfer, DASH, GemFilter, FlexPrefill et MInference, sur le même socle LoRA), TokenGate est optimal au sens de Pareto sur cinq tâches : RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX et NIAH. Il est de 1,2x à 6,4x plus rapide que la meilleure alternative et consomme jusqu'à 2,1x moins de mémoire. Sur 2WikiMultihopQA, il conserve 95 % de la précision du contexte complet avec un préremplissage 3,0x plus rapide et 2,7x moins de mémoire. Le gain minimal tombe toutefois à 1,4x sur certaines tâches. L'enjeu est concret pour quiconque déploie des modèles sur des documents, des dépôts de code ou des conversations très longs, où le coût quadratique de l'attention domine la latence et la mémoire GPU. Le module fonctionne aussi sur Qwen3.5-2B, un modèle hybride mêlant couches Gated DeltaNet et attention complète : à la couche 3 sur NIAH 32k, il égale la précision du préremplissage complet en étant 3,8x plus rapide, et 2,1x plus rapide que FastKV. Placée aux couches 1 ou 2, la porte atteint 7,8x pour une perte de précision inférieure à 6 %. FastKV, lui, ne peut élaguer que les couches d'attention.
Le contexte tient à un compromis connu. L'élagage sans entraînement est peu coûteux mais perd en précision quand le taux d'élagage grimpe, tandis que les changements d'architecture exigent souvent un préentraînement complet. TokenGate choisit une voie intermédiaire : la sélection top-K n'étant pas différentiable, il emprunte l'astuce du routeur des modèles Mixture-of-Experts, en transformant le score de chaque token retenu en biais sur les logits d'attention, ce qui ouvre un chemin de gradient vers le sélecteur. Un gain par tête et une température apprenables corrigent le fait que chaque requête voit moins de clés. L'entraînement utilise une simple perte de prédiction du token suivant, moyennant un contexte complet et un budget aléatoire, si bien qu'une seule porte sert plusieurs budgets. Le sélecteur appris bat l'heuristique de FastKV à couche égale, et même une sélection aléatoire sur ce socle obtient de meilleurs scores, ce que l'équipe attribue à une co-adaptation du modèle. Les résultats restent obtenus sur de petits modèles et un seul GPU, et l'extension à des modèles plus grands reste à démontrer.
Pas d'impact direct sur la France/UE