🤖 machine learning

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

Questo articolo introduce la Looped Latent Attention (LLA), un codec post-training che sfrutta la struttura a basso rango delle cache KV indicizzate a loop nei Transformer con pesi condivisi (weight-tied) per ottenere una compressione estrema (fino a 32x) e aumentare significativamente la capacità di batch, mantenendo prestazioni quasi prive di perdite attraverso l'inizializzazione SVD e la distillazione.

Autori originali: James O' Neill, Fergal Reid

Pubblicato 2026-07-20

📖 6 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: James O' Neill, Fergal Reid

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Il Problema della Memoria nei Cervelli Artificiali

Immaginate di cercare di insegnare a un robot come scrivere una storia. Il robot ha un cervello fatto di milioni di minuscoli interruttori (parametri) che sanno come le parole si incastrano tra loro. Ma ecco il problema: per scrivere una storia lunga, il robot non può limitarsi a guardare il proprio cervello; ha bisogno di un blocco per gli appunti. Ogni volta che scrive una nuova parola, deve ricordare tutto ciò che ha appena scritto per non ripetersi o perdere il filo del discorso. Nel mondo dell'Intelligenza Artificiale, questo blocco per gli appunti è chiamato KV cache (cache Key-Value). È un taccuino privato per ogni singola conversazione, e cresce di dimensioni con ogni parola che il robot pronuncia.

Ora, immaginate un tipo speciale di robot progettato per essere super efficiente. Invece di avere un cervello enorme con molte stanze diverse (layer) per compiti differenti, questo robot ha una singola stanza piccola e intelligente che visita ripetutamente. Riutilizza lo stesso set di istruzioni, ciclando attraverso di esse per pensare più a fondo. Questo è chiamato un Looped Transformer (Transformer a ciclo). Risparmia spazio nel cervello stesso, il che è un bene. Ma c'è un problema subdolo: anche se il robot riutilizza la stessa stanza, scrive comunque una nuova pagina nel suo blocco per appunti ogni singola volta che esegue il ciclo. Se il robot cicla quattro volte per risolvere un problema di matematica, finisce per avere quattro pagine separate di note per lo stesso momento della storia. Ciò significa che il robot "a ciclo" ha ancora bisogno di una quantità massiccia di memoria per funzionare, vanificando lo scopo di essere efficiente. La grande domanda per gli scienziati è: possiamo rimpicciolire quel blocco per appunti senza far dimenticare al robot come pensare?

La Scoperta del Paper: Una Scorciatoia Segreta nelle Note

Questo articolo introduce un trucco astuto chiamato Looped Latent Attention (LLA) per risolvere quel problema di memoria. I ricercatori hanno scoperto che le note del robot non sono così disordinate come sembrano. Quando il robot cicla attraverso il suo processo di pensiero, le note che scrive per la stessa parola non cambiano drasticamente; seguono un percorso fluido e prevedibile, come una biglia che rotola giù da una collina dolce. Invece di salvare ogni singola pagina del blocco per appunti del robot, il nuovo metodo salva solo un piccolo "riassunto" di quel percorso e un set di istruzioni su come ricostruire le pagine specifiche solo quando il robot ha effettivamente bisogno di leggerle.

Pensatelo in questo modo: Immaginate di guardare un film in cui il personaggio principale cammina lungo un corridoio. In un computer normale, salvereste una foto completa ad alta definizione del personaggio ad ogni singolo passo che compie. Questo occupa una enorme quantità di spazio di archiviazione. Ma con LLA, vi rendete conto che il personaggio sta solo camminando in linea retta. Quindi, invece di salvare migliaia di foto, salvate una foto del personaggio e una semplice nota che dice: "Muoviti di 1 pollice in avanti per il passo 2, 2 pollici per il passo 3". Quando il computer ha bisogno di vedere il personaggio al passo 3, disegna rapidamente quella figura usando la nota. Questo è esattamente ciò che fa il paper per l'IA: comprime la parte del "ciclo" della memoria, trasformando una pila di quattro pagine in un singolo, minuscolo riassunto che può essere espanso nuovamente nelle pagine complete istantaneamente.

