Forscher von Stanford und UC Berkeley haben TokenGate veröffentlicht, ein Modul für das langes Kontext-Prefill. TokenGate lernt die Token-Pruning-Methode, sodass ein bestehendes LLM seine Eingabe effizient abarbeiten kann. Nur die relevanten Tokens werden vollständig verarbeitet. Das Forschungsteam berichtet über 1,4x bis 11,3x schnellere Zeit bis zum ersten Token (TTFT) ohne Verlust der Genauigkeit. Die maximale Prefill-Memorie sinkt um 1,5x bis 3,4x. Der Code und die Checkpoints befinden sich auf GitHub unter Apache-2.0.
TL;DR
- MLP-Gate: Ein kleiner MLP-Gate und LoRA (r=16) wurden zu Qwen3-1.7B oder Qwen3.5-2B hinzugefügt. Getestet wurde in Kontexten bis 128k Tokens.
- Empfohlene Hardware: Für die Zeitmessungen werden 1 A100-Cluster mit Größe 1 verwendet. Für das Training wurden 4 A100-Cluster verwendet.
- Leistung: Pareto-optimale Leistung bei 5 Aufgaben gegenüber 6 Baseline-Szenarien.
- Bester: NIAH 16k mit 11,3-mal schnellerer TTFT und 3,42-mal geringerem Speicherbedarf. Die TTFT-Erholung erreicht 13,0-mal bei 128k.
- Schlechteste: Die Geschwindigkeit steigt auf 1,4x bei einigen Aufgaben.
Was ist TokenGate?
TokenGate ist eine Ersatzkomponente, die eine tragbare Token-Pruning-Funktion zu den vorgebildeten LLMs hinzufügt. Transformers führen jedes Token durch alle Schichten. Die Kosten für die Aufmerksamkeit steigen quadratisch mit der Länge. Pruning ohne Training ist kostengünstig, verliert aber die Genauigkeit bei hohen Pruning-Raten. Architekturenaustausche erfordern oft eine vorgebildete Phase von Grund auf.
Wie hat TokenGate das Pruning von Tokens verstanden?
Der Vorwärtspass teilt sich in 3 Phasen auf:
- Standard-Frontung: Die Schichten unterhalb der Gate-Schicht L verarbeiten alle N-Tokens wie gewohnt.
- Gate-Modul: eine MLP-Architektur berechnet den Score jedes Tokens aus seinem post-Attentionsresiduum in Schicht L. Ein Top-K-Selector hält die am höchsten bewerteten Tokens bei.
- Prune-Front-Behandlung: Die Schichten oberhalb von L verarbeiten nur die K Tokens und erzeugen keine KV-Einträge für die geprünften Elemente.
Nur der Dokumententeil des Promptes kann entfernt werden. Die Frage wird niemals entfernt.
Was ist token-gated Attention?
Die Top-K-Selection ist nicht unterscheidbar, was die End-to-End-Training blockiert. TokenGate nutzt den Trick des Routers „Mixture-of-Experts“. Der Score jedes behaltenen Tokens wird zu einem Bias auf den Attention-Logits über dem Gate. Dieser Bias gibt den Gradienten einen Weg zurück zum Scorer.
Eine lernbare Gewinnrate pro Kopf und eine lernbare Temperatur korrigieren dies. Die Temperatur korrigiert für jede Anfrage, bei der weniger Schlüssel vorhanden sind. Die Trainierung verwendet die einfache Next-Token-Verlustmessung. Jeder Schritt berechnet einen Durchschnitt aus einem vollständigen- Kontext-Verlust und einem zufälligen-Budget-Verlust. Somit funktioniert 1 trainiertes Gate für alle Budgets.
Wie zeigt TokenGate sich bei Benchmarks?
Das Forschungsteam testete Qwen3-1.7B auf RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX und needle-in-a-haystack (NIAH). Die Baselines waren FastKV, SlimInfer, DASH, GemFilter, FlexPrefill und MInference, alle auf demselben LoRA-Backbone.
- Vorheriges vollständiges Vollfüllen: 1,4x bis 11,3x schnellere TTFT und 1,5x bis 3,4x weniger Speicherverbrauch.
- Gegen die nächstbeste Grundlage: 1,2x bis 6,4x schnellere TTFT und bis zu 2,1x weniger Speicher.
- 2WikiMultihopQA: Bei 95% der vollen Kontextgenauigkeit läuft das Vorspannen 3,0x schneller und verbraucht 2,7x weniger Speicher.
- Längere Kontexte: eine Gate, die mit 16k trainiert wurde, erreicht bei 32k 12,8x und bei 64k sowie 128k 13,0x. Die Speicherersparnis bleibt bei etwa 3,43x.
Der ausgebildete Selektor übertrifft die Heuristik von FastKV auf derselben Ebene. Selbst eine zufällige Selektion erzielt höhere Punkte im TokenGate-Backbone. Das Forschungsteam verweist auf die gemeinsame Anpassung des Backbone.
Wird es auf Hybrid-Attention-Modellen funktionieren?
Ja, Qwen3.5-2B kombiniert Gated DeltaNet und vollständige Attention-Schichten. Auf NIAH 32k erreichte die TokenGate in der Schicht 3 die gleiche Genauigkeit wie das vollständige Vollfüllen. Es lief 3,8-mal schneller als das vollständige Vollfüllen und 2,1-mal schneller als FastKV. Die Gating in den Schichten 1 oder 2 erreichte eine Genauigkeit von 7,8-mal, innerhalb von 6 %. FastKV kann nur die Attention-Schichten entfernen.
TokenGate gegenüber FastKV, SlimInfer, GemFilter und DASH
| Funktion | TokenGate | FastKV | SlimInfer | GemFilter | DASH |
|---|---|---|---|---|---|
| Tokenauswahl | Lernte MLP-Gate | Aufmerksamkeit aus den letzten Anfragen | Aufmerksamkeitsblock, mehreckstegig | Letzte-Anfrage beachten bei der Filterschicht | Verborgene-Zustandsstabilität |
| Trainings erforderlich | Ja (Gate + LoRA r=16) | Keine | Keine | Keine | Keine |
| Aufpflanzschicht bei Tokentag’s Qwen3-1.7B-Version | 1 | 13 (Standardwert) | 8, 17, 25 | 12 | 11 |
| Prunes in den Gated DeltaNet-Layern | Ja | Nein (nur Schichten der Aufmerksamkeit) | Nicht offenlegt | Nicht offenlegt | Nicht offenlegt |
| Bester selbst gemeldeter Vollständigkeitsgewinn (eigene Einstellung) | 11,3x TTFT, 3,4x Speicher | 1,82x Vorschub | 2,53x TTFT (LLaMA3.1-8B, RTX 4090) | 2,4x schneller, 30% weniger GPU-Memory im Vergleich zu SOTA | Nicht offenlegt |
| Code | GitHub | GitHub | GitHub | GitHub | Veröffentlicht |
Die selbst gemelten Vorteile werden mit unterschiedlichen Methoden erfasst, daher sind sie nicht direkt vergleichbar. Die „Pruning-Layer“ stammen aus Anhang C.2 von TokenGate.
Wie wird TokenGate implementiert?
Der GitHub-Repo erfordert Python 3.10+ und CUDA PyTorch 2.9+. Ein optionaler FlashAttention-2 Key-bias-Kernel beschleunigt die Verarbeitung von dense Qwen3. Ohne diesen Kernel erfolgt die Inferenz auf FlexAttention. Der Preis ist 1,0x bis 1,1x höher als bei einer einfachen Anfrage.
Ein {state} Platzhalter markiert den zu entfernenden Dokumenten. generate nimmt ein max_kept_prefill_tokens Budget in Anspruch. Checkpoints umfassen alle 5 Aufgaben für Qwen3-1.7B sowie eine Qwen3.5-2B NIAH-Gate. Beispiel: liana313/tokengate-niah-qwen3-1.7b.
Das Training folgt dem TRL/PEFT-Rezept. Jeder Schritt ist 1,3-mal langsamer als ein einfacher LoRA-Prozess. Die Ausführung kam bei 4 von 5 Aufgaben schneller zum Erfolg. Die durchschnittliche Trainingszeit betrug 1,8 statt 2,2 A100 Stunden.
Wichtigste Erkenntnisse
- TokenGate fügt bestehenden LLMs durch einen kleinen Gate und LoRA die erlerte Token-Pruning-Funktion hinzu.
- Die Ausführung ist 1,4x bis 11,3x schneller und verbraucht 1,5x bis 3,4x weniger maximale Speicherkapazität.
- Die durch Tokens gesteuerte Aufmerksamkeit macht die Auswahl der Top-K-Menge trainiert verfügbar, im Stil eines MoE-Router.
- Es übertrifft die Nächstbeste Pruning-Basisline um bis zu 6,4x bei TTFT.
- Es funktioniert mit dem Hybriden Qwen3.5-2B, bei dem FastKV die DeltaNet-Layer nicht entfernen kann.
Schauen Sie sich das Paper, den GitHub-Repo und die Hugging Face-Paper-Seite an. Folgen Sie uns auch gerne auf Twitter, und vergessen Sie nicht, sich bei unserem 150k+ML SubReddit anzumelden und unseren Newsletter zu abonnieren. Warte! Sind Sie auf Telegram? Jetzt können Sie uns auch auf Telegram beitreten.
Müssen Sie eine Partnerschaft mit uns eingehen, um Ihre GitHub-Repo, Hugging Face-Seite, Produktveröffentlichung oder Webinar zu promoten? Mit uns in Kontakt treten
Die Nachricht von Forschern von Stanford und UC Berkeley: TokenGate – Token-Gated Attention macht das Vollfüllen von langen Kontexten bis zu 11,3-mal schneller erschien zunächst auf MarkTechPost.