LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

1. Il Problema: L'Inflazione KV nei Modelli Diffusivi a Blocchi

I modelli linguistici basati sulla diffusione (Diffusion Language Models - DLM) offrono un'alternativa promettente ai decoder autoregressivi tradizionali, permettendo la generazione di token in qualsiasi ordine all'interno di un blocco. Tuttavia, l'inferenza su contesti lunghi rimane un collo di bottiglia a causa dell'attenzione basata sulla memoria.

Il paper identifica un problema critico chiamato "KV Inflation" (Inflazione della Cache Key-Value) quando si applicano metodi di attenzione sparsa standard (progettati per modelli autoregressivi) ai DLM:

• Meccanismo: In un blocco di $B$ query, ogni query seleziona un sottoinsieme diverso di posizioni KV dal prefisso.
• Conseguenza: Poiché l'attenzione è vincolata dalla memoria (memory-bound), il costo effettivo non è determinato dalla somma delle selezioni individuali, ma dall'unione di tutte le posizioni KV selezionate dalle $B$ query.
• Risultato: Anche se ogni query guarda solo $k$ token, l'unione può espandersi fino a $B \times k$, annullando i vantaggi di velocità previsti dalla sparsità e rendendo l'approccio "naïve" inefficiente.

2. Metodologia: LoSA (Locality-Aware Sparse Attention)

Gli autori osservano che, tra due step consecutivi di denoising, solo una piccola frazione di token subisce cambiamenti significativi nello stato nascosto (hidden state), mentre la maggior parte rimane stabile. Sfruttando questa località dei cambiamenti di rappresentazione, LoSA propone un approccio ibrido:

1. Classificazione dei Token:

   • Token Attivi (Active Tokens): Token che mostrano grandi variazioni nello stato nascosto (e quindi nei vettori di query) tra lo step $t-1$ e $t$.
   • Token Stabili (Stable Tokens): Token con variazioni minime o nulle.

2. Meccanismo di Ottimizzazione:

   • Per i Token Stabili: Invece di ricalcolare l'attenzione sparsa, LoSA riutilizza direttamente i risultati dell'attenzione del prefisso calcolati nello step precedente. Questo preserva l'informazione completa del prefisso denso senza alcun costo di accesso alla memoria KV.
   • Per i Token Attivi: Viene applicata l'attenzione sparsa (utilizzando l'algoritmo QUEST come selector) solo su questo sottoinsieme ridotto di token.

3. Riduzione dell'Unione KV:

   • Invece di calcolare l'unione delle posizioni KV per tutti i $B$ token del blocco, LoSA calcola l'unione solo per il set di token attivi $|A|$ (dove $|A| \ll B$).
   • Questo riduce drasticamente il numero di indici KV unici che devono essere caricati dalla memoria, abbassando il traffico di memoria e la latenza.

4. Flusso di Lavoro:

   • Inizializzazione: Al primo step di denoising di un blocco, viene eseguita un'attenzione densa per inizializzare le statistiche.
   • Iterazione Successiva: Si calcola il punteggio di località (MSE tra query $t$ e $t-1$), si selezionano i token attivi, si riutilizza la cache per gli stabili e si calcola l'attenzione sparsa solo per gli attivi.

3. Contributi Chiave

• Identificazione del Problema KV Inflation: Dimostrazione che l'applicazione diretta di tecniche di sparsità ai DLM fallisce a causa dell'espansione dell'unione delle posizioni KV.
• Osservazione della Località: Rilevamento che i cambiamenti di rappresentazione nei DLM sono altamente localizzati: solo pochi token cambiano significativamente tra gli step di denoising.
• Proposta LoSA: Un nuovo metodo di attenzione sparsa che combina il riutilizzo della cache per i token stabili e l'attenzione sparsa dinamica solo per i token attivi.
• Efficienza e Accuratezza: LoSA riduce la densità dell'attenzione mantenendo o migliorando l'accuratezza rispetto alle basi di riferimento, risolvendo il compromesso tipico tra velocità e qualità.

4. Risultati Sperimentali

Il metodo è stato valutato su modelli come Trado-8B, Trado-4B e SDAR-8B su benchmark come LongBench e task di ragionamento.

• Accuratezza:

  • Su LongBench, LoSA supera significativamente i metodi basati su sparsità naïve (come QUEST).
  • A livelli di sparsità aggressivi (budget 128), LoSA ottiene un miglioramento di +10.43 punti di accuratezza media rispetto a QUEST su Trado-8B.
  • Mantiene un'accuratezza vicina a quella dell'attenzione densa anche con sparsità elevate.

• Efficienza e Velocità:

  • Densità KV: LoSA riduce la densità della cache KV di un fattore medio di 1.54x rispetto alle baseline.
  • Velocità (Speedup): Su GPU NVIDIA RTX A6000, LoSA raggiunge uno speedup di 4.14x rispetto all'attenzione densa (contro 3.26x di QUEST).
  • Scalabilità: Il trend di velocità si mantiene anche su hardware più recente (RTX 5090), con uno speedup di 3.67x, confermando che il collo di bottiglia rimane il trasferimento dati (memory-bound).

5. Significato e Impatto

LoSA rappresenta un avanzamento fondamentale per l'inferenza efficiente dei modelli linguistici basati sulla diffusione.

• Superamento dei Limiti Attuali: Risolve il problema fondamentale che impediva l'uso efficace della sparsità nei DLM, permettendo loro di scalare a contesti lunghi senza sacrificare la qualità.
• Paradigma Ibrido: Introduce un nuovo paradigma che non si limita a selezionare quali token guardare, ma decide dinamicamente quali token meritano un ricalcolo e quali possono essere "congelati" riutilizzando la memoria.
• Applicabilità Pratica: I risultati dimostrano che è possibile ottenere accelerazioni hardware significative (fino a 4x) mantenendo le prestazioni del modello, rendendo i DLM più competitivi rispetto ai decoder autoregressivi tradizionali per compiti di ragionamento complesso e contesti estesi.

In sintesi, LoSA trasforma un problema di inefficienza (KV Inflation) in un'opportunità di ottimizzazione sfruttando la natura iterativa e locale delle modifiche negli stati nascosti dei modelli diffusivi.