Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention

Titolo: Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention

Autori: Yuzhen Mao, Michael Y. Li, Emily B. Fox (Stanford University)

1. Il Problema

L'adattamento dei Large Language Models (LLM) a contesti lunghi (long-context) è ostacolato dal costo computazionale quadratico ($O(n^2)$) dei meccanismi di attenzione standard. Mentre le applicazioni reali (come la generazione di codice a livello di repository, il ragionamento complesso e i sistemi di agenti autonomi) richiedono la gestione di milioni di token, l'attenzione completa diventa proibitiva sia per l'addestramento che per l'inferenza.

Le soluzioni esistenti si dividono in due categorie principali, entrambe con limiti significativi:

• Selezione/Compressione KV-cache (Inferenza): Metodi come H2O o StreamingLLM rimuovono o selezionano token durante l'inferenza su modelli pre-addestrati con attenzione completa. Tuttavia, questi modelli non hanno imparato a gestire la sparsità, limitando la loro capacità di identificare i pattern ottimali.
• Attenzione Sparsa (Addestramento): Metodi recenti come NSA, DSA o MoBA introducono sparsità durante il pre-addestramento. Tuttavia, spesso richiedono modifiche architetturali (es. ramificazioni parallele), l'uso di moduli esterni non differenziabili (es. indexer) o operazioni non differenziabili (es. pooling medio) che impediscono l'ottimizzazione end-to-end del processo di selezione.

Inoltre, nelle configurazioni RAG (Retrieval-Augmented Generation), l'attenzione completa su tutti i documenti recuperati introduce rumore e distrazioni, poiché la maggior parte dei documenti è irrilevante per la query specifica.

2. Metodologia: Gist Sparse Attention (GSA)

Gli autori propongono GSA, un framework che colma il divario tra la compressione basata su "gist" (sintesi) e l'attenzione sparsa apprendibile, senza modificare l'architettura del Transformer.

Idea Chiave: Token Gist come Segnali di Routing

Il cuore di GSA risiede nell'uso di token "gist" intercalati (interleaved gist tokens). Questi token sono rappresentazioni apprese che riassumono blocchi locali di token grezzi. A differenza di semplici statistiche (come il pooling medio), i token gist sono ottimizzati durante l'addestramento per catturare il contenuto semantico del loro blocco.

Il processo avviene in quattro fasi:

1. Compressione: Il contesto viene partizionato in blocchi di lunghezza $L$. A ogni blocco viene associato un token gist ($g_m$) che apprende una sintesi compressa del blocco.
2. Valutazione della Rilevanza: Durante la decodifica, per ogni token di query ($q_t$), il modello calcola un punteggio di rilevanza per ogni token gist basato sul prodotto scalare tra la query e la chiave del gist ($s_{t,m} = q_t^\top k_{g_m}$).
3. Selezione Top-k: Vengono selezionati i $k$ blocchi (chunk) con i punteggi di gist più alti.
4. Svolgimento Selettivo (Selective Unfolding): Solo per i blocchi selezionati, il modello "svolge" (unfolds) la rappresentazione compressa, reintroducendo i token grezzi originali (raw tokens) nel contesto di attenzione insieme al loro token gist. I blocchi non selezionati rimangono compressi (solo gist) o vengono ignorati.

Questo crea un meccanismo coarse-to-fine (dal grezzo al fine): il modello mantiene una rappresentazione globale compatta e recupera i dettagli granulari solo dove sono necessari.

Estensione Gerarchica (H-GSA)

Per gestire contesti estremamente lunghi con complessità logaritmica, gli autori introducono una struttura gerarchica Gist-of-Gist:

• I token gist vengono a loro volta compressi in meta-gist tokens.
• La selezione avviene in modo ricorsivo: prima si selezionano i segmenti più rilevanti a livello di meta-gist, poi i sotto-blocchi a livello di gist, e infine i token grezzi.
• Questo riduce la complessità di decodifica per passo da lineare a logaritmica ($O(\log n)$).

