ReLaX: Reasoning with Latent Exploration for Large Reasoning Models

1. Il Problema: Collasso dell'Entropia e Esplorazione Inefficace

Il paper affronta una limitazione fondamentale nell'addestramento di Large Reasoning Models (LRM) tramite Reinforcement Learning with Verifiable Rewards (RLVR).

• Il Fenomeno: Sebbene l'RLVR abbia dimostrato grandi potenzialità nel migliorare il ragionamento, tende a spingere la politica del modello verso un eccessivo determinismo. Questo porta a un "collasso dell'entropia", dove la distribuzione della politica si concentra su un sottospazio ristretto, causando una convergenza prematura e un'esplorazione inefficace.
• Il Limite delle Soluzioni Attuali: Le metodologie esistenti tentano di mitigare questo problema aumentando l'entropia a livello di token (es. regolarizzazione basata sull'entropia, reshaping del reward). Tuttavia, l'obiettivo di mantenere un'alta entropia a livello di token è intrinsecamente in conflitto con la natura dell'RL, che favorisce politiche deterministiche. Inoltre, per i modelli multimodali (MLLM), il feedback a livello di token non riesce a riflettere accuratamente le complesse computazioni interne cross-modali, rendendo queste soluzioni inefficienti.
• L'Iperbole: Il collasso dell'entropia è visto come un sintomo superficiale di una patologia più profonda: le dinamiche latenti interne che governano la generazione dei token perdono flessibilità e convergono in pattern rigidi.

2. Metodologia: ReLaX e Dinamiche Latenti

Gli autori propongono ReLaX (Reasoning with Latent eXploration), un framework che sposta il focus dall'esplorazione nello spazio dei token all'esplorazione nello spazio latente continuo delle rappresentazioni interne del modello.

A. Teoria dell'Operatore di Koopman

Per rendere analizzabili e intervenibili le dinamiche non lineari e ad alta dimensionalità degli stati nascosti (hidden states) degli LRM, il paper utilizza la Teoria dell'Operatore di Koopman.

• Questo approccio permette di rappresentare l'evoluzione non lineare del sistema in uno spazio infinito-dimensionale di osservabili, dove l'evoluzione è governata da un operatore lineare ($K$).
• Viene utilizzato un ResKoopNet (estensione di ResDMD) per apprendere un dizionario neurale che mappa gli stati nascosti non lineari in una rappresentazione lineare stabile e accurata.

B. Metrica Chiave: DSD (Dynamic Spectral Dispersion)

Sulla base della rappresentazione lineare ottenuta, gli autori introducono una nuova metrica chiamata Dynamic Spectral Dispersion (DSD).

• Definizione: Il DSD è calcolato come la varianza dei moduli degli autovalori dell'operatore di Koopman ($K$) derivati dagli stati nascosti di una traiettoria.
• Significato: Un DSD alto indica una ricca diversità di modi dinamici (spettro disperso), suggerendo che il modello mantiene una flessibilità computazionale interna e non è caduto in pattern ripetitivi. Un DSD basso indica un collasso in comportamenti rigidi e degenerati.
• Vantaggio: A differenza dell'entropia dei token, il DSD cattura l'eterogeneità dei processi computazionali interni, fornendo un indicatore più fondamentale dell'esplorazione della politica.

C. Funzione Obiettivo e Regularizzazione

ReLaX integra il DSD nell'obiettivo di ottimizzazione della politica (basato su GRPO, Group Relative Policy Optimization):

1. Regularizzazione basata su DSD: Viene aggiunto un termine di regolarizzazione ($\mathcal{L}_{xp}$) che penalizza la convergenza verso dinamiche latenti omogenee (basso DSD).
2. Adaptive KL & Advantage Shaping: Per evitare un'esplorazione eccessiva e caotica, il metodo applica due controlli:
   • Advantage Shaping: La regolarizzazione DSD è ponderata solo per le traiettorie con vantaggi positivi (reward alto), evitando di esplorare direzioni dannose.
   • KL Adattivo: Viene applicata una penalità KL solo alle traiettorie che mostrano una dispersione dinamica eccessiva (DSD troppo alto), stabilizzando l'addestramento senza bloccare l'esplorazione utile.

