Sintesi Tecnica: Selezione con Gate di Incertezza per l'Attention Block-Sparse

Problema

I meccanismi di attenzione block-sparse (ad esempio, SSA, Quest, H2O, SnapKV) scalano i modelli linguistici a lungo contesto sostituendo la softmax $O(N^2)$ con una selezione top-$k$ per query su blocchi di chiavi. Sebbene efficiente, questo processo di selezione è miope. Quando i punteggi del $k$-esimo e del $(k+1)$-esimo blocco sono quasi in parità, il selettore compie una decisione senza budget aggiuntivo. Se il blocco scartato contiene prove critiche (ad esempio, nel ragionamento multi-hop o nel recupero "needle-in-a-haystack"), l'informazione viene persa in modo irreversibile, poiché gli strati a valle non possono recuperarla. Questo fallimento è strutturale, non meramente rumoroso, e degrada le prestazioni specialmente in sceni dove la rilevanza dipende dal contesto latente piuttosto che dai match superficiali tra query e key.

I metodi esistenti trattano il cutoff come una regola decisionale fissa. Essi non tengono conto dell'incertezza di tale decisione o del valore dell'informazione (VoI, Value of Information) che potrebbe essere guadagnato espandendo l'insieme di attenzione quando il cutoff è ambiguo.

Metodologia

Il documento propone un router con gate di incertezza che si posiziona sopra gli esistenti backbone di scoring per blocchi (come il mean-pooling in stile SSA o i limiti superiori K-min/K-max di Quest) senza richiedere riaddestramento o parametri aggiuntivi. Il metodo opera in cinque fasi:

