Riepilogo Tecnico: T2MLR: Transformer con Riconcorrenza Temporale negli Strati Medi

Problema

Gli autori identificano un limite fondamentale nelle architetture Transformer standard per quanto riguarda le capacità di ragionamento. Sebbene i Transformer eccellano nel recupero in-context, il loro meccanismo di decodifica autoregressiva crea un collo di bottiglia informativo. Ad ogni passo di decodifica, rappresentazioni latenti ricche e ad alta dimensionalità vengono ripetutamente compresse in uno spazio di token sparso (vettori one-hot) prima di essere immesse come input al passo successivo. Questo processo impedisce agli stati di ragionamento astratto intermedio di persistere nel tempo, costringendo il modello a ricostruire le informazioni necessarie indirettamente tramite l'attenzione o attraverso embedding di input altamente compressi.

Gli approcci esistenti al ragionamento latente tentano di mitigare questo problema operando interamente nello spazio continuo (spesso nell'ultimo strato) o propagando l'incertezza attraverso combinazioni lineari di embedding di token. Tuttavia, questi metodi tipicamente reinseriscono i segnali ricorrenti attraverso l'embedding di input, bypassando gli strati intermedi dove avviene il ragionamento astratto. In alternativa, le architetture Transformer "a loop" aumentano la profondità effettiva iterando su tutti gli strati, ma ciò comporta costi di inferenza che scalano linearmente con il numero di loop, rendendoli poco pratici per la distribuzione.

Metodologia

Il documento introduce T2MLR (Transformers with Temporal Middle-Layer Recurrence), un'architettura innovativa progettata per consentire alle rappresentazioni astratte intermedie di persistere ed evolversi attraverso i passi di decodifica senza aumentare la profondità durante l'inferenza.

Design Architetturale

T2MLR modifica un decoder-only Transformer standard a $L$ strati introducendo un percorso ricorrente a porta (gated) che collega uno strato profondo del token precedente a uno strato precoce del token corrente.

1. Confini della Ricorrenza: L'architettura è parametrizzata da due indici di strato, $\ell_{start}$ e $\ell_{end}$, che definiscono un blocco di strati "medi".
2. Cache Ricorrente ($R_t$): Per ogni passo di decodifica $t$, viene mantenuta una cache di dimensione costante. Essa memorizza la rappresentazione dopo lo strato $\ell_{end}$ del passo precedente.
3. Modulo di Fusione ($\Phi$): Prima di elaborare il token corrente allo strato $\ell_{start}$, il modello fonde la rappresentazione del token corrente ($h^{(\ell_{start}-1)}_t$) con la cache ricorrente ($R_{t-1}$) del token precedente. La fusione è governata da un meccanismo a porta apprendibile:
   $$\Phi(h, R) = h + \tanh(\gamma_{cur}) \sigma(f_{cur}([h, R])) \odot h + \tanh(\gamma_{rec}) \sigma(f_{rec}([h, R])) \odot W_{rec}R$$
   Questo design permette al modello di apprendere quando fare affidamento sullo stream del token corrente rispetto allo stato latente accumulato temporalmente.
4. Aggiornamento della Cache: Dopo il passaggio attraverso lo strato $\ell_{end}$, la cache ricorrente viene aggiornata utilizzando la nuova rappresentazione e la cache precedente tramite una connessione residua e RMSNorm.

Strategia di Addestramento: Parallelismo Temporale

Una sfida chiave per le architetture di ragionamento latente è che l'addestramento standard in parallelo ai token (sequence-parallel) si interrompe quando esistono dipendenze tra i passi temporali. Per affrontare questo, gli autori propongono uno schema di addestramento approssimato in parallelo temporale utilizzando iterazioni a punto fisso di Jacobi:

• Invece dell'elaborazione sequenziale dei token, il modello esegue $d_{forward}$ iterazioni di raffinamento in parallelo.
• Parte da una cache iniziale derivata da un passaggio forward standard, poi fonde iterativamente questa cache con la rappresentazione del token corrente, affina il blocco centrale e aggiorna la cache.
• I gradienti vengono retropropagati attraverso queste iterazioni con una profondità troncata ($d_{backward}$), analogamente alla Truncated Backpropagation Through Time (TBPTT).
• Questo approccio consente di addestrare il modello con dense teacher forcing e parallelismo di sequenza, una capacità che gli autori sostengono sia unica tra le varianti di chain-of-thought continuo.

Contributi Chiave

1. Innovazione Architetturale: T2MLR introduce un percorso ricorrente mirato specificamente negli strati medi del Transformer, piuttosto che nell'embedding di input o nell'output finale. Ciò preserva l'interfaccia autoregressiva standard consentendo al contempo la persistenza dello stato latente.
2. Addestramento Scalabile: Gli autori dimostrano un metodo per addestrare queste dipendenze ricorrenti utilizzando il parallelismo temporale approssimato, consentendo il pre-addestramento e il fine-tuning su grandi dataset senza sacrificare la scalabilità.
3. Capacità di Retrofitting: Il documento mostra che T2MLR non richiede l'addestramento da zero. Il percorso ricorrente può essere inserito in un modello pre-addestrato esistente (es. SmolLM2-1.7B) e sottoposto a fine-tuning per ottenere significativi miglioramenti nel ragionamento.
4. Efficienza: A differenza dei Transformer a loop, T2MLR aggiunge un overhead di inferenza trascurabile (al massimo ~8% per token) perché la ricorrenza è gestita tramite una cache di dimensione costante e un singolo passaggio forward durante l'inferenza.

Risultati Sperimentali

Gli autori valutano T2MLR attraverso benchmark sintetici, pre-addestramento su linguaggio naturale e compiti di ragionamento downstream.

• Ragionamento Sintetico (S5-Retrieval): T2MLM supera significativamente sia i Transformer standard che gli LSTM in un compito che richiede il tracciamento simultaneo dello stato e il recupero in-context. Mentre i Transformer falliscono a lunghezze di sequenza elevate a causa dei vincoli di profondità e gli LSTM falliscono nel recupero a causa della compressione dello stato fisso, T2MLR mantiene un'alta accuratezza.
• Pre-addestramento (FineWeb-Edu): Su modelli da 135M di parametri, le varianti T2MLR eguagliano o superano costantemente i baseline Transformer con parametri equivalenti su benchmark NLP zero-shot (ARC, HellaSwag, ecc.). Notevolmente, la ricorrenza negli strati medi (effettuando il loop su circa il 20% degli strati, es. strati 13–18) spesso supera la ricorrenza a tutti gli strati.
• Ragionamento Downstream: Su compiti come GSM-Aug (matematica), HotPotQA (multi-hop QA) e Variable Assignment, T2MLR mostra guadagni costanti rispetto ai baseline. Le varianti con ricorrenza negli strati medi superano nuovamente le varianti a tutti gli strati, suggerendo che la ricorrenza mirata è più efficace del looping globale.
• Scalabilità: I guadagni di prestazione persistono e crescono scalando a 361M e 1B di parametri ed estendendo il pre-addestramento a 50B di token.
• Retrofitting: Applicare T2MLR a un modello pre-addestrato SmolLM2-1.7B-Instruct e sottoporlo a fine-tuning su dati matematici ha migliorato l'accuratezza di GSM8K dal 35.8% al 39.9% e di MATH500 dal 12.8% al 18.0%, dimostrando l'applicabilità pratica senza un pre-addestramento completo.

Significato e Rivendicazioni

Il documento sostiene che un ragionamento latente efficace nei Transformer non richiede il looping su tutti gli strati (che aumenta il costo di inferenza) né la limitazione della ricorrenza all'embedding di input (che bypassa il calcolo astratto). Inveve, gli autori affermano che la ricorrenza mirata negli strati medi è sufficiente per sbloccare significative capacità di ragionamento.

Il significato di T2MLR risiede nella sua capacità di:

1. Decouplare la Profondità di Ragionamento dal Costo di Inferenza: Introduce un percorso latente temporale che persiste attraverso i passi temporali senza la penalità lineare di inferenza associata alle architetture a loop.
2. Abbassare la Barriera all'Adozione: Dimostrando che l'architettura può essere retrofittata su modelli pre-addestrati esistenti, gli autori sostengono che le avanzate capacità di ragionamento possono essere aggiunte ai modelli attuali a una frazione del costo del pre-addestramento da zero.
3. Validare i Bias Induttivi: I risultati suggeriscono che il "punto ottimale" per la ricorrenza nei Transformer è negli strati medi, dove avviene il calcolo astratto, piuttosto che nelle interfacce di input o output.

Gli autori riconoscono che T2MLR comporta un overhead in termini di tempo di addestramento (2–4x il tempo di calcolo wall-clock a seconda della dimensione del blocco ricorrente) dovuto ai passi di approssimazione di Jacobi. Lo inquadrano come un compromesso: T2MLR sacrifica il calcolo di addestramento per ottenere un'architettura di inferenza ad alte prestazioni e basso overhead per i compiti di ragionamento.