Technische Zusammenfassung: TokenButler

Problemstellung

Large Language Models (LLMs) sehen sich bei Kontextlängen von 64k–1M Token einem erheblichen Engpass bei Speicher und Rechenleistung gegenüber. Der Key-Value (KV) Cache, der die Token-Historie für eine effiziente Dekodierung speichert, wächst linear mit der Kontextlänge und verbraucht erhebliche GPU-Speicherbandbreite und -bandbreite.

Bestehende Lösungen zur Milderung dieses Engpasses lassen sich im Allgemeinen in drei Kategorien einteilen, die jeweils spezifische Einschränkungen aufweisen:

1. Statische Strategien (z. B. StreamingLLM): Verwenden feste Regeln (z. B. gleitende Fenster), um den Cache zu begrenzen. Diese sind nicht abfragebewusst und können sich nicht an dynamische Eingabebedürfnisse anpassen.
2. Adaptive Verwerfung (z. B. H2O, SnapKV): Werfen Token dauerhaft ab, die basierend auf Metriken wie Aufmerksamkeitsscores als unwichtig erachtet werden. Dies birgt das Risiko, kritische Token zu verlieren, die später in koreferenziellen Kontexten (wo eine früh erwähnte Entität viel später referenziert wird) relevant werden könnten, da die Wichtigkeit oft nicht monoton ist.
3. Adaptive dynamische Strategien (z. B. Quest): Bewahren den gesamten Cache, holen aber Teilmengen von Token (Seiten) zur Dekodierungszeit nach. Obwohl sie eine dauerhafte Verwerfung vermeiden, scheitern sie aufgrund ihrer Abhängigkeit von grobgranularer Seiten-zentrierter Abrufung bei dichten, kontextreichen Aufgaben, bei denen kritische Informationen über Seitengrenzen hinweg verteilt sein können oder eine feingranulare Auswahl erfordern.

Darüber hinaus verlassen sich viele bestehende Methoden auf ungenaue Proxy-Metriken für die Token-Wichtigkeit (z. B. Aktualität oder aggregierte Aufmerksamkeitsscores) oder verursachen einen hohen Rechenoverhead bei der Berechnung vollständiger Aufmerksamkeitsscores für jeden Token.

Methodik: TokenButler

Die Autoren schlagen TokenButler vor, einen hochgranularen, abfragebewussten Prädiktor, der darauf ausgelegt ist, kritische Token zu identifizieren, ohne sie dauerhaft zu verwerfen oder auf grobe Seiteneinteilung zurückzugreifen.

Kernarchitektur

TokenButler arbeitet, indem er niedrigdimensionale Wichtigkeitsabfragen bei festen Tiefenschritten (Produzenten-Schichten) vorhersagt und diese mit gelernten Projektionen der echten KV-Cache-Schlüssel in nachfolgenden Schichten (Verbraucher-Schichten) kombiniert.

• Abfragevorhersage: Ein leichtgewichtiges Multi-Layer Perceptron (MLP) nimmt verborgene Zustände aus einer „Produzenten"-Schicht entgegen und generiert $G$ slotspezifische Abfragevektoren pro Aufmerksamkeitskopf. Diese Vektoren sind niedrigdimensional ($d' \ll E$).
• Schlüssel-Cache-Projektion: Für jede Verbraucher-Schicht werden die tatsächlichen hochdimensionalen KV-Cache-Schlüssel unter Verwendung einer gelernten Matrix in denselben niedrigdimensionalen Raum projiziert.
• Bewertung: Die vorhergesagten Abfragen werden mit den projizierten Schlüsseln skalar multipliziert, um Token-Wichtigkeitsscores kostengünstig zu generieren und eine dynamische Token-Auswahl unter einem festen Budget zu ermöglichen.

Training

• Distillation: TokenButler wird durch Distillation der maskierten kausalen Aufmerksamkeitsverteilungen eines eingefrorenen Basis-LLMs trainiert. Das Ziel ist die Minimierung des Cross-Entropy-Verlusts zwischen den Attention-Logits des Lehrers und den vorhergesagten Logits des Schülers (TokenButler).
• Effizienz: Um lange Kontexte während des Trainings zu handhaben, wird der Verlust auf einer unterprobenen Menge von Abfragepositionen berechnet, wobei der Fokus auf dem späten Kontextbereich und dem finalen Token liegt, wodurch die Komplexität von $O(L^2)$ auf $O(RL)$ reduziert wird.
• Overhead: Der Prädiktor fügt dem Basismodell weniger als 1% Parameter-Overhead hinzu.

Inferenzstrategie

TokenButler partitioniert den KV-Cache in drei Puffer, um den Speicherzugriff zu optimieren:

1. Sink-Puffer: Die ersten $S$ Token (immer behalten).
2. Lokaler Fenster-Puffer: Ein zirkulärer Puffer der neuesten $N$ Token (immer behalten).
3. Wichtiger Puffer: Wird von TokenButler dynamisch basierend auf vorhergesagten Scores gefüllt.