I ricercatori hanno testato questo approccio su diversi modelli di IA, incluso uno chiamato Ouro-1.4B. Hanno scoperto che questo metodo funziona incredibilmente bene. Infatti, su un potente chip per computer chiamato H200, sono stati in grado di far contenere 768 diverse conversazioni nella memoria contemporaneamente, mentre prima potevano contenerne solo 32. Questo è un salto enorme in termini di quante persone possono usare l'IA nello stesso momento. Hanno anche dimostrato che questo trucco funziona anche quando l'IA risolve problemi matematici molto lunghi e difficili, anche se ha bisogno di un po' di addestramento extra per assicurarsi di non confondersi quando sta scrivendo la propria storia invece di limitarsi a copiare un insegnante.

A Cosa Questo Paper Dice "No"

È importante sapere cosa questo metodo non è. I ricercatori hanno testato un'idea molto semplice per prima: "E se buttassimo via le prime tre pagine di note e tenessimo solo l'ultima?". Pensavano che forse il robot si stabilizzi verso una risposta finale dopo alcuni cicli, quindi le note precedenti non sarebbero state importanti. Ci hanno provato, ed è stato un disastro. L'IA ha completamente dimenticato come fare matematica, ottenendo zero nei test. Questo dimostra che il "percorso" che le note compiono è importante; non puoi semplicemente prendere il punto finale. Devi catturare il viaggio.

Hanno anche confrontato il loro metodo con altri modi per rimpicciolire la memoria, come schiacciare le note insieme (quantizzazione) o condividere le note tra diverse parti del cervello (compressione head-axis). Sebbene quei metodi aiutino un po', il paper mostra che comprimere l'asse del "ciclo" è la formula magica. Se si prova a comprimere le altre parti invece dei cicli, l'IA perde la capacità di ragionare. Il paper è molto chiaro: il ciclo è la parte più ridondante della memoria, ed è l'unico posto in cui si dovrebbe stringere.

Quanto Sono Sicuri?

Gli autori sono piuttosto fiduciosi in questi risultati perché li hanno misurati direttamente. Non hanno solo tirato a indovinare; hanno eseguito l'IA su compiti reali come la risoluzione di problemi matematici (GSM8K) e la scrittura di codice (HumanEval). Hanno dimostrato che con una compressione di 4x (rendendo la memoria 4 volte più piccola), l'IA si comporta quasi altrettanto bene della versione originale, non compressa. Quando hanno spinto la compressione all'estremo (21.3x), l'IA era ancora in grado di gestire molti compiti, sebbene iniziasse a faticare con i passaggi di ragionamento più lunghi e complessi.

Hanno anche testato questo su un tipo diverso di modello di IA chiamato Huginn-3.5B, che ha una struttura cerebrale diversa e cicla fino a 32 volte. Anche lì, il metodo ha funzionato, dimostrando che questo non è solo un colpo di fortuna per un robot specifico. Il paper suggerisce che questa "traiettoria a basso rango" (il percorso fluido delle note) è una proprietà fondamentale del modo in cui questi robot a ciclo pensano. Sebbene abbiano dimostrato che funziona per risparmiare memoria e aumentare la capacità, notano che per la prestazione assoluta migliore su compiti molto lunghi, l'IA ha bisogno di un secondo round di addestramento dove pratica con le proprie note compresse, invece di limitarsi a copiare un insegnante. Questo addestramento "on-policy" corregge i piccoli errori che accadono quando il robot inizia a scrivere la propria storia.

In breve, questo paper mostra che possiamo rendere i modelli di IA a ciclo molto più efficienti rendendoci conto che la loro memoria è ripetitiva in un modo specifico. Archivando un piccolo riassunto di quella ripetizione invece delle note complete, possiamo far contenere centinaia di volte più conversazioni su un singolo chip per computer, rendendo l'IA potente accessibile a più persone senza la necessità di hardware super costosi.

Sintesi Tecnica: Looped Latent Attention (LLA)

Problema