1. Scoring dei Blocchi (Agnostico rispetto al Backbone): Le chiavi vengono raggruppate in blocchi. I punteggi sono calcolati utilizzando un backbone esistente (ad esempio, prodotto interno mediato tramite mean-pooling o limiti superiori per coordinata).
2. Selezione per Tile: Per ridurre l'overhead computazionale, la selezione viene eseguita per Query Tile (un gruppo di righe di query consecutive) anziché per singola riga. Il punteggio del tile è il punteggio massimo tra le righe del tile stesso. Una selezione top-$k$ determina l'insieme iniziale di blocchi ($kv\_idx$).
3. Misurazione dell'Incertezza (Segnale VoI): Il metodo calcola un margine di cutoff normalizzato ($\sigma$) per ogni tile e testa:
   $$\sigma = \frac{s_{(k-1)} - s_{(k)}}{s_{(0)} - s_{(k)}}$$
   dove $s_{(i)}$ sono i punteggi dei blocchi ordinati.
   • $\sigma \to 1$: Il cutoff è decisivo (l'insieme mantenuto è molto al di sopra della coda).
   • $\sigma \to 0$: Il cutoff è ambiguo (il $k$-esimo e il $(k+1)$-esimo blocco sono in parità).
     Questo $\sigma$ funge da proxy per la perdita attesa della decisione top-$k$; un $\sigma$ basso indica un alto valore nell'acquisire ulteriori informazioni.
4. Aggregazione e Triggering: I valori di $\sigma$ per testa vengono mediati in un segnale per tile ($\bar{\sigma}$). Un trigger viene attivato se $\bar{\sigma}$ scende al di sotto di una soglia di quantile empirica specifica per lo strato ($q$). Ciò identifica i tile "rischiosi" in cui la decisione era ambigua.
5. Espansione Controllata: I tile attivati espandono il loro insieme di attenzione per un fattore $\rho$ (ad esempio, $\rho=2$), prestando attenzione ai blocchi top-$\rho k$ invece di top-$k$. I tile non attivati mantengono il budget di base. L'implementazione utilizza un kernel fuso che esegue il padding dei tile non attivati con sentinelle per mantenere forme di dispatch uniformi, garantendo che il costo scalino in modo sub-lineare rispetto all'espansione.

Contributi Chiave

1. Formulazione VoI: Una formulazione del cutoff del selettore come decisione sotto incertezza, nuova e agnostica rispetto al backbone. Aggiunge solo uno scalare ($\sigma$) per tile e una soglia di quantile per strato, senza richiedere riaddestramento o parametri extra.
2. Componibilità: Dimostrazione empirica che il router si compone con distinti backbone di scoring (K-mean in stile SSA e K-max di Quest). Il router eleva costantemente le prestazioni del backbone che risulta superiore per un determinato modello, provando che "miglior scoring" e "migrazione del budget migliore" sono miglioramenti ortogonali.
3. Implementazione Fusa: Un kernel Triton personalizzato che gestisce selezione ed espansione in un unico percorso di codice, mantenendo il dispatch a forma uniforme. Ciò consente al sistema di girare a $0.62\times$ fino a $0.80\times$ il tempo di esecuzione dell'attenzione densa a 128K di contesto.
4. Benchmarking Diagnostico: Introduzione del benchmark Pointer-Chase Haystack (PCH) per isolare la qualità del selettore dalla capacità del modello, e l'uso di LongBench-v1 come controllo negativo per dimostrare che i benefici del router sono condizionati dal fatto che il budget del selettore sia stretto rispetto alla densità delle evidenze.

Risultati

Il metodo è stato valutato su quattro modelli (Qwen2.5-14B, Qwen2.5-7B-1M, Mistral-Nemo, Qwen3.6-35B) su contesti da 32K a 128K.

• Guadagni di Performance:
  • Su LongBench-v2 medium (32K contesto), la configurazione router-su-Quest ha ottenuto un recall accoppiato di 0.75 rispetto allo 0.47 del baseline top-$k$ (+28 punti percentuali), con un McNemar $p < 0.01$.
  • Su RULER NIAH multi-key, il router-su-Quest ha raggiunto un'accuratezza di 0.98 (rispetto a 0.93 per top-$k$) su Qwen2.5-14B e 1.00 su Qwen2.5-7B-1M.
  • Il router ha costantemente migliorato il backbone vincente per ogni modello. In particolare, il backbone ottimale cambiava in base all'architettura: il K-max di Quest vinceva su Qwen2.5/Mistral, mentre il K-mean in stile SSA vinceva su Qwen3.6 (a causa della regolarizzazione della varianza di magnitudo tramite QK-Norm). Il router ha elevato entrambi.
• Efficienza:
  • A 128K di contesto, il router ha preservato 0.81 dell'accuratezza densa su Qwen2.5-7B-1M e 0.89 su Qwen3.6, girando a 0.62$\times$ e 0.80$\times$ il tempo di esecuzione dell'attenzione densa, rispettivamente.
  • La "vittoria di Pareto" (migliore velocità a parità di qualità) emerge a contesti più lunghi (64K–128K) man mano che la quota di attenzione della computazione totale aumenta, in linea con la legge di Amdahl.
• Budget vs. Selettività: Gli studi di ablazione mostrano che il miglioramento del router deriva sia dall'aumento del budget che dall'allocazione selettiva. Sui modelli vicini al tetto delle prestazioni (es. Qwen-14B), il guadagno è principalmente guidato dal budget; sui modelli con più margine (es. Mistral-Nemo), la selettività (spendere il budget solo sui tile ambigui) fornisce la maggior parte del guadagno.

Significato e Rivendicazioni

Il documento afferma che il router con gate di incertezza fornisce uno strato generico e plug-and-play per l'attenzione block-sparse che mitiga il fallimento strutturale della selezione miope top-$k$. Il suo significato risiede in:

• Robustezza: Si adatta alle specifiche caratteristiche di scoring di diverse architetture di modelli (ad esempio, gestendo l'effetto QK-Norm in Qwen3.6) senza tuning manuale.
• Efficienza: Raggiunge un'accuratezza quasi densa a 128K di contesto con un tempo di esecuzione significativamente ridotto, rendendo più pratica l'inferenza a lungo contesto.
• Fondamento Teorico: Inquadra il problema della selezione attraverso la lente del Valore dell'Informazione (VoI), fornendo un limite di regret a campione finito sotto assunzioni di rumore Gaussiano (sebbene il documento noti che tale limite è attualmente debole nei regimi dominati dal rumore in cui opera il metodo).

Gli autori dichiarano esplicitamente che il miglioramento del router è condizionale: si attiva solo quando il budget per query è piccolo rispetto alla posizione delle evidenze rilevanti per la risposta. In scenari in cui il budget è generoso (ad esempio, LongBench-v1 con prompt nativi brevi), il router fornisce un beneficio trascurabile, fungendo da controllo negativo che valida la specificità del metodo per scenari a lungo contesto con budget ristretto.