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

Riepilogo Tecnico: SpotAttention

Problema

Le capacità di contesto lungo sono diventate uno standard nei Large Language Models (LLM) pre-addestrati, abilitando compiti come l'elaborazione di interi codebase o l'ancoraggio delle risposte su evidenze estese. Tuttavia, il costo computazionale dell'auto-attenzione (self-attention) scala in modo inefficiente con la lunghezza della sequenza. Durante la fase di prefill, il calcolo cresce quadraticamente con la lunghezza della sequenza e, durante la fase di decode, ogni nuovo token richiede la rilettura di una cache Key-Value (KV) che cresce linearmente con la dimensione del contesto. Sebbene FlashAttention abbia migliorato l'efficienza di I/O evitando la materializzazione dell'intera matrice di attenzione, non ha alterato il costo fondamentale quadratico del prefill né la crescita lineare delle letture di memoria per singolo step.

I metodi esistenti di attenzione sparsa tentano di mitigare questo problema tentando di prestare attenzione solo a un sottoinsieme rilevante di token passati. Tuttavia, la selezione di questo sottoinsieme è spesso costosa. Metodi come Quest valutano i blocchi di token e Twilight aggiunge il pruning top-$p$, ma il selettore stesso diventa un nuovo collo di bottiglia. Inoltre, i selettori privi di addestramento (training-free) si affidano spesso a euristiche progettate a mano che potrebbero non generalizzare bene o richiedere più fasi di selezione, incorrendo in un overhead cumulativo.

Metodologia: SpotAttention

SpotAttention è un meccanismo di routing block-sparse "plug-in", progettato per essere applicato a transformer già addestrati e congelati (frozen). Consiste in un "selettore" (o indizzatore) leggero attaccato a ogni strato di full-attention.

1. Architettura del Selettore

Il selettore è un minuscolo punteggiatore Q-K multi-head con larghezza ridotta ($d_{idx}$) e $H_{idx}$ teste. Invece di punteggiare i singoli token, opera su blocchi di query e chiavi (ad esempio, 16 token per blocco).

• Meccanismo di Punteggio: Per una query alla posizione $q$ e un blocco di chiavi alla posizione $t$, il selettore calcola uno score ReLU pesato per testa:
  $$s_s[q, t] = \sum_{h=1}^{H_{idx}} w_h \text{ReLU}(\mathbf{q}^I_h \cdot \mathbf{k}^I_{h,t})$$
  dove $\mathbf{q}^I_h$ e $\mathbf{k}^I_{h,t}$ sono proiezioni per singola testa, e $w_h$ sono pesi di miscelazione derivati dallo stato nascosto.
• Efficienza: Punteggiando a livello di blocco, il processo di selezione viene eseguito come una singola moltiplicazione di matrici su tensor-core, evitando l'overhead delle operazioni a livello di token.

2. Addestramento tramite Distillazione KL

Il selettore è addestrato indipendentamente contro la distribuzione di attenzione densa del backbone congelato (il "teacher").

• Obiettivo: L'addestramento minimizza la divergenza di Kullback-Leibler (KL) tra la distribuzione di attenzione del teacher ($p_\tau$) e la distribuzione stimata dal selettore ($p_\sigma$).
• Varianti: Il paper esplora due varianti KL:
  • DenseKL: Calcola la perdita su tutti i $T$ chiavi.
  • SparseKL: Calcola la perdita solo sul proprio set selezionato top-$K$ del selettore, con distribuzioni rinormalizzate.
• Convergenza: Il selettore converge rapidamente (entro circa 200 step) attraverso varie dimensioni di backbone (da 4B a 32B) e architetture, imparando a corrispondere alla forma della distribuzione del teacher.

3. Regola di Selezione Dual Top-$p$

In fase di inferenza, il selettore valuta tutti i blocchi passati. Per determinare a quali blocchi prestare attenzione, SpotAttention introduce una regola dual top-$p$ che legge il budget direttamente dalla distribuzione calibrata del selettore:

1. Blocchi Riservati: Un prefisso fisso di blocchi "sink" (es. i primi 128 token) e un suffisso fisso di blocchi di "recency" (es. gli ultimi 256 token) sono sempre mantenuti per evitare che siano dominati dalla selezione nucleare.
2. Selezione Nucleus: Il campionamento top-$p$ (nucleus) viene applicato alla distribuzione residua (blocchi esclusi da sink e recency). La selezione si ferma quando la massa di probabilità cumulativa raggiunge una soglia $p$.
3. Budget Dinamico: Questo approccio produce un budget variabile per query e per strato, adattandosi al contenuto senza una fase di pruning separata. Un pavimento minimo ($K_{min}$) evita casi degeneri in cui il nucleus è troppo piccolo.