Nei Transformer con pesi condivisi e ciclici (Looped, weight-tied Transformers), un singolo blocco viene applicato ripetutamente (ricorrenza) per ridurre il numero di parametri. Sebbene ciò comprima efficacemente i pesi del modello, il collo di bottiglia della memoria durante l'inferenza rimane la cache Key/Value (KV). Nell'inferenza ciclica standard, il modello ricalcola l'attenzione ad ogni passo di ricorrenza, memorizzando una cache KV separata per ogni passo, strato e testa. Di conseguenza, un modello con $T$ cicli su $D$ strati trasporta una cache KV equivalente a un decoder di $T \times D$ strati, annullando i benefici di memoria derivanti dalla condivisione dei pesi.

Il paper postula che questa cache indicizzata per ciclo sia altamente strutturata. Per un dato token, strato e testa, i vettori K/V attraverso i passi di ricorrenza tracciano una traiettoria breve e a basso rango, piuttosto che esplorare uno spazio ad alta dimensionalità. La ridondanza è specifica rispetto all'asse del ciclo; gli spettri lungo gli assi delle teste e degli strati rimangono significativamente più piatti (rango più elevato). La sfida consiste nel comprimere questo asse di ricorrenza senza far collassare la traiettoria in un singolo punto fisso, il che degraderebbe le prestazioni.

Metodologia: Looped Latent Attention (LLA)

Gli autori introducono la Looped Latent Attention (LLA), un codec per la cache post-training che sostituisce i $T$ vettori per-ciclo per K/V con una rappresentazione latente compatta e mappe di ricostruzione specifiche per il ciclo.

Meccanismo Core

Memorizzazione Latente: Invece di memorizzare $T$ $T$ vettori ( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ ) per un determinato token, strato e testa, LLA memorizza un singolo vettore latente centrato a basso rango $c \in \mathbb{R}^r$ $c \in R^{r}$ .
- $c = (x - \mu)W_{\text{down}}$
- Qui, $\mu$ è la media di un set di calibrazione, e $W_{\text{down}}$ è una matrice di proiezione verso il basso.
Ricostruzione Specifica per Ciclo: Il vettore K/V specifico per il passo di ricorrenza $t$ $t$ viene ricostruito utilizzando una proiezione verso l'alto specifica per il ciclo:
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
Ranchi Asimmetrici: Il metodo utilizza ranghi separati per Chiavi ( $r_k$ ) e Valori ( $r_v$ ). L'analisi empirica mostra che i vettori Value si stabilizzano più lentamente attraverso i cicli rispetto ai vettori Key, motivando l'uso di $r_v > r_k$ .
Variante LLA-2D: Per una compressione estrema, LLA-2D ripiega l'asse delle teste nel latente, memorizzando un singolo latente per token e strato per ricostruire tutte le $H$ teste. Ciò riduce la dimensione della cache di un fattore $H$ ma accoppia le teste.

Addestramento e Inizializzazione

Inizializzazione tramite SVD: Il codec viene inizializzato tramite la Decomposizione ai Valori Singolari (SVD) delle attivazioni del modello "teacher" su un set di calibrazione. La proiezione verso il basso è impostata sui primi $r$ vettori singolari destri. Questa inizializzazione è critica; un'inizializzazione casuale risulta in una divergenza KL significativamente più alta.
Distillazione: Il codec viene raffinato congelando il modello teacher e minimizzando la divergenza KL forward tra la distribuzione del teacher e la distribuzione del modello con il codec applicato, aggiungendo un termine di matching dell'output dell'attenzione.
Raffinamento On-Policy: Per affrontare il bias di esposizione nei lunghi rollout autoregressivi (dove il modello condiziona sui propri prefissi generati invece che su quelli del teacher), viene impiegata una seconda fase di distillazione on-policy. Il codec viene aggiornato utilizzando prefissi campionati dal modello studente stesso, stabilizzando le prestazioni nelle generazioni lunghe.

Contributi Chiave

