Test-time training end-to-end: comprimere il contesto lungo nei pesi del modello

Un gruppo di Stanford, NVIDIA e Berkeley riformula il contesto lungo come apprendimento continuo: invece di ricordare ogni token, il modello continua ad allenarsi mentre legge e comprime il testo nei propri pesi. Guadagni reali su costo e loss, un limite netto sul recupero esatto.

Ripensate alla prima lezione di un corso seguito anni fa. Non ricordate la frase con cui il docente aprì, né l’ordine esatto delle slide. Eppure qualcosa è rimasto: un’intuizione, un modo di ragionare che oggi usate ancora. Il cervello non ha archiviato la trascrizione integrale di quelle ore; ha compresso l’esperienza, buttando via i dettagli e trattenendo ciò che serviva. È un modo economico di ricordare, ed è esattamente il contrario di come funziona un Transformer con la sua attention. Un paper firmato da ricercatori di Astera Institute, NVIDIA, Stanford, Berkeley e UC San Diego — End-to-End Test-Time Training for Long Context, dicembre 2025 — prova a insegnare ai modelli linguistici proprio questa forma di memoria per compressione.

Il costo di ricordare tutto

Per capire la proposta serve sapere cosa scardina. Un Transformer con full attention tiene in memoria le key e i value di ogni token già letto: per predire la parola successiva li riscorre tutti. È un archivio quasi fotografico, con un difetto strutturale. Il costo per token cresce linearmente con la lunghezza del contesto — in termini di complessità, $O(T^2)$ per il prefill e $O(T)$ per generare ogni token — e su sequenze di centinaia di migliaia di token diventa proibitivo. Le alternative note sono le reti ricorrenti come Mamba 2 e Gated DeltaNet: costo per token costante, ma efficacia che cala man mano che il contesto si allunga. Finestre scorrevoli e architetture ibride mitigano il problema senza risolverlo: restano meno efficaci della full attention nello sfruttare il testo lungo.

La domanda del paper è netta: si può ottenere un modello che migliora davvero con più contesto, ma con costo per token costante e senza dover ricordare ogni dettaglio? La risposta ruota attorno a una parola: compressione. Gli autori notano che addestrare un modello con next-token prediction è già, di per sé, un atto di compressione: enormi quantità di dati finiscono schiacciate dentro i pesi. E se allora continuassimo semplicemente ad addestrare il modello anche a runtime, sul contesto specifico che ci viene dato?

Allenarsi mentre si legge

Questa è l’idea del Test-Time Training (TTT). Mentre legge il contesto, il modello lo usa come materiale di allenamento: per ogni token prova a indovinare il successivo, misura quanto ha sbagliato e aggiusta i propri pesi di conseguenza. La perdita è la solita cross-entropy della predizione,

$$\ell_t(W) = \mathrm{CE}\big(f(x_{t-1};\, W),\ x_t\big),$$

e a ogni passo il modello fa un piccolo passo di discesa del gradiente su quella perdita:

$$W_t = W_{t-1} – \eta\, \nabla \ell_t(W_{t-1}).$$

L’immagine è quella dello studente che, prima di affrontare l’esercizio nuovo, ripassa risolvendo quelli già svolti: dopo aver «esercitato» i pesi sui token letti, l’informazione di quel contesto resta impressa nei pesi aggiornati. Il modello non conserva una cache che si gonfia; conserva un unico stato di dimensione fissa. Da qui il vantaggio: costo per token costante, $O(T)$ per il prefill e $O(1)$ per la decodifica, come una RNN.

Il pezzo mancante: prepararsi a imparare

C’è una crepa in questo schema, ed è il vero contributo del lavoro. L’idea di aggiornare i pesi a runtime è vecchia (si chiamava dynamic evaluation). Il problema è che il modello, in fase di addestramento, viene ottimizzato per la sua perdita «a freddo», non per la perdita dopo che sarà stato aggiornato al test. Le due fasi non parlano la stessa lingua. Gli autori chiudono questa distanza con il meta-learning: durante l’addestramento ogni sequenza viene trattata come se fosse già una sequenza di test, si esegue TTT su di essa in un «ciclo interno», e poi si ottimizza la perdita media a valle del TTT rispetto all’inizializzazione, in un «ciclo esterno» che richiede gradienti di gradienti. La differenza è tutta qui:

$$L(W_0; X) = \frac{1}{T}\sum_{t=1}^{T} \ell_t(W_{t-1}) \qquad\text{contro}\qquad L_{\text{naive}}(W_0; X) = \frac{1}{T}\sum_{t=1}^{T} \ell_t(W_0).$$

A sinistra la versione end-to-end: ogni termine misura la perdita con i pesi già aggiornati dal test-time training, cioè con la $W_{t-1}$ che il TTT produrrà davvero. A destra la versione ingenua, che finge che i pesi restino l’inizializzazione $W_0$. È una differenza sottile e decisiva: nel loro esperimento-giocattolo, la variante ingenua migliora appena rispetto a un modello senza memoria, mentre la variante end-to-end arriva quasi al livello della full attention. Il metodo si chiama TTT-E2E perché è end-to-end due volte: nel ciclo interno ottimizza direttamente la next-token prediction (E2E al test), nel ciclo esterno ottimizza la perdita finale dopo il TTT (E2E in addestramento).