4. Quantizzazione

La cache K del selettore può essere quantizzata a microscala INT4 o FP4. Poiché la cache del selettore è distinta dalla cache KV del backbone, ciò riduce l'impronta di memoria del selettore di circa 3,5$\times$ senza influenzare la precisione o l'accuratezza del backbone.

Contributi Chiave

1. Attenzione Block-Sparse Plug-In: Un selettore leggero addestrato tramite distillazione KL che si attacca a transformer congelati. Esegue una singola moltiplicazione di matrici su tensor-core, evitando l'oneroso overhead del selettore visto nel lavoro precedente.
2. Regola Dual Top-$p$: Un meccanismo di selezione adattivo al contenuto che deriva il budget per query e per strato direttamente dalla distribuzione appresa. Riserva i blocchi sink e di recency per garantire la stabilità, eguagliando l'accuratezza del top-$K$ statico con un budget medio inferiore.
3. Comportamento di Training e Inferenza: Il paper dimostra che il selettore apprende l'asimmetria del budget per strato del backbone (gli strati iniziali prestano attenzione a più contesto rispetto a quelli successivi), un pattern che il $K$ statico non può replicare. Il metodo generalizza da una lunghezza di addestramento di 16K a una di valutazione di 128K senza degradazione dell'accuratezza.

Risultati

Il metodo è stato valutato su cinque backbone (Qwen3: 4B–32B; Qwen3.5: 4B–9B) su quattro benchmark (RULER, BABILong, InfiniteBench, LongBench-v2) con lunghezze di contesto fino a 128K.

• Accuratezza: SpotAttention eguaglia l'accuratezza del backbone denso a lunghezze di contesto fino a 128K (8$\times$ la lunghezza di addestramento). Con un budget di $K=0,5L$, rimane entro un errore standard di bootstrap del baseline denso in tutti i modelli.
• Latenza e Throughput:
  • A $L=128K$ su Qwen3-8B, SpotAttention è 3,9$\times$ più veloce di FlashAttention e 1,8$\times$ più veloce di Twilight (il più forte baseline training-free).
  • Il throughput di decode è sostenuto a 70 token/s a 128K e 36 token/s a 512K, mentre FlashAttention e Twilight vanno in Out-of-Memory (OOM) rispettivamente a 512K e 256K.
  • Il selettore stesso è significativamente più veloce del selettore di Twilight (1,55 ms vs 4,59 ms a 128K).
• Quantizzazione: La quantizzazione della cache K del selettore a INT4 o FP4 riduce la sua impronta di memoria di 3,5$\times$ senza alcuna perdita misurabile di accuratezza. Sebbene ciò non migliori la latenza di decode a singolo stream (dominata dalla larghezza di banda della KV del backbone), riduce significativamente l'uso della memoria e i costi di prefill.

Significato e Affermazioni

Il paper posiziona SpotAttention come una soluzione al "collo di bottiglia del selettore" nell'attenzione sparsa. Imparando una distribuzione calibrata tramite distillazione KL, il selettore abilita un budget basato sul nucleus dinamico che si adatta alle specifiche esigenze di ogni query e strato.

Gli autori affermano che:

• Efficienza: Raggiunge il throughput più elevato tra i metodi confrontati ai contesti lunghi (fino a 512K), superando le limitazioni di memoria che bloccano i baseline densi e altri metodi sparsi.
• Generalizzazione: Un selettore addestrato su un contesto fisso di 16K generalizza efficacemente a 128K e oltre senza ri-addestramento.
• Modularità: È una soluzione "plug-in" che non richiede il ri-addestramento del backbone, rendendolo applicabile a modelli pre-addestrati esistenti.
• Limitazioni: Gli autori notano che il backbone più grande valutato è 32B, e i risultati sono attualmente limitati alla famiglia Qwen su benchmark inglesi. Riconoscono inoltre che la latenza di sparse-prefill non è stata testata, sebbene l'architettura lo supporti, e che la regola dual top-$p$ dipende dal fatto che il selettore sia ben calibrato (cosa che si verifica sotto l'addestramento SparseKL ma collassa sotto DenseKL).