Addestramento End-to-End

GSA è addestrato in due fasi (opzionali ma raccomandate):

1. Continued Pretraining: Il modello impara a comprimere il contesto in token gist usando una maschera causale specifica. A questo punto, il modello può già effettuare la selezione e lo svolgimento selettivo in inferenza.
2. Selective Finetuning (Opzionale): Il meccanismo di selezione top-k viene integrato direttamente nella maschera di attenzione durante l'addestramento, permettendo al modello di ottimizzare congiuntamente la qualità della compressione e l'accuratezza della selezione tramite backpropagation.

3. Contributi Chiave

• Pivot Apprendibili: Dimostrano che i token gist intercalati possono fungere da segnali di routing efficaci per l'attenzione sparsa, superando i limiti delle sintesi statistiche fisse.
• Framework Unificato: GSA unisce compressione e attenzione sparsa in un unico framework end-to-end, senza richiedere modifiche architetturali, moduli esterni o operazioni non differenziabili.
• Complessità Logaritmica: L'estensione gerarchica (H-GSA) permette di scalare a contesti illimitati mantenendo una complessità di decodifica logaritmica.
• Efficienza nel RAG: Elimina la necessità di attenzione incrociata completa tra documenti, utilizzando i token gist come surrogati per filtrare i documenti irrilevanti prima di accedere ai dettagli.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su LongBench e su cinque benchmark RAG (Multi-document QA), utilizzando modelli come Qwen2-7B e Llama3.2-1B.

• Prestazioni Superiori: GSA supera costantemente le baseline di compressione (come ActivationBeacon, UniGist, KVLink) e i metodi di attenzione sparsa solo in inferenza.
  • Su LongBench, GSA ottiene guadagni significativi (es. +3.7 punti su ActivationBeacon a 8x compression) mantenendo prestazioni vicine al modello con attenzione completa.
  • Su RAG, i risultati sono ancora più eclatanti: GSA supera di oltre 11 punti le baseline di compressione a 8x, e addirittura supera il modello con attenzione completa (Full-PT) durante il continued pretraining, dimostrando una capacità superiore a filtrare i documenti distrattori.
• Hierarchical GSA (H-GSA): A compressioni elevate (16x, 32x), la versione gerarchica supera la versione a singolo livello, confermando l'efficacia della selezione gerarchica.
• Ablation Study: La combinazione di token gist selezionati + token grezzi (SG+SR) si è rivelata la strategia migliore, confermando che i gist forniscono informazioni complementari utili anche quando i token grezzi sono ripristinati.
• Selezione Top-k vs Top-p: La selezione adattiva basata su Top-k si è dimostrata più stabile e performante rispetto alle soglie Top-p (nucleus sampling) su diversi task.

5. Significato e Impatto

Il lavoro di Mao et al. rappresenta un passo avanti significativo nell'efficienza dei LLM per contesti lunghi:

1. Semplicità Implementativa: Non richiede di riscrivere l'architettura del Transformer o di integrare componenti esterni complessi, rendendo GSA facilmente adottabile su modelli esistenti.
2. Ottimizzazione del Budget di Token: Trasforma la compressione da un processo di perdita irreversibile a un meccanismo dinamico di recupero delle informazioni, massimizzando l'utilità dei token disponibili.
3. Scalabilità: La soluzione gerarchica offre una via praticabile per gestire contesti di milioni di token con costi computazionali gestibili, avvicinandosi all'efficienza dei modelli a stato spaziale (come Mamba) pur mantenendo la potenza espressiva dell'attenzione.

In sintesi, GSA dimostra che "dimenticare" strategicamente (compressione) e "ricordare" selettivamente (svolgimento) è una strategia superiore rispetto all'elaborazione uniforme di tutti i token, specialmente in scenari dove il contesto è ricco di informazioni irrilevanti.