Les chercheurs de Stanford et de UC Berkeley ont publié TokenGate, un module pour le préchargement de longue contexte. TokenGate a appris la réduction des tokens, permettant à un LLM existant de passer rapidement par son entrée. Seuls les tokens importants sont traités en détail. L’équipe de recherche indique 1,4x à 11,3x plus de vitesse pour le temps jusqu’au premier token (TTFT), sans perte d’exactitude. La mémoire de préchargement maximale diminue de 1,5x à 3,4x. Le code et les checkpoints se trouvent sur GitHub sous licence Apache-2.0.
TL;DR
- Portail MLP : Un petit portail MLP et un LoRA (r=16) ajoutés à Qwen3-1.7B ou Qwen3.5-2B. Testé sur des contextes allant jusqu’à 128k tokens.
- Compatible avec : Les temps de papier utilisent 1 A100 à taille de lot 1. L’entraînement a utilisé 4 A100.
- Performance : Optimum selon le principe de Pareto pour 5 tâches par rapport à 6 lignes de base.
- Meilleur : NIAH 16k avec une vitesse de traitement 11,3 fois supérieure et une consommation de mémoire 3,42 fois inférieure. La gain en vitesse atteint 13,0 fois à 128k.
- Le pire : Le gain de performance tombe à 1,4x pour certaines tâches.
Qu'est-ce que TokenGate ?
TokenGate est un composant intégré qui ajoute une prétraînement de tokens traçable aux LLMs pré-trainés. Les Transformers font passer chaque token à travers toutes les couches. Le coût d’attention augmente quadratiquement avec la longueur. La prétraînement sans entraînement est économique, mais perd en précision à des taux de prétraînement élevés. Les changements architecturaux nécessitent souvent un pré-traînement à partir de zéro.
Comment TokenGate a-t-il appris le fonctionnement de la réduction des tokens ?
Le passage en avant se divise en 3 phases :
- Transfert standard : les couches en dessous de la couche de portail L traitent tous les N tokens comme d’habitude.
- Module de portail : un MLP calcule les scores de chaque token à partir de son résidu post-attention au niveau L. Un sélecteur top-K garde les tokens ayant les meilleurs scores.
- Avant le prélèvement : les couches au-dessus de L ne traitent que ces K tokens et ne créent jamais d’enregistrements KV pour les éléments prélévés.
Seule la partie document du prompt peut être supprimée. La question ne est jamais supprimée.
Qu'est-ce que l'attention gérée par des tokens ?
La sélection des top-K est non différenciable, ce qui bloque l’entraînement en continu. TokenGate utilise la technique du routage Mixture-of-Experts. Le score de chaque token conservé devient un biais sur les logits d’attention au-dessus du seuil. Ce biais permet aux gradients de revenir vers le scorer.
Une gain apprenable par tête et une température apprenable améliorent cela. La température corrige pour chaque requête en voyant moins de clés. L’entraînement utilise une perte simple du prochain mot. Chaque étape moylie une perte de contexte complet et une perte de budget aléatoire. Ainsi, 1 porte entraînée fonctionne avec tous les budgets.
Comment TokenGate se performance sur les benchmarks ?
L’équipe de recherche a testé Qwen3-1.7B sur RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX et needle-in-a-haystack (NIAH). Les benchmarks utilisés étaient FastKV, SlimInfer, DASH, GemFilter, FlexPrefill et MInference, basés sur le même backbone LoRA.
- Versus plein remplissage : 1,4 à 11,3 fois plus rapide en TTFT et 1,5 à 3,4 fois moins de mémoire.
- Contre le niveau de base suivant : 1,2x à 6,4x plus rapide en TTFT et jusqu’à 2,1x moins de mémoire.
- 2WikiMultihopQA : À 95 % d’exactitude totale, le préchargement se fait 3,0 fois plus rapidement et consomme 2,7 fois moins de mémoire.
- Contextes plus longs : une porte entraînée à 16 k utilisant 32 k atteint 12,8 fois, à 64 k atteint 13,0 fois, et à 128 k atteint 13,0 fois. Le gain en mémoire reste près de 3,43 fois.
Le sélecteur expert bat l’heuristique de FastKV au même niveau. Même la sélection aléatoire-K obtient des scores plus élevés sur le backbone de TokenGate. L’équipe de recherche attribue cela à la co-adaptation du backbone.
Est-ce que cela fonctionne avec les modèles d'attention hybrides ?
Oui, Qwen3.5-2B combine des couches Gated DeltaNet et des couches d’attention complètes. Sur NIAH 32k, le TokenGate au niveau de la couche 3 a atteint l’exacte précision du pré-fill complet. Il fonctionne 3,8 fois plus rapidement que le pré-fill complet et 2,1 fois plus rapidement que FastKV. La gestion des couches au niveau 1 ou 2 atteint 7,8 fois, avec une précision inférieure de 6 %. FastKV ne peut que supprimer les couches d’attention.
TokenGate contre FastKV, SlimInfer, GemFilter et DASH
| Caractéristique | TokenGate | FastKV | SlimInfer | GemFilter | DASH |
|---|---|---|---|---|---|
| Sélection des tokens | Appris la porte MLP | Attention provenant des dernières requêtes | Alerte, plusieurs étapes | Dernière requête attention au niveau du filtre | Stabilité de l’état caché |
| Formation nécessaire | Oui (gate + LoRA r=16) | Aucun | Aucun | Aucun | Aucun |
| Étage de taille dans la mise en œuvre de Qwen3-1.7B de TokenGate | 1 | 13 (par défaut) | 8, 17, 25 | 12 | 11 |
| Prunes à Gated DeltaNet couches | Oui | Non (seulement les couches d’attention) | Non divulgué | Non divulgué | Non divulgué |
| Meilleur gain de remplissage auto-reporté (configuration personnelle) | 11,3x TTFT, 3,4x mémoire | 1.82x préchargement | 2.53x TTFT (LLaMA3.1-8B, RTX 4090) | 2,4x de vitesse, 30 % moins de mémoire GPU par rapport au SOTA | Non divulgué |
| Code | GitHub | GitHub | GitHub | GitHub | Annoncé |
Les gains rapportés par l’utilisateur utilisent différentes configurations, donc ils ne sont pas directement comparables. Les couches de taillement proviennent de l’Appendice C.2 de TokenGate.
Comment déployer TokenGate ?
Le repo GitHub nécessite Python 3.10+ et CUDA PyTorch 2.9+. Un noyau FlashAttention-2 key-bias optionnel accélère l’inference de Qwen3 dense. Il coûte de 1,0x à 1,1x par appel simple. Sans lui, l’inference se fait via FlexAttention.
Un espace réservé {state} indique le document qui peut être supprimé. generate utilise un budget de max_kept_prefill_tokens. Les points de contrôle couvrent toutes les 5 tâches sur Qwen3-1.7B, ainsi que une porte NIAH pour Qwen3.5-2B. Exemple : liana313/tokengate-niah-qwen3-1.7b.
L’entraînement correspond à une recette TRL/PEFT. Chaque étape est 1,3 fois plus lente que le LoRA classique. Les exercices se sont convergencenés plus rapidement sur 4 des 5 tâches. Le temps moyen d’entraînement était de 1,8 contre 2,2 A100 heures.
Résumés clés
- TokenGate ajoute une prise de raccourcissement des tokens appris aux LLM existants par l’intermédiaire d’un petit portail et de LoRA.
- Les performances sont 1,4 fois à 11,3 fois plus rapides, avec une consommation de mémoire maximale 1,5 fois à 3,4 fois inférieure.
- L’attention filtrée par des tokens rend la sélection des top-K trainable, dans le style de MoE-router.
- Cela est jusqu’à 6,4 fois supérieur au niveau de prise de raccourcis suivant le mieux.
- Cela fonctionne avec le Qwen3.5-2B hybride, où FastKV ne peut pas supprimer les couches de DeltaNet.
Veuillez consulter le paper, le repo GitHub et la page de paper sur Hugging Face. Veuillez également nous suivre sur Twitter, et n’oubliez pas de rejoindre notre subreddit 150k+ML et de vous abonner à notre bulletin email. Attends ! êtes-vous sur Telegram ? Maintenant, vous pouvez également nous rejoindre sur Telegram.
Voulez-vous collaborer avec nous pour promouvoir votre repository GitHub, votre page Hugging Face, la sortie d'un produit ou un webinaire, etc. ? Connectez-vous à nous
La publication découverte par des chercheurs de Stanford et de UC Berkeley du TokenGate : l’attention gérée par les tokens rend le remplissage du contexte long jusqu’à 11,3 fois plus rapide a été publiée en premier sur MarkTechPost.