Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers

Riepilogo Tecnico: Convergenza a Punto Fisso Per-Token nei Transformer a Profondità Ricorrente

Definizione del Problema
I transformer standard a profondità fissa vincolano la qualità del modello a un singolo punto di calcolo operativo. Servire diversi livelli di qualità/costo richiede l'addestramento e l'archiviazione di modelli separati, mentre l'adattamento del calcolo alla difficoltà dell'input durante l'inferenza sacrifica la cache del prompt. I transformer a profondità ricorrente (che applicano un nucleo con pesi condivisi un numero variabile di volte) offrono una soluzione teorica, permettendo a un singolo set di parametri di operare su un intervallo di profondità. Le opere precedenti spesso confondono due proprietà distinte: la robustezza alla profondità (produrre output utili ad altezze non addestrate) e l'allocazione della profondità (determinare quanta profondità richiede ogni specifico token). Gli approcci esistenti all'allocazione per-token soffrono di limitazioni significative: i metodi di early-exit su stack non condivisi falliscono a causa delle inconsistenze della cache key-value; i metodi di mixture-of-recursions fissano il budget di calcolo nei pesi durante il pretraining, impedendo il retargeting post-hoc; e le policy di halting addestrate congiuntamente possono collassare in profondità degenerate o richiedere una calibrazione delicata. Crucialmente, nessun lavoro precedente ha misurato direttamente il calcolo per-token di un modello elastico per determinare se una semplice policy priva di addestramento potesse superare quelle apprese.

Metodologia
Gli autori propongono un approccio in due fasi che separa la robustezza alla profondità dall'allocazione della profondità:

1. Robustezza alla Profondità (Pretraining): Un backbone a profondità randomizzata viene addestrato usando una ricetta minima. Ad ogni step dell'optimizer, viene campionato un conteggio di ricorsione $r$ da una distribuzione log-normale e applicato all'intero step. Il modello utilizza la standard cross-entropy loss senza unità di halting, ponder loss o penalità FLOP. L'architettura consiste in un preludio unico, un nucleo con pesi condivisi (con reiniezione dell'input e sandwich normalization per prevenire il collasso dello stato) e una coda unica.
2. Allocazione della Profondità (Policy Post-hoc): Una volta congelato il backbone, vengono valutate tre policy sullo stesso modello per determinare la profondità per-token:
   • A0 (Profondità Uniforme): Applica un numero di loop $r$ fisso a tutti i token (baseline).
   • A1 (Uscita per Convergenza): Una regola priva di addestramento che interrompe un token una volta che la divergenza KL tra le distribuzioni di output successive ($KL(p_i || p_{i-1})$) scende al di sotto di una soglia $\epsilon$.
   • A2 (Router con Etichette Oracle): Un liner probe addestrato sugli stati nascosti del modello stesso per predire se l'argmax di un token cambierà in loop futuri.

Contributi Chiave

1. Misurazione Diretta Per-Token: Lo studio fornisce la prima caratterizzazione della convergenza per-token in un modello linguistico a profondità ricorrente. Dimostra che lo stato ricorrente converge a un punto fisso per-token. Mentre il token mediano converge al sesto loop, una coda di circa il 10% dei token continua ad aggiornarsi alla profondità media di addestramento di otto. La profondità di convergenza è ordinata per tipo di token (gli spazi bianchi convergono più velocemente, le parole di contenuto più profondamente).
2. Leggere la Convergenza vs. Impararla: Il paper dimostra che leggere direttamente la dinamica dello stato del modello supera l'apprendimento della predizione delle stesse. La uscita per convergenza priva di addestramento (A1) raggiunge la qualità di una profondità uniforme di 8 con una media di 4,94 loop (una riduzione del 38% della profondità media). Al contrario, il router lineare (A2), addestrato sulle etichette di convergenza raccolte dallo stesso modello, richiede quasi l'intera profondità (7,99 loop) e non produce risparmi.
3. Ipotesi di Saturazione della Media di Addestramento: Supportati da dati interni e controlli di coerenza esterni, gli autori ipotizzano che la profondità utile di un modello ricorrente saturi alla sua distribuzione media di loop durante l'addestramento. I modelli addestrati con una media di 4 saturano alla profondità 4; quelli con una media di 8 saturano alla profondità 8.

Risultati

• Elasticità: Il backbone a profondità randomizzata mostra un miglioramento monotono della qualità dalla profondità 1 a 8, dopo la quale la loss di validazione si stabilizza (3,20 nats) anche fino a 32 loop. Ciò conferma che il modello è robusto rispetto alla profondità senza richiedere una calibrazione specifica per ogni profondità.
• Dinamiche di Convergenza: La divergenza KL media tra successivi output scende da $3,9 \times 10^{-1}$ al loop 2 a $8,5 \times 10^{-6}$ al loop 16. Il cambiamento di stato decade per step, confermando che il sistema raggiunge un punto fisso piuttosto che girare inutilmente.
• Efficienza di Allocazione: La policy A1 priva di addestramento eguaglia la qualità di un modello a profondità uniforme di 8 con 4,94 loop medi. Supera la baseline uniforme alle profondità ridotte e eguaglia la stessa nell'intervallo di profondità media. Il router appreso (A2) non riesce a fornire risparmi, richiedendo quasi l'intera profondità per eguagliare la baseline.
• Comportamento a Livello di Token: I token strutturali (spazi bianchi) escono prima, mentre i token di contenuto computano più a lungo. L'ordine relativo di numeri e parole si inverte tra i dataset (TinyStories vs. FineWeb-Edu), attribuito a shift di distribuzione piuttosto che a fallimenti architettonici.
• Limiti di Efficienza: Sebbene lo studio riporti una riduzione del 38% della profondità media (un proxy dei FLOP), nota esplicitamente che il miglioramento reale del tempo di esecuzione dipende dal batching continuo per profondità, che non è stato implementato. Il miglior caso di velocità wall-clock è stimato tra 1,48–1,51x, mentre l'attuale harness mostra uno 0% di velocità wall-clock realizzata a causa del masking invece che dello skipping del lavoro.

Significato e Rivendicazioni
Il paper sostiene che l' "elasticità" dei transformer a profondità ricorrente è una proprietà stabile che può essere sfruttata tramite semplici policy prive di addestramento. La scoperta centrale è che la variazione per-token nella profondità di convergenza è direttamente leggibile e sfruttabile, rendendo superflue e meno efficaci le complesse policy di halting apprese. Gli autori posizionano questo lavoro come uno studio di misurazione e allocazione piuttosto che un contributo al backbone elastico sottostante (che attribuiscono a lavori precedenti come Huginn e LoopFormer).

Lo studio ha un ambito modesto:

• I risultati sono stabiliti a una singola scala (parametri della classe 135M) e un singolo seed (per il modello principale).
• Non viene rivendicata la parità di FLOP rispetto ai transformer standard; i confronti sono fatti contro un riferimento ceiling (SmolLM2-135M) che ha un vantaggio di 150x nel conteggio dei token.
• La velocità wall-clock realizzata non è riportata come metrica finale, ma come un intervallo tra la riduzione dei FLOP e il potenziale del miglior caso di batching.
• L'ipotesi che la profondità di saturazione sia uguale alla profondità media di addestramento è presentata come un'ipotesi aperta supportata da dati limitati, non come una legge dimostrata.

Lo studio completo è stato riproducibile su una singola RTX 4090 in circa 100 ore GPU.