SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Technische Samenvatting: SpotAttention

Probleemstelling

Long-context capaciteiten zijn de standaard geworden voor vooraf getrainde Large Language Models (LLM's), wat taken mogelijk maakt zoals het verwerken van volledige codebases of het onderbouwen van antwoorden met uitgebreide bewijsvoering. Echter, de computationele kosten van self-attention schalen slecht met de sequentielengte. Tijdens de prefill-fase groeit de rekenkracht kwadratisch met de sequentielengte, en tijdens de decode-fase vereist elke nieuwe token het opnieuw lezen van een Key-Value (KV) cache die lineair groeit met de contextgrootte. Hoewel FlashAttention de I/O-efficiëntie verbeterde door het materialiseren van de volledige attention-matrix te vermijden, veranderde dit de fundamentele kwadratische prefill-kosten of de lineaire groei van per-stap geheugenleesacties niet.

Bestaande sparse attention-methoden proberen dit te verzachten door alleen aandacht te besteden aan een relevante subset van het verleden. Het selecteren van deze subset is echter vaak kostbaar. Methoden zoals Quest scoren token-blokken, en Twilight voegt top-$p$ pruning toe, maar de selector zelf wordt een nieuwe bottleneck. Boverevens vertrouwen training-vrije selectors vaak op handmatig ontworpen heuristieken die mogelijk niet goed generaliseren of meerdere selectiestadia vereisen, wat cumulatieve overhead veroorzaakt.

Methodologie: SpotAttention

SpotAttention is een plug-in, block-sparse routeringsmechanisme dat is ontworpen om te worden toegevoegd aan reeds getrainde, bevroren transformers. Het bestaat uit een lichtgewicht "selector" (of indexer) die aan elke full-attention laag is bevestigd.

1. Selector Architectuur

De selector is een minuscule multi-head Q-K scorer met een gereduceerde breedte ($d_{idx}$) en $H_{idx}$ heads. In plaats van individuele tokens te scoren, opereert het op blokken van queries en keys (bijv. 16 tokens per blok).

• Scoring Mechanisme: Voor een query op positie $q$ en een key-blok op positie $t$, berekent de selector een head-gewogen ReLU-score:
  $$s_s[q, t] = \sum_{h=1}^{H_{idx}} w_h \text{ReLU}(\mathbf{q}^I_h \cdot \mathbf{k}^I_{h,t})$$
  waarbij $\mathbf{q}^I_h$ en $\mathbf{k}^I_{h,t}$ per-head projecties zijn, en $w_h$ de mixing weights zijn afgeleid van de hidden state.
• Efficiëntie: Door bloksgewijs te scoren, draait het selectieproces als een enkele tensor-core matrixvermenigvuldiging, waardoor de overhead van token-niveau operaties wordt vermeden.

2. Training via KL Distillatie

De selector wordt onafhankelijk getraind tegen de dense attention-distributie van de bevroren backbone (de "teacher").

• Doelstelling: De training minimaliseert de Kullback-Leibler (KL) divergentie tussen de attention-distributie van de teacher ($p_\tau$) en de geschatte distributie van de selector ($p_\sigma$).
• Varianten: Het paper verkent twee KL-varianten:
  • DenseKL: Berekent de loss over alle $T$ keys.
  • SparseKL: Berekent de loss alleen over de eigen top-$K$ geselecteerde set van de selector, met genormaliseerde distributies.
• Convergentie: De selector convergeert snel (binnen ~200 stappen) over diverse backbone-groottes (4B tot 32B) en architecturen, waarbij het leert de distributie-vorm van de teacher te matchen.

3. Dual Top-$p$ Selectieregel

Tijdens inferentie scoort de selector alle voorgaande blokken. Om te bepalen aan welke blokken er aandacht moet worden besteed, introduceert SpotAttention een dual top-$p$ regel die het budget direct afleest uit de gekalibreerde distributie van de selector:

1. Reserved Blocks: Een vast prefix van "sink" blokken (bijv. de eerste 128 tokens) en een vast suffix van "recency" blokken (bijv. de laatste 256 tokens) worden altijd behouden om te voorkomen dat ze worden gedomineerd door de nucleus-selectie.
2. Nucleus Selectie: Top-$p$ (nucleus) sampling wordt toegepast op de residuele distributie (blokken exclusief sinks en recency) . De selectie stopt wanneer de cumulatieve waarschijnlijkheidsmassa een drempelwaarde $p$ bereikt.
3. Dynamisch Budget: Deze aanpak levert een variabel budget per query en per laag op, dat zich aanpast aan de inhoud zonder een aparte pruning-fase. Een minimumvloer ($K_{min}$) voorkomt degeneratieve gevallen waarin de nucleus te klein is.

4. Kwantisatie

De K-cache van de selector kan worden gekwantiseerd naar INT4 of FP4 microscale. Omdat de cache van de selector onderscheidend is van de KV-cache van de backbone, vermindert dit het geheugengebruik van de selector met ongeveer 3.5$\times$ zonder de precisie of nauwkeurigheid van de backbone te beïnvloeden.

Belangrijkste Bijdragen

1. Plug-In Block-Sparse Attention: Een lichtgewicht selector getraind via KL-distillatie die aan bevroren transformers kan worden toegevoegd. Het draait als een enkele tensor-core matmul, waardoor de zware selector-overhead van eerdere werken wordt vermeden.
2. Dual Top-$p$ Regel: Een content-adaptief selectiemechanisme dat het per-query, per-laag budget direct afleidt uit de geleerde distributie. Het reserveert sink- en recency-blokken om stabiliteit te garanderen, waarbij de statische top-$K$ nauwkeurigheid wordt geëvenaard met een kleiner gemiddeld budget.
3. Training en Inferentie Gedrag: Het paper demonstreert dat de selector de per-laag budget-asymmetrie van de backbone leert (vroege lagen besteden meer aandacht aan context dan latere lagen), een patroon dat een statische $K$ niet kan repliceren. De methode generaliseert van een 16K trainingslengte naar 128K evaluatie zonder verlies van nauwkeurigheid.

Resultaten

De methode is geëvalueerd op vijf backbones (Qwen3: 4B–32B; Qwen3.5: 4B–9B) over vier benchmarks (RULER, BABILong, InfiniteBench, LongBench-v2) bij contextlengtes tot 128K.

• Nauwkeurigheid: SpotAttention evenaart de nauwkeurigheid van de dense backbone bij contextlengtes tot 128K (8$\times$ de trainingslengte). Bij een budget van $K=0.5L$ blijft het binnen één bootstrap standaardfout van de dense baseline over alle modellen.
• Latency en Doorvoersnelheid:
  • Bij $L=128K$ op Qwen3-8B draait SpotAttention 3.9$\times$ sneller dan FlashAttention en 1.8$\times$ sneller dan Twilight (de sterkste training-vrije baseline).
  • De decode-doorvoersnelheid wordt constant gehouden op 70 tokens/s bij 128K en 36 tokens/s bij 512K, terwijl FlashAttention en Twilight out-of-memory (OOM) gaan bij respectievelijk 512K en 256K.
  • De selector zelf is aanzienlijk sneller dan de selector van Twilight (1.55 ms vs. 4.59 ms bij 128K).
• Kwantisatie: Het kwantiseren van de K-cache van de selector naar INT4 of FP4 verkleint het geheugengebruik met 3.5$\times$ zonder meetbaar verlies van nauwkeurigheid. Hoewel dit de single-stream decode-latency niet verbetert (die wordt gedomineerd door de backbone KV-bandbreedte), vermindert het het geheugengebruik en de prefill-kosten aanzienlijk.

Betekenis en Claims

Het paper positioneert SpotAttention als een oplossing voor de "selector bottleneck" in sparse attention. Door een gekalibreerde distributie te leren via KL-distillatie, maakt de selector een dynamisch, nucleus-gebaseerd budget mogelijk dat zich aanpast aan de specifieke behoeften van elke query en laag.

De auteurs beweren dat:

• Efficiëntie: Het de hoogste doorvoersnelheid bereikt onder de vergeleken methoden bij lange contexten (tot 512K), waarbij de geheugenbeperkingen worden overwonnen die dense en andere sparse baselines stoppen.
• Generalisatie: Een selector getraind op een vaste 16K context generaliseert effectief naar 128K en verder zonder hertraining.
• Modulariteit: Het een "plug-in" oplossing is die geen hertraining van de backbone vereist, waardoor het toepasbaar is op bestaande voorgetrainde modellen.
• Beperkingen: De auteurs merken op dat de grootste geëvalueerde backbone 32B is, en de resultaten momenteel beperkt zijn tot de Qwen-familie op Engelse benchmarks. Ze erkennen ook dat de sparse-prefill latency niet is gemeten, hoewel de architectuur dit ondersteunt, en dat de dual top-$p$ regel afhankelijk is van het feit dat de selector goed gekalibreerd is (wat standhoudt onder SparseKL training, maar instort onder DenseKL).