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

1. Enunciazione del Problema

I Large Language Models (LLM) affrontano un collo di bottiglia fondamentale nell'inferenza a contesto lungo a causa della complessità quadratica ($O(L^2)$) del meccanismo di attenzione durante la fase di prefill. Sebbene le ottimizzazioni hardware come FlashAttention riducano il sovraccarico di I/O della memoria, non alterano la complessità algoritmica.

I metodi di accelerazione esistenti soffrono di due limitazioni principali:

• Sparsità Strutturata: Metodi come Minference o FlexPrefill applicano una sparsità a livello di blocco. Questo spesso mantiene token irrilevanti se sono raggruppati con token importanti nello stesso blocco, portando a una sparsità sub-ottimale.
• Espulsione dei Token: Metodi come FastKV o GemFilter scartano permanentemente i token nei primi livelli. Questo assume che l'importanza del token sia statica, ignorando la natura dinamica della rilevanza dei token attraverso i diversi livelli e le diverse testate di attenzione. Una volta espulso, un token non può essere riesaminato, portando a un degrado dell'accuratezza.

2. Metodologia: Token Sparse Attention

Gli autori propongono Token Sparse Attention, un meccanismo di sparsificazione a livello di token leggero, dinamico e reversibile. A differenza dei metodi basati sull'espulsione, non rimuove permanentemente i token ma comprime lo spazio di calcolo e successivamente ripristina la sequenza.

Architettura Principale

Il metodo opera in due fasi per ogni testata di attenzione:

1. Compressione (Selezione QKV):
   • Per ogni testata di attenzione $h$, viene selezionato un sottoinsieme di indici di token $S_h$ basato su punteggi di importanza.
   • I tensori Query ($Q$), Key ($K$) e Value ($V$) vengono raccolti per formare tensori compressi $\hat{Q}, \hat{K}, \hat{V}$ con una lunghezza di sequenza ridotta $L' \ll L$.
   • Caratteristica Chiave: La selezione viene eseguita indipendentemente per ogni testata, consentendo a diverse testate di concentrarsi su diversi sottoinsiemi di token, affrontando l'eterogeneità tra le testate.
   • I tensori compressi rimangono densi e contigui, consentendo l'uso di kernel ottimizzati come FlashAttention.
2. Decompressione (Ripristino dell'Output):
   • L'output di attenzione $\hat{O}$ (forma $L' \times d$) viene disperso nuovamente in un tensor inizializzato a zero della lunghezza di sequenza originale $L \times d$ utilizzando gli indici specifici della testata.
   • Caratteristica Chiave: Le posizioni non selezionate vengono azzerate, ma la connessione residua dal livello precedente preserva l'informazione dei token saltati. Questo permette al modello di "rivedere" e rivalutare l'importanza del token nei livelli successivi.

Copertura Dinamica dei Token

Per determinare quali token mantenere e quanti, il documento introduce la Copertura Dinamica dei Token:

• Punteggio: Viene calinata una mappa di attenzione proxy leggera utilizzando le query recenti contro tutte le chiavi. I punteggi di importanza dei token sono derivati sommando i pesi di attenzione verticalmente.
• Allocazione del Budget: Invece di mantenere i top-$k$ token, il metodo ordina i token per importanza crescente e pota la "coda" (i token meno importanti) fino a quando la massa cumulativa dei token potati non supera una soglia di copertura $\tau$. Questo tratta la potatura come regolarizzazione strutturale contro il rumore di attenzione.
• Selezione per Testata: Dopo aver determinato il budget globale, ogni testata seleziona indipendentemente i suoi top-$k$ token in base ai suoi punteggi specifici.

Selezione del Livello Sparsa

Gli autori osservano che applicare la sparsità a tutti i livelli degrada le prestazioni. Introducono il Deriva della Rappresentazione Inter-Livello ($R_\ell$) per misurare la stabilità delle rappresentazioni dei token tra i livelli.

• I livelli con bassa deriva (rappresentazioni stabili) vengono selezionati per la sparsificazione.
• I livelli con alta deriva (dove le rappresentazioni cambiano significativamente) vengono mantenuti densi per preservare il comportamento del modello.
• Questa selezione è un passaggio di pre-elaborazione una tantum per modello.

3. Contributi Chiave

1. Sparsificazione a Livello di Token Reversibile: Un design innovativo "Comprimi e Decomprimi" che riduce il costo computazionale senza espellere permanentemente i token, consentendo una rivalutazione dinamica dell'importanza dei token attraverso i livelli.
2. Eterogeneità per Testata: Il metodo supporta la selezione indipendente dei token per ogni testata di attenzione, catturando il fatto che diverse testate danno priorità a diverse caratteristiche semantiche.
3. Compatibilità Senza Attriti: L'approccio è pienamente compatibile con i kernel esistenti densi (FlashAttention) e sparsi (Minference, FlexPrefill), agendo come un acceleratore complementare piuttosto che come un sostituto.
4. Budgeting Dinamico: Un approccio guidato dai dati per determinare i livelli di sparsità basati sulle distribuzioni di importanza dei token e sulla stabilità per livello, piuttosto che su rapporti fissi.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su LLaMA-3.1-8B-Instruct e Mistral-Nemo-12B-Instruct utilizzando benchmark come RULER, InfiniteBench e Needle-in-a-Haystack.

• Compromesso Accuratezza-Latenza:
  • Token Sparse Attention migliora costantemente il compromesso velocità-accuratezza.
  • A 128K di contesto, raggiunge un'accelerazione dell'attenzione fino a $\times 3.23$ rispetto a FlashAttention con meno dell'1% di degrado dell'accuratezza.
  • Se combinato con FlexPrefill, migliora l'accelerazione da $\times 2.44$ a $\times 2.76$ mantenendo un'accuratezza media identica (87,27%).
• Confronto con Metodi di Espulsione:
  • Supera significativamente i metodi di espulsione dei token (PyramidInfer, FastKV, GemFilter). Ad esempio, a 128K, Token Sparse Attention raggiunge un'accuratezza dell'86,84% contro il 78,49% di PyramidInfer con accelerazioni simili.
• Scalabilità:
  • I guadagni di accelerazione aumentano con la lunghezza del contesto (ad esempio, guadagno minimo a 4K, sostanziale a 128K) perché la sparsità dell'attenzione aumenta con la lunghezza della sequenza.
• Sovraccarico:
  • Il sovraccarico per la valutazione e la gestione degli indici è minimo, rappresentando <11% della latenza totale di attenzione anche a livelli elevati di sparsità.

5. Significato

Questo documento affronta un divario critico nell'inferenza LLM a contesto lungo spostandosi dall'espulsione irreversibile dei token verso una sparsificazione dinamica e intercalata.

• Insight Teorico: Dimostra che l'importanza del token è altamente dinamica attraverso i livelli e le testate, rendendo necessari meccanismi di selezione flessibili piuttosto che una potatura statica.
• Impatto Pratico: Fornisce un'ottimizzazione "plug-and-play" che può essere sovrapposta ai kernel di attenzione all'avanguardia esistenti (FlashAttention, Minference, ecc.) per ottenere accelerazioni massive in scenari a contesto lungo senza riaddestrare i modelli o sacrificare un'accuratezza significativa.
• Direzione Futura: Stabilisce un nuovo punto di progettazione per l'inferenza scalabile, suggerendo che la compressione reversibile è una strategia superiore per gestire il collo di bottiglia $O(L^2)$ rispetto alla rimozione permanente dei token.