L'obiettivo finale è:
$$J(\theta) = J_{GRPO}(\theta) + \alpha \tilde{\mathcal{L}}_{xp} + \beta \sum_{i \in I} D_{KL}(\pi_\theta(o_i) \parallel \pi_{ref}(o_i))$$

3. Contributi Chiave

1. Nuova Metrica (DSD): Introduzione di un indicatore che quantifica l'eterogeneità delle dinamiche latenti di ragionamento, offrendo una caratterizzazione più profonda dell'esplorazione rispetto alle statistiche superficiali dei token.
2. Framework ReLaX: Sviluppo di un nuovo paradigma di addestramento RLVR che utilizza la regolarizzazione basata sul DSD per bilanciare esplorazione e sfruttamento, mitigando la saturazione delle prestazioni.
3. Teoria Applicata: Applicazione innovativa della teoria dell'operatore di Koopman per analizzare e controllare le dinamiche interne dei Large Reasoning Models.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli LLM (solo testo) e MLLM (multimodali) di diverse dimensioni (3B e 7B).

• Ragionamento Multimodale (VLM):

  • Su 7 benchmark multimodali (inclusi MMMU, MathVista, EMMA), ReLaX-VL-7B ha raggiunto un punteggio medio di 53.2, stabilendo un nuovo stato dell'arte (SOTA), superando il precedente miglior modello (VL-Rethinker-7B con 52.5).
  • Il modello 3B di ReLaX ha superato diversi modelli precedenti di livello 7B.
  • L'analisi delle dinamiche di addestramento mostra che ReLaX mantiene un'entropia e un DSD più elevati e stabili rispetto al GRPO standard, evitando il collasso prematuro.

• Ragionamento Matematico (Solo Testo):

  • Su 6 benchmark matematici (MATH500, AMC, AIME), ReLaX ha superato tutti i metodi basati su token-level (inclusi FR3E, DAPO, KL-Cov).
  • Su Qwen2.5-7B-Math, ReLaX ha migliorato il punteggio medio del 49.1% rispetto al baseline, superando il precedente SOTA (FR3E) di un margine significativo (+6.3 punti).

• Analisi Qualitativa:

  • I casi studio mostrano che i modelli addestrati con ReLaX eseguono verifiche di ragionamento più robuste e pertinenti (es. applicazione corretta di leggi fisiche/matematiche) rispetto ai modelli che usano regolarizzazione a livello di token, che tendono a generare codice di verifica inutile o allucinazioni.
  • ReLaX dimostra una maggiore robustezza nelle varianti visive dei problemi (DynaMath), mantenendo la coerenza dove altri modelli falliscono.

5. Significato e Impatto

Il lavoro di ReLaX segna un cambio di paradigma significativo nel campo dell'RLVR per i modelli di ragionamento:

• Oltre lo Spazio dei Token: Dimostra che l'intervento diretto sulle dinamiche latenti interne è più efficace e fondamentale rispetto alla manipolazione delle statistiche dei token.
• Scalabilità e Generalizzazione: Il metodo funziona efficacemente sia su modelli puramente testuali che multimodali, suggerendo che la flessibilità computazionale interna è un prerequisito universale per un ragionamento avanzato.
• Soluzione al Trade-off: Offre una soluzione più "principiata" al dilemma esplorazione-sfruttamento, utilizzando la teoria dei sistemi dinamici per guidare l'esplorazione in modo controllato e significativo.

In sintesi, ReLaX apre una nuova direzione per migliorare le capacità di ragionamento degli LLM, spostando l'attenzione dalla superficie (token) alla struttura profonda (dinamiche latenti) del processo di pensiero del modello.