Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Riepilogo Tecnico: Tensor Cache

Enunciato del Problema
L'inferenza autoregressiva dei Transformer si basa sulla memorizzazione di coppie Chiave-Valore (KV) per evitare di ricalcolare i prefissi. Tuttavia, lo stato KV mantenuto cresce linearmente con la lunghezza del contesto, la profondità e la concorrenza, creando un collo di bottiglia nella memoria. La memorizzazione a finestra scorrevole limita l'uso della memoria trattenendo solo i $W$ token più recenti, ma questo approccio soffre di "dimenticanza rigida": una volta che un token viene espulso dalla finestra, le sue informazioni vengono completamente scartate, rendendo inaccessibili prove rilevanti al di fuori della finestra. I metodi esistenti a stato limitato o scartano completamente i token espulsi o tentano di comprimere ogni token in una media mobile, il che può introdurre rumore o non riuscire a distinguere efficacemente tra contesto recente e distante.

Metodologia: Tensor Cache (TC)
Tensor Cache introduce un'architettura a stato limitato a due livelli che trasforma l'espulsione KV da un evento di cancellazione a un'operazione di scrittura in memoria associativa.

1. Architettura a Due Livelli:

   • L1 (Cache di Primo Livello): Un buffer ad anello KV di dimensioni fisse che contiene i $W$ token più recenti. Questo supporta l'attenzione softmax esatta per il contesto locale.
   • L2 (Cache di Secondo Livello): Una matrice di prodotto esterno $A$ di dimensioni fisse (memoria a pesi veloci) che memorizza informazioni esclusivamente dai token espulsi dalla finestra L1.
   • Meccanismo di Scrittura: Quando il buffer ad anello sovrascrive una voce, la coppia $(k, v)$ spostata non viene scartata. Invece, viene scritta nella matrice L2 $A$ tramite un aggiornamento del prodotto esterno: $A \leftarrow \lambda A + \eta (k \otimes v)$. Qui, $\lambda$ (decadimento) e $\eta$ (tasso di scrittura) sono scalari appresi per testa.
   • Meccanismo di Lettura: Le future query $q_t$ accedono al passato compresso calcolando $r_t = q_t A$. Sfruttando l'identità dell'attenzione lineare $q(k \otimes v) = \langle q, k \rangle v$, questa operazione recupera efficacemente informazioni da tutti i token espulsi senza memorizzarli singolarmente.
   • Fusione: Una porta scalare appresa $g$ fonde le uscite dell'attenzione locale (L1) e della lettura della memoria (L2): $y_t = y_{local} + \sigma(g) m_t$.

2. Ottimizzazione dell'Addestramento (Scansione Parallela):
   Il documento identifica un problema critico con le scorciatoie standard di addestramento a blocchi. Un approccio ingenuo riassume un blocco tramite la sua media ($\bar{k} \otimes \bar{v}$), il che introduce $C^2 - C$ prodotti esterni spurii tra token ($k_i \otimes v_j$ dove $i \neq j$) che non esistono durante l'inferenza in streaming.

   • Soluzione: Gli autori implementano una scansione parallela a somma pesata matematicamente equivalente alle scritture sequenziali per token. Questo colma il divario tra addestramento e inferenza entro l'epsilon float32 (errore relativo < $10^{-7}$), garantendo che il modello apprenda le dinamiche associative corrette senza il rumore dell'interferenza tra token.

Contributi Chiave

• Memoria Associativa Condizionata all'Espulsione: A differenza delle precedenti memorie a prodotto esterno (es. Infini-attention, mLSTM) che ingeriscono ogni token, TC utilizza la matrice a prodotto esterno rigorosamente come cache L2 alimentata solo dalle espulsioni della finestra scorrevole. Questo partiziona chiaramente i token in "residenti in L1" e "compressi in L2".
• Correzione di Artefatti di Addestramento: Il documento identifica e risolve una scorciatoia di addestramento precedentemente trascurata (media a blocchi) che introduce prodotti spurii tra token. La scansione parallela a somma pesata proposta elimina questa discrepanza allo stesso costo computazionale.
• Efficienza a Stato Limitato: TC fornisce un meccanismo in cui il contesto più vecchio influenza le previsioni senza mantenere singole voci espulse, offrendo un'alternativa pratica alla piena ritenzione KV.

Risultati Sperimentali
Gli autori hanno valutato TC su sistemi di scaling, richiamo associativo sintetico e modellazione linguistica su testo reale a lungo contesto (OpenWebText e Shakespeare).

• Richiamo Associativo Sintetico: In compiti che richiedono il recupero di informazioni oltre la finestra scorrevole, TC ha raggiunto una precisione del 100%, pareggiando le prestazioni di KV Completo. Ha ridotto lo stato di inferenza mantenuto del 72–84% rispetto a KV Completo, superando significativamente Infini-attention (che si degradava all'aumentare dei gap) e le baseline a finestra scorrevole (che performavano a livelli casuali).
• Modellazione Linguistica a Lungo Contesto: Su OpenWebText, TC ha ottenuto il più basso Negative Log-Likelihood (NLL) medio a ogni lunghezza di contesto di valutazione da 1.024 a 32.768 token, con l'unica eccezione di $L=2.048$ dove era statisticamente indistinguibile da Infini-attention. A $L=32.768$, TC ha raggiunto un NLL di 5,14, superando KV Completo (6,00) e Infini-attention (5,42), utilizzando 2,4$\times$ meno memoria di picco GPU rispetto a KV Completo.
• Throughput: TC ha comportato un sovraccarico di throughput di circa il 30% rispetto ai metodi puri a finestra scorrevole (dovuto alla lettura/aggiornamento L2), ma è rimasto costantemente più veloce di Infini-attention su tutte le lunghezze di contesto testate.
• Diagnostica della Capacità: Gli esperimenti sulla capacità pura di memoria associativa hanno mostrato che le prestazioni di recupero si degradano man mano che più associazioni non correlate vengono compresse nello stato fisso, confermando la natura limitata della memoria.

Significato e Affermazioni
Il documento afferma che Tensor Cache colma con successo il divario tra l'efficienza della memoria dell'attenzione a finestra scorrevole e le capacità di richiamo dell'attenzione a contesto completo. Trattando l'espulsione KV come un'operazione di memoria strutturata piuttosto che come una cancellazione distruttiva, TC permette ai Transformer di accedere a prove rilevanti al di fuori della finestra immediata senza il costo di memoria lineare delle cache KV complete. Gli autori sottolineano che il loro contributo principale non è l'invenzione della memoria a prodotto esterno in sé, ma la sua specifica applicazione come cache L2 condizionata all'espulsione e la rigorosa correzione degli artefatti di addestramento che in precedenza ostacolavano tali metodi. TC è presentato come un'alternativa pratica a stato limitato che migliora il confine memoria-qualità per l'inferenza a lungo contesto.