Latent Reasoning in TRMs is Secretly a Policy Improvement Operator

Sintesi Tecnica: Il Ragionamento Latente nei TRM è Segretamente un Operatore di Miglioramento della Policy

1. Definizione del Problema

Recenti modelli piccoli che utilizzano la ricorsione latente, come i Tiny Recursive Models (TRM) e gli Hierarchical Reasoning Models (HRM), hanno dimostrato risultati promettenti su compiti di ragionamento complesso come ARC-AGI. La teoria prevalente suggerisce che questi modelli ottengano guadagni di prestazione aumentando la profondità effettiva della rete, permettendo loro di emulare modelli più grandi con meno parametri. Tuttavia, le osservazioni empiriche rivelano una discrepanza: i layer aggiunti ricorsivamente spesso performano peggio rispetto ai modelli a passaggio singolo di equivalente profondità feed-forward. Ciò indica che, nelle architetture a loop, non ogni passo ricorsivo contribuisce significativamente alla soluzione, risultando in "computazione morta" (dead compute).

Le domande centrali affrontate da questo lavoro sono:

• Quando e perché il ragionamento latente migliora le prestazioni?
• Quando risulta in una computazione ridondante?
• Come può essere migliorato l'addestramento di questi modelli per garantire che ogni passo sia efficace?

I metodi di addestramento attuali per i TRM applicano la supervisione principalmente all'output finale. Ciò lascia i primi passi di raffinamento con un debole assegnamento del credito (credit assignment), rendendo l'ottimizzazione fragile, specialmente in regimi con pochi dati.

2. Metodologia

Gli autori propongono un inquadramento teorico del ragionamento latente e una corrispondente metodologia di addestramento.

Formalizzazione Teorica: Il Ragionamento Latente come Miglioramento della Policy
Il paper formalizza il ragionamento ricorsivo latente come un operatore di miglioramento della policy.

• Policy di Riferimento ($\hat{\pi}_t$): La distribuzione dell'output del modello prima di un passo di ragionamento latente (la "prima ipotesi").
• Policy Migliorata ($\pi^+_t$): La distribuzione dell'output dopo un aggiornamento del ragionamento latente.
• Segnale di Advantage: La differenza tra i log-probabilità della policy migliorata e della policy di riferimento funge da segnale di advantage implicito ($A_t(s, a) = \log \pi^+_t(a|s) - \log \hat{\pi}_t(a|s)$).
• Meccanismo: Gli autori dimostrano che un singolo passo TRM implementa un aggiornamento della policy che ri-pesa la policy di riferimento basandosi su una condizione di ottimalità. Nello specifico, il passo di ragionamento latente si comporta come una policy condizionata all'evento che la risposta venga migliorata. Questo collega l'algoritmo nello spazio latente al noto framework del miglioramento della policy nel Reinforcement Learning (RL) e del Classifier-Free Guidance (CFG) nei modelli di diffusione.

Schema di Addestramento Proposto: Deep Improvement Supervision (DIS)
Basandosi sull'intuizione del miglioramento della policy, gli autori introducono la Deep Improvement Supervision (DIS) per eliminare i passi di computazione morta.

• Target per Passo: Invece di supervisionare solo l'output finale, la DIS fornisce target intermedi $\{y^\dagger_s\}$ per ogni passo di ricorsione $s$. Questi target sono generati tramite uno schema di corruzione discreta monotona dell'output reale $y^*$.
• Obiettivo di Addestramento: Il modello viene addestrato a mappare il target precedente $y^\dagger_{s-1}$ (tramite la policy di riferimento) al successivo target meno corrotto $y^\dagger_s$ (tramite la policy migliorata).
• Funzione di Perdita: L'addestramento minimizza una perdita di cross-entropy tra i logit post-ragionamento e il target specifico per il passo $y^\dagger_s$. Ciò forza il rapporto logaritmico dell'advantage ad ogni passo a essere positivo per il token corretto, garantendo che ogni passo di ricorsione esegua un miglioramento diretto piuttosto che una computazione ridondante.