Definizione di LLA: Un codec K/V latente cross-loop per Transformer con pesi condivisi che raggiunge un rapporto di compressione $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ . Memorizza $(r_k + r_v)H$ scalari per token per strato invece di $2Td_{\text{head}}$ .
Compressione Specifica per Asse: Il paper dimostra che l'asse di ricorrenza (ciclo) è la dimensione più comprimibile nei modelli ciclici (Ouro-1.4B, Ouro-2.6B-Thinking, e Huginn-3.5B). Comprimere l'asse del ciclo supera le prestazioni della compressione dell'asse delle teste (MLA), dell'asse degli strati (condivisione cross-layer) o dell'applicazione della quantizzazione a parità di budget della cache.
Rifiuto del Collasso in Punto Fisso: Gli autori mostrano che la traiettoria del ciclo non può essere sostituita da un singolo stato (l'ultimo ciclo). Un controllo a zero parametri che riutilizza solo il K/V dell'ultimo ciclo provoca il collasso delle prestazioni di generazione (l'accuratezza su GSM8K scende a 0.000), provando che la "breve traiettoria a basso rango" è necessaria.
Impatto sul Deployment: La riduzione esatta degli scalari si traduce direttamente in un aumento della capacità di serving. Su una singola GPU H200, LLA aumenta la dimensione massima del batch per Ouro-1.4B a 4k di contesto da 32 a 768 sequenze (con una compressione di 21.3×).

Risultati Sperimentali

Compressione vs. Accuratezza

Prestazioni GSM8K: Su Ouro-1.4B, una compressione leggera (da 1.33× a 2×) eguaglia l'accuratezza del teacher (circa 0.79–0.80). Con una compressione di 4×, LLA mantiene un'accuratezza di 0.59, superando significativamente la compressione dell'asse delle teste (0.52) e il riutilizzo dell'ultimo ciclo (0.00).
Generalizzazione di Scala: Su Ouro-2.6B-Thinking, LLA rimane quasi privo di perdite (0.81–0.83) tra 4× e 10.67× di compressione, mentre i baseline dell'asse delle teste perdono sostanziale accuratezza.
Trasferimento a Huginn: Il metodo si trasferisce a Huginn-3.5B (32 passi di ricorrenza). Un codec SVD senza addestramento raggiunge una fedeltà quasi perfetta fino a 32× in metriche indipendenti dal decoder.

Stabilità delle Generazioni Lunghe

MATH-500: I lunghi rollout matematici sono sensibili al drift del prefisso. L'addestramento off-policy (usando i prefissi del teacher) porta alla degradazione dell'accuratezza e ad alti tassi di "nessuna risposta" all'aumentare della lunghezza della generazione.
Guadagno On-Policy: L'applicazione del raffinamento on-policy eleva l'accuratezza su MATH-500 da 0.43 a 0.66 con 4× di compressione e riduce significativamente le generazioni senza risposta.

Capacità di Serving

Riduzione della Memoria: La riduzione della cache è esatta. Con una compressione di 21.3×, la cache KV per Ouro-1.4B a 1M di contesto scende da 768 GiB a 36 GiB.
Throughput: La "via del percorso di memorizzazione latente" (caching solo i latenti) è limitata dal calcolo (compute-bound) ma consente batch e contesti precedentemente impossibili a causa di OOM (Out of Memory). La "via veloce" (ricostruzione e caching dei full K/V) preserva il throughput ma non realizza il risparmio di memoria.

Significato e Claim

Il paper sostiene che la ricorrenza è l'asse della cache più comprimibile nei Transformer ciclici con pesi condivisi. Il significato risiede in:

Insight Meccanicistico: Identifica che l'asse del ciclo è spettralmente distinto (rango inferiore) rispetto agli assi delle teste e degli strati, permettendo una compressione efficiente senza far collassare il comportamento dinamico del modello in un punto statico.
Deployment Pratico: Fornisce una soluzione post-training che aumenta direttamente la dimensione del batch e le capacità di lunghezza del contesto su hardware esistente senza riaddestrare l'architettura base del modello.
Direzione Futura: I risultati suggeriscono che le architetture ricorsive native potrebbero beneficiare di cache latenti disaccoppiate, dove la ricorrenza calcola un contenuto condiviso più delta economici per ciclo, invece di rimaterializzare interi stati K/V.

Gli autori sottolineano che il metodo non si basa su implementazioni di attenzione assorbita per le sue affermazioni di accuratezza; il risultato centrale è la validità dell'assunzione della traiettoria a basso rango e l'efficacia del codec nel preservare il comportamento del teacher sotto vincoli di memoria rigorosi.