Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

1. Problemstellung

Large Language Models (LLMs) stoßen bei der Inferenz mit langen Kontexten aufgrund der quadratischen Komplexität ($O(L^2)$) des Aufmerksamkeitsmechanismus während der Prefill-Phase auf eine fundamentale Engstelle. Während Hardware-Optimierungen wie FlashAttention den Overhead bei der Speicher-I/O reduzieren, ändern sie nicht die algorithmische Komplexität.

Bestehende Beschleunigungsmethoden leiden unter zwei Hauptbeschränkungen:

• Strukturierte Sparsität: Methoden wie Minference oder FlexPrefill wenden Sparsität auf Blockebene an. Dies behält oft irrelevante Token bei, wenn sie mit wichtigen Token im selben Block gruppiert sind, was zu suboptimaler Sparsität führt.
• Token-Eviction: Methoden wie FastKV oder GemFilter verwerfen Token in frühen Schichten dauerhaft. Dies geht von einer statischen Token-Wichtigkeit aus und ignoriert die dynamische Natur der Token-Relevanz über verschiedene Schichten und Aufmerksamkeitsköpfe hinweg. Einmal verworfene Token können nicht erneut berücksichtigt werden, was zu einer Verschlechterung der Genauigkeit führt.

2. Methodik: Token Sparse Attention

Die Autoren schlagen Token Sparse Attention vor, einen leichten, dynamischen und reversiblen Sparsifizierungsmechanismus auf Token-Ebene. Im Gegensatz zu Eviction-basierten Methoden entfernt er Token nicht dauerhaft, sondern komprimiert den Berechnungsraum und stellt die Sequenz anschließend wieder her.

Kernarchitektur

Die Methode arbeitet pro Aufmerksamkeitskopf in zwei Stufen:

1. Kompression (QKV-Auswahl):
   • Für jeden Aufmerksamkeitskopf $h$ wird eine Teilmenge von Token-Indizes $S_h$ basierend auf Wichtigkeitsscores ausgewählt.
   • Die Query-($Q$), Key-($K$) und Value-($V$)-Tensoren werden gesammelt, um komprimierte Tensoren $\hat{Q}, \hat{K}, \hat{V}$ mit einer reduzierten Sequenzlänge $L' \ll L$ zu bilden.
   • Schlüsselmerkmal: Die Auswahl erfolgt unabhängig pro Kopf, was es ermöglicht, dass verschiedene Köpfe sich auf unterschiedliche Teilmengen von Token konzentrieren, wodurch die Heterogenität zwischen den Köpfen adressiert wird.
   • Die komprimierten Tensoren bleiben dicht und zusammenhängend, was die Verwendung optimierter Kernel wie FlashAttention ermöglicht.