Semplificazioni Architetturali
Il metodo proposto semplifica l'architettura TRM:

• Rimuovendo la necessità di un meccanismo di halting appreso (Adaptive Computation Time).
• Riducendo il numero di passi di supervisione e di passi di ragionamento latente.
• Utilizzando un numero fisso di passi di ricorsione ($N_{sup}$) determinato dallo schema di corruzione.

3. Contributi Chiave

1. Intuizione Teorica: Il paper stabilisce che il ragionamento latente nei TRM è matematicamente equivalente a un operatore di miglioramento della policy. Esso colma il divario tra gli algoritmi di ragionamento nello spazio latente e i framework di RL, spiegando il comportamento del modello attraverso la lente della massimizzazione dell'advantage.
2. Deep Improvement Supervision (DIS): Un nuovo schema di addestramento che fornisce target intermedi per passo. Questo trasforma il difficile problema dell'assegnazione del credito a lungo termine in un compito di apprendimento supervisionato trattabile, dove ogni passo è esplicitamente addestrato a migliorare rispetto allo stato precedente.
3. Guadagni di Efficienza: Il metodo ottiene prestazioni competitive riducendo significativamente l'overhead computazionale. Evita l'addestramento di moduli di halting e riduce il numero totale di passaggi forward di un fattore di 18× rispetto alle configurazioni standard dei TRM, mantenendo o migliorando l'accuratezza.

4. Risultati Sperimentali

Gli autori hanno valutato il loro metodo sul problema N-Queens e sui benchmark ARC-AGI (ARC-AGI-1 e ARC-AGI-2).

• N-Queens: Utilizzando un modello compatto (0.8M di parametri), il metodo DIS ha raggiunto un'accuratezza di 0.69, eguagliando le prestazioni di un TRM standard con una testa di halting e 16 passi di supervisione, ma utilizzando significativamente meno passi di inferenza ($T=1, n=2$ vs. $T=3, n=6$).
• ARC-AGI:
  • Modello Compatto (0.8M di parametri): La DIS ha raggiunto il 24.0% di accuratezza su ARC-AGI-1, superando l'originale TRM-compact (12.0%) e la maggior parte degli esistenti LLM open-source senza conoscenza esterna.
  • Modello Medio (7M di parametri): La DIS ha raggiunto il 41.3% su ARC-AGI-1 e il 6.0% su ARC-AGI-2, superando il TRM-medium standard (40.4% / 3.3%) e pareggiando le prestazioni del TRM originale (40.4% / 3.3%) pur utilizzando meno passi di ragionamento latente.
• Riduzione della Computazione: Il metodo ha ridotto il numero totale di passaggi forward di un fattore 18 rispetto alla configurazione baseline del TRM, dimostrando che la supervisione esplicita degli step intermedi elimina la necessità di eccessivi loop di ragionamento latente.

5. Significato e Rivendicazioni

Il paper sostiene che, vedendo il ragionamento latente come un processo di miglioramento della policy, sia possibile spiegare perché i modelli ricorsivi standard a volte falliscono (a causa della "computazione morta" dove i passi non migliorano il margine di advantage) e come risolvere il problema.

La principale importanza risiede nel framework di Deep Improvement Supervision, che:

• Fornisce una regola di addestramento costruttiva per garantire che ogni passo di ricorsione esegua un miglioramento non banale.
• Semplifica l'addestramento dei modelli ricorsivi eliminando la necessità di complessi meccanismi di halting e di un fragile assegnamento del credito.
• Dimostra che modelli piccoli ed efficienti nei parametri possono raggiungere risultati allo stato dell'arte su compiti di ragionamento complesso (come ARC-AGI) quando l'obiettivo di addestramento è allineato con il meccanismo teorico del miglioramento della policy.

Gli autori concludono che il loro approccio offre una via per modelli di ragionamento più efficienti ed efficaci, allineando il segnale di addestramento con la natura algoritmica intrinseca della ricorsione latente.