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

1. Probleemstelling

Grote Taalmodellen (LLM's) ondervinden een fundamentele bottleneck bij inferentie met lange contexten vanwege de kwadratische complexiteit ($O(L^2)$) van het attentiemechanisme tijdens de prefill-fase. Hoewel hardware-optimalisaties zoals FlashAttention de overhead van geheugen-I/O verminderen, veranderen ze de algoritmische complexiteit niet.

Bestaande versnellingsmethoden lijden onder twee hoofdbeperkingen:

• Gestructureerde Sparsiteit: Methoden zoals Minference of FlexPrefill passen sparsiteit op blokniveau toe. Dit behoudt vaak irrelevante tokens als deze in hetzelfde blok zijn gegroepeerd met belangrijke tokens, wat leidt tot suboptimale sparsiteit.
• Tokenverwijdering (Eviction): Methoden zoals FastKV of GemFilter gooien tokens in vroege lagen permanent weg. Dit gaat ervan uit dat tokenbelang statisch is, en negeert het dynamische karakter van tokenrelevantie over verschillende lagen en attentiekoppen heen. Eenmaal verwijderd, kan een token niet opnieuw worden overwogen, wat leidt tot een afname in nauwkeurigheid.

2. Methodologie: Token Sparse Attention

De auteurs stellen Token Sparse Attention voor, een lichtgewicht, dynamisch en reversibel mechanisme voor sparsificatie op token-niveau. In tegenstelling tot methoden op basis van verwijdering, verwijdert het tokens niet permanent, maar comprimeert het de rekenruimte en herstelt het vervolgens de sequentie.

Kernarchitectuur

De methode werkt in twee fasen per attentiekop:

1. Compressie (QKV-selectie):
   • Voor elke attentiekop $h$ wordt een subset van tokenindices $S_h$ geselecteerd op basis van belangrijkheidsscores.
   • De Query ($Q$), Key ($K$) en Value ($V$) tensoren worden verzameld om gecomprimeerde tensoren $\hat{Q}, \hat{K}, \hat{V}$ te vormen met een verkorte sequentielengte $L' \ll L$.
   • Kernkenmerk: Selectie wordt onafhankelijk per kop uitgevoerd, waardoor verschillende koppen zich kunnen richten op verschillende subsets van tokens, wat de heterogeniteit per kop aanpakt.
   • De gecomprimeerde tensoren blijven dicht en aaneengesloten, waardoor het gebruik van geoptimaliseerde kernels zoals FlashAttention mogelijk is.
2. Decompressie (Herstel van output):
   • De attentie-output $\hat{O}$ (vorm $L' \times d$) wordt teruggeplaatst in een met nullen geïnitialiseerde tensor van de oorspronkelijke sequentielengte $L \times d$ met behulp van kop-specifieke indices.
   • Kernkenmerk: De niet-geselecteerde posities worden op nul gezet, maar de residuele verbinding uit de vorige laag behoudt de informatie van de overgeslagen tokens. Dit stelt het model in staat om in volgende lagen tokenbelang opnieuw te "bezoeken" en te evalueren.

Dynamische Tokendekking

Om te bepalen welke tokens te behouden en hoeveel, introduceert het paper Dynamische Tokendekking:

• Scorebepaling: Een lichtgewicht proxy-attentiekarte wordt berekend met recente queries tegen alle keys. Tokenbelangrijke scores worden afgeleid door attentiegewichten verticaal op te tellen.
• Budgettoewijzing: In plaats van de top-$k$ tokens te behouden, sorteert de methode tokens op stijgende belangrijkheid en snoeit de "staart" (minder belangrijke tokens) tot de cumulatieve massa van gesnoeide tokens een dekkingsthorshold $\tau$ overschrijdt. Dit behandelt snoeien als structurele regularisatie tegen attentieruis.
• Selectie per kop: Na het bepalen van het globale budget selecteert elke kop onafhankelijk zijn top-$k$ tokens op basis van zijn specifieke scores.

Selectie van Sparselagen

De auteurs merken op dat het toepassen van sparsiteit op alle lagen de prestaties verslechtert. Ze introduceren Inter-Layer Representation Drift ($R_\ell$) om de stabiliteit van tokenrepresentaties tussen lagen te meten.

• Lagen met lage drift (stabiele representaties) worden geselecteerd voor sparsificatie.
• Lagen met hoge drift (waar representaties significant veranderen) worden dicht gehouden om het modelgedrag te behouden.
• Deze selectie is een eenmalige voorverwerkingsstap per model.

3. Belangrijkste Bijdragen

1. Reversibele Sparsificatie op Token-niveau: Een nieuw "Comprimeren en Decomprimeren"-ontwerp dat de rekenkosten verlaagt zonder tokens permanent te verwijderen, waardoor dynamische herbeoordeling van tokenbelang over lagen mogelijk is.
2. Heterogeniteit per Kop: De methode ondersteunt onafhankelijke tokenselectie voor elke attentiekop, wat het feit benadrukt dat verschillende koppen verschillende semantische kenmerken prioriteren.
3. Naadloze Compatibiliteit: De aanpak is volledig compatibel met bestaande dichte (FlashAttention) en sparsere (Minference, FlexPrefill) kernels, en fungeert als een complementaire versneller in plaats van een vervanging.
4. Dynamisch Budgetteren: Een datagedreven aanpak om sparsiteitsniveaus te bepalen op basis van verdelingen van tokenbelang en stabiliteit per laag, in plaats van vaste ratios.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op LLaMA-3.1-8B-Instruct en Mistral-Nemo-12B-Instruct met behulp van benchmarks zoals RULER, InfiniteBench en Needle-in-a-Haystack.

• Afweging Nauwkeurigheid-Latentie:
  • Token Sparse Attention verbetert consistent de snelheid-nauwkeurigheid-afweging.
  • Bij 128K context bereikt het tot $\times 3.23$ versnelling in attentie ten opzichte van FlashAttention met minder dan 1% afname in nauwkeurigheid.
  • In combinatie met FlexPrefill verbetert het de versnelling van $\times 2.44$ naar $\times 2.76$ terwijl de gemiddelde nauwkeurigheid identiek blijft (87,27%).
• Vergelijking met Verwijderingsmethoden:
  • Het presteert aanzienlijk beter dan tokenverwijderingsmethoden (PyramidInfer, FastKV, GemFilter). Bijvoorbeeld, bij 128K bereikt Token Sparse Attention 86,84% nauwkeurigheid versus 78,49% voor PyramidInfer bij vergelijkbare versnellingen.
• Schaalbaarheid:
  • Versnellingswinsten nemen toe met de contextlengte (bijvoorbeeld minimaal voordeel bij 4K, aanzienlijk bij 128K) omdat attentiesparsiteit toeneemt met de sequentielengte.
• Overhead:
  • De overhead voor scorebepaling en indexbeheer is minimaal en staat voor <11% van de totale attentielatentie, zelfs bij hoge sparsiteitsniveaus.

5. Betekenis

Dit paper adresseert een kritieke kloof in inferentie van lange contexten bij LLM's door weg te bewegen van irreversibele tokenverwijdering naar dynamische, verweven sparsificatie.

• Theoretisch Inzicht: Het toont aan dat tokenbelang zeer dynamisch is over lagen en koppen heen, wat flexibele selectiemechanismen vereist in plaats van statisch snoeien.
• Praktische Impact: Het biedt een "plug-and-play"-optimalisatie die bovenop bestaande state-of-the-art attentiekernels (FlashAttention, Minference, enz.) kan worden gelegd om enorme versnellingen te bereiken in scenario's met lange contexten zonder modellen opnieuw te trainen of aanzienlijke nauwkeurigheid op te offeren.
• Toekomstige Richting: Het vestigt een nieuw ontwerppunt voor schaalbare inferentie, wat suggereert dat reversibele compressie een superieure strategie is voor het hanteren van de $O(L^2)$-bottleneck in vergelijking met permanente tokenverwijdering.