Optimierungen:

• Batching: Projektionen von Schlüsseln werden in Batches von $N$ Token durchgeführt, nicht bei jedem Dekodierschritt, wodurch die zeitliche Lokalität genutzt wird.
• Vorhersageintervall & Nachbarn-Abruf: Um die Kosten des Prädiktors zu amortisieren, wird der Prädiktor nicht bei jedem Schritt ausgeführt. Stattdessen läuft er alle $N$ Schritte (Vorhersageintervall). Um das Risiko zu mindern, Token zu verpassen, deren Wichtigkeit sich zwischen den Intervallen verschiebt, setzt das System „Nachbarn-Abruf" ein, der ausgewählte Token erweitert, um ihre räumlichen Nachbarn einzubeziehen, und so die Abdeckung von Multi-Token-Entitäten oder Denkketten sicherstellt.

Hauptbeiträge

1. Leichtgewichtiger Prädiktor: Eine gelernte Metrik für die Token-Wichtigkeit, die feingranulare Schätzungen mit minimalem Parameter-Overhead (<1%) liefert und über Attention-Distillation trainiert wird.
2. Synthetischer Benchmark: Eine neuartige koreferenzielle Abrufaufgabe, die zeigt, dass bestehende Verwerfungs- und seitenbasierte Methoden selbst bei kurzen Kontextlängen (300 Token) aufgrund von dauerhaftem Token-Verlust oder Page-Misses versagen. TokenButler erreicht auf dieser Aufgabe eine nahezu orakelähnliche Genauigkeit.
3. Leistung bei langen Kontexten: Konkurrierende oder überlegene Ergebnisse auf den RULER- und LongBench-Benchmarks (64K Kontext) im Vergleich zu state-of-the-art Sparsitäts- und Komprimierungsmethoden.
4. Effizienzgewinne:
   • Bis zu 1,6× Geschwindigkeitssteigerung auf der GPU gegenüber Dichter Aufmerksamkeit unter Verwendung von Vorhersageintervallen.
   • Bis zu 7,6× Latenzreduktion im Vergleich zu Dichter Aufmerksamkeit mit CPU-Offloading, was eine effiziente Dekodierung auf speicherbegrenzten Systemen ermöglicht.
   • Die Genauigkeit bleibt selbst bei einer 16-fachen Reduzierung der Prädiktor-Aufrufe innerhalb von ≈1,1% der schrittweisen Baseline.

Experimentelle Ergebnisse

• Synthetische Koreferenz-Aufgabe: TokenButler erreichte auf Llama-3.1-8B und Qwen2.5-7B eine Genauigkeit und Abdeckung von ~80–95 % und übertraf damit deutlich verwerfungsbasierte Methoden (H2O, SnapKV), die unter 50 % Genauigkeit fielen, sowie seitenbasierte Methoden (Quest), die unter partiellen Page-Misses litten.
• RULER-Benchmark (64K Kontext): TokenButler entsprach oder übertraf die Leistung von Baselines wie SnapKV, xKV und Quest auf Llama-3.1-8B und Qwen2.5-7B und erzielte den besten Durchschnittswert auf Qwen2.5-7B.
• LongBench: TokenButler zeigte über verschiedene Datensätze hinweg eine bessere Genauigkeit als frühere KV-Sparsitäts- und Komprimierungsmethoden.
• Reasoning-Modelle (AIME24): Auf DeepSeek-R1-Distill-Llama-8B erreichte TokenButler 30,0 % Genauigkeit und näherte sich der dichten Dekodierungs-Baseline (33,3 %) an, wobei es andere Sparsitätsmethoden unter demselben 8K-Token-Budget übertraf.
• Latenz: In CPU-Offloading-Szenarien (256K+ Kontext) reduzierte TokenButler die Latenz von ~3,2 s pro Token (Dicht) auf ~0,6 s (TokenButler).

Bedeutung

Die Arbeit geht davon aus, dass die Token-Wichtigkeit hochdynamisch und abfrageabhängig ist, wodurch statische Heuristiken und grobgranulare Abrufmechanismen für komplexe, langkontextuelle Aufgaben unzureichend werden. TokenButler zeigt, dass ein gelernter, feingranularer Selektionsmechanismus diese Heuristiken effektiv ersetzen kann. Durch die Bewahrung des gesamten KV-Caches bei gleichzeitiger selektiver Zugriffnahme nur auf die relevantesten Token mildert TokenButler den Speicherbandbreiten-Engpass ohne die mit dauerhafter Verwerfung verbundene Qualitätsminderung. Die vorgeschlagenen Strategien für Vorhersageintervalle und Nachbarn-Abruf zeigen zudem, dass der Rechenkosten eines solchen gelernten Prädiktors effektiv amortisiert werden können, was ihn zu einer praktischen Lösung für Echtzeit-Inferenz mit langem Kontext auf begrenzter Hardware macht.