Nella pratica servono alcuni accorgimenti. L’aggiornamento non avviene token per token ma a mini-batch (per la sequenza da 128K, finestra di attention $k = 8\text{K}$ e mini-batch $b = 1\text{K}$); si aggiornano solo i livelli MLP e solo l’ultimo quarto dei blocchi, per contenere il costo; e a quei blocchi si aggiunge un secondo MLP statico che fa da «deposito sicuro» per la conoscenza acquisita nel pre-training. L’architettura, per il resto, resta un normale Transformer: il paper insiste che il design architetturale ha un ruolo minore e di supporto, il lavoro lo fa la formulazione come apprendimento continuo.

Contesto lungo: ricordare tutto o comprimere? Full attention token letti cache key + value di OGNI token cresce con il contesto costo per token cresce con T prefill O(T·T) · decode O(T) TTT-E2E token letti pesi del modello aggiornati mentre legge il contesto dimensione fissa costo per token costante prefill O(T) · decode O(1) A 128K token: latenza costante, 2,7× più veloce di full attention (H100)

Due modi di gestire il contesto lungo. La full attention conserva key e value di ogni token in una cache che cresce, con costo per token crescente. TTT-E2E comprime il contesto nei pesi, di dimensione fissa: il costo resta costante e a 128K token la latenza di inferenza è 2,7 volte più bassa.

I numeri, senza abbellirli

Il paper misura soprattutto le proprietà di scala. Con modelli da 3 miliardi di parametri addestrati su 164 miliardi di token, TTT-E2E scala con la lunghezza del contesto nello stesso modo della full attention, mentre Mamba 2 e Gated DeltaNet peggiorano progressivamente al crescere del contesto. Non solo: nel breakdown per posizione, TTT-E2E è l’unico metodo che ottiene sempre una perdita più bassa della full attention lungo tutta la finestra, con un vantaggio che si concentra sui token iniziali. E come le RNN, ha latenza di inferenza costante indipendente dalla lunghezza: a 128K token risulta 2,7 volte più veloce della full attention su una GPU H100. Interessante anche un dettaglio delle ablazioni: applicare TTT-E2E sopra la full attention migliora comunque la perdita di test di 0,018, segno che l’effetto è ortogonale, non un semplice rattoppo di una finestra ridotta.

Poi arriva la doccia fredda, riportata dagli stessi autori. Sul Needle in a Haystack — il test che chiede di ritrovare una stringa-bersaglio nascosta in un testo lungo — la full attention stravince su tutti gli altri metodi, TTT-E2E compreso, e il divario si allarga con il contesto: nel compito di recupero di una pass-key a 128K token, la full attention azzecca 0,99 dei casi contro 0,06 di TTT-E2E. Non è un bug, è il rovescio della medaglia della compressione: buttare via i dettagli «irrilevanti» significa buttare via anche l’ago nel pagliaio. La forza della self-attention, ammette il paper, sta proprio nel suo recupero quasi senza perdite.

I limiti dichiarati

Oltre al recupero esatto, gli autori mettono in chiaro un secondo problema, questa volta di ingegneria: la latenza di addestramento. Calcolare gradienti di gradienti è una procedura molto meno ottimizzata dell’addestramento di un Transformer normale. Nella loro implementazione TTT-E2E è 1,2 volte più veloce della full attention a 128K, ma 3,4 volte più lento a 8K — e siccome quasi tutto il budget di addestramento si spende sul contesto corto, questa resta una limitazione seria. Indicano due vie d’uscita, lasciate al lavoro futuro: un kernel di attention su misura (l’attuale FlashAttention non supporta i gradienti di gradienti) e l’inizializzazione da un Transformer già pre-addestrato, che ridurrebbe il TTT a una piccola frazione del calcolo complessivo. Vale la pena aggiungere che i risultati riguardano modelli base, senza instruction tuning né reinforcement learning, e a una scala «di ricerca»: gli autori sono espliciti nel dire che alcuni dettagli implementativi potrebbero essere artefatti del loro setup.

Perché guardarlo

La mossa concettuale conta più del singolo numero. Per anni il contesto lungo è stato affrontato come un problema di architettura: rendere l’attention più economica, impilare strati ricorrenti, inventare finestre. Qui invece diventa un problema di apprendimento continuo, risolto con un Transformer quasi standard che continua a studiare mentre lavora. Ne esce una gerarchia che ricorda la memoria biologica, come nota la conclusione: i pesi aggiornati al test sono la memoria a lungo termine, la finestra scorrevole è quella a breve termine. Le due si completano, e non si escludono. La posta in gioco è capire se il futuro dei modelli con contesti enormi somiglierà più a un archivio che conserva ogni parola o a una mente che, di quel corso seguito anni fa, ha tenuto solo ciò che serviva.

Il paper

Arnuv Tandon, Karan Dalal, Xinhao Li, Daniel Koceja, Marcel Rød, Yu Sun e altri, End-to-End Test-Time Training for Long Context, arXiv:2512.23675 (dicembre 2025). Codice pubblico: github.com/test-time-training/e2e.

I commenti sono riservati agli iscritti.

Accedi per commentare