2. Dekompression (Wiederherstellung der Ausgabe):
   • Die Aufmerksamkeitsausgabe $\hat{O}$ (Form $L' \times d$) wird unter Verwendung der kopfspezifischen Indizes in einen mit Nullen initialisierten Tensor der ursprünglichen Sequenzlänge $L \times d$ gestreut.
   • Schlüsselmerkmal: Die nicht ausgewählten Positionen werden auf Null gesetzt, aber die Residualverbindung aus der vorherigen Schicht bewahrt die Information der übersprungenen Token. Dies ermöglicht dem Modell, die Token-Wichtigkeit in nachfolgenden Schichten erneut zu „besuchen" und neu zu bewerten.

Dynamische Token-Abdeckung

Um zu bestimmen, welche Token behalten und wie viele beibehalten werden sollen, führt das Papier Dynamic Token Coverage ein:

• Bewertung: Eine leichte Proxy-Aufmerksamkeitskarte wird unter Verwendung aktueller Queries gegen alle Keys berechnet. Token-Wichtigkeits-Scores werden durch vertikales Summieren der Aufmerksamkeitsgewichte abgeleitet.
• Budgetzuweisung: Anstatt die Top-$k$-Token zu behalten, werden die Token nach aufsteigender Wichtigkeit sortiert und das „Schwanzende" (die unwichtigsten Token) so lange beschneidet, bis die kumulative Masse der beschneideten Token einen Abdeckungsschwellenwert $\tau$ überschreitet. Dies behandelt das Beschneiden als strukturelle Regularisierung gegen Aufmerksamkeitsrauschen.
• Kopfspezifische Auswahl: Nach Bestimmung des globalen Budgets wählt jeder Kopf unabhängig seine Top-$k$-Token basierend auf seinen spezifischen Scores aus.

Selektion sparer Schichten

Die Autoren beobachten, dass die Anwendung von Sparsität auf alle Schichten die Leistung verschlechtert. Sie führen Inter-Layer Representation Drift ($R_\ell$) ein, um die Stabilität der Token-Repräsentationen zwischen den Schichten zu messen.

• Schichten mit niedrigem Drift (stabile Repräsentationen) werden für die Sparsifizierung ausgewählt.
• Schichten mit hohem Drift (wo sich Repräsentationen signifikant ändern) bleiben dicht, um das Modellverhalten zu erhalten.
• Diese Auswahl ist ein einmaliger Vorverarbeitungsschritt pro Modell.

3. Hauptbeiträge

1. Reversible Sparsifizierung auf Token-Ebene: Ein neuartiges „Komprimieren und Dekomprimieren"-Design, das die Rechenkosten senkt, ohne Token dauerhaft zu verwerfen, und eine dynamische Neubewertung der Token-Wichtigkeit über die Schichten hinweg ermöglicht.
2. Heterogenität zwischen Köpfen: Die Methode unterstützt eine unabhängige Token-Auswahl für jeden Aufmerksamkeitskopf und erfasst die Tatsache, dass verschiedene Köpfe unterschiedliche semantische Merkmale priorisieren.
3. Nahtlose Kompatibilität: Der Ansatz ist vollständig kompatibel mit bestehenden dichten (FlashAttention) und sparsamen (Minference, FlexPrefill) Kerneln und fungiert als ergänzender Beschleuniger statt als Ersatz.
4. Dynamisches Budgeting: Ein datengesteuerter Ansatz zur Bestimmung der Sparsitätsgrade basierend auf Verteilungen der Token-Wichtigkeit und Schicht-stabilität anstatt fester Verhältnisse.

4. Experimentelle Ergebnisse

Experimente wurden an LLaMA-3.1-8B-Instruct und Mistral-Nemo-12B-Instruct mit Benchmarks wie RULER, InfiniteBench und Needle-in-a-Haystack durchgeführt.

• Trade-off zwischen Genauigkeit und Latenz:
  • Token Sparse Attention verbessert den Trade-off zwischen Geschwindigkeit und Genauigkeit konsistent.
  • Bei 128K Kontext wird eine bis zu $\times 3,23$-fache Beschleunigung der Aufmerksamkeit im Vergleich zu FlashAttention erreicht, bei weniger als 1% Genauigkeitsverlust.
  • In Kombination mit FlexPrefill verbessert sich die Beschleunigung von $\times 2,44$ auf $\times 2,76$, bei Beibehaltung der identischen durchschnittlichen Genauigkeit (87,27%).
• Vergleich mit Eviction-Methoden:
  • Deutliche Überlegenheit gegenüber Token-Eviction-Methoden (PyramidInfer, FastKV, GemFilter). Beispielsweise erreicht Token Sparse Attention bei 128K eine Genauigkeit von 86,84% gegenüber 78,49% für PyramidInfer bei ähnlicher Beschleunigung.
• Skalierbarkeit:
  • Die Beschleunigungsgewinne nehmen mit der Kontextlänge zu (z. B. minimaler Gewinn bei 4K, beträchtlich bei 128K), da die Aufmerksamkeits-Sparsität mit der Sequenzlänge zunimmt.
• Overhead:
  • Der Overhead für Bewertung und Indexverwaltung ist minimal und macht selbst bei hohen Sparsitätsgraden <11% der gesamten Aufmerksamkeitslatenz aus.

5. Bedeutung

Dieses Papier schließt eine kritische Lücke in der Inferenz von LLMs mit langen Kontexten, indem es von einer irreversiblen Token-Eviction hin zu einer dynamischen, interleaved Sparsifizierung übergeht.

• Theoretische Einsicht: Es zeigt, dass die Token-Wichtigkeit über Schichten und Köpfe hinweg hochdynamisch ist und flexible Auswahlmechanismen erfordert, anstatt statisches Beschneiden.
• Praktische Auswirkung: Es bietet eine „Plug-and-Play"-Optimierung, die auf bestehende state-of-the-art Aufmerksamkeits-Kernel (FlashAttention, Minference usw.) aufgelegt werden kann, um massive Beschleunigungen in Szenarien mit langen Kontexten zu erzielen, ohne Modelle neu zu trainieren oder signifikante Genauigkeit zu opfern.
• Zukunftsrichtung: Es etabliert einen neuen Designpunkt für skalierbare Inferenz und legt nahe, dass reversible Kompression eine überlegene Strategie zur Bewältigung der $O(L^2)$-Engstelle im Vergleich zur permanenten Token-Entfernung darstellt.