Scaling Reasoning Efficiently via Relaxed On-Policy Distillation

1. Il Problema

L'articolo affronta la sfida di trasferire capacità di ragionamento avanzate dai grandi modelli linguistici (LLM, "teacher") a modelli più piccoli e con risorse limitate (SLM, "student").

• Limiti del RL diretto: I piccoli modelli faticano ad ottimizzare direttamente tramite Reinforcement Learning (RL) a causa della loro capacità rappresentativa limitata e della variabilità dei reward, rendendo l'approccio RL standard inefficace.
• Instabilità della Distillazione On-Policy: Le tecniche attuali di distillazione on-policy (dove lo studente impara dai propri percorsi generati sotto la guida del teacher) soffrono di gravi instabilità. In particolare, tendono a causare:
  • Trasferimento negativo: Il modello studente può degradare rispetto alla sua versione iniziale.
  • Collasso dell'entropia: Il modello converge prematuramente su un insieme ristretto di output, perdendo la capacità di esplorare diverse soluzioni.
  • Reward instabili: La distribuzione dei reward basata sul rapporto di verosimiglianza (log-likelihood ratio) presenta code pesanti negative (quando il teacher assegna probabilità quasi zero a un token scelto dallo studente) e valori vicini allo zero (rumore), che destabilizzano l'ottimizzazione.

2. Metodologia: REOPOLD

Gli autori propongono REOPOLD (Relaxed On-Policy Distillation), un framework che interpreta la distillazione on-policy come un problema di ottimizzazione della politica (Policy Optimization), trattando il rapporto di verosimiglianza teacher-studente come un reward. Per stabilizzare l'addestramento, REOPOLD introduce tre componenti chiave:

A. Equivalenza Teorica e Stop-Gradient

Il paper dimostra teoricamente che la distillazione on-policy è equivalente a un metodo di gradiente della politica. Per stabilizzare l'ottimizzazione, viene applicato un operatore stop-gradient sul termine di reward (il log-rapporto). Questo trasforma il reward in un segnale estrinseco fisso, agendo come un "control variate" che riduce la varianza del gradiente e previene l'instabilità causata dalla dipendenza del reward dai parametri del modello.

B. Clipping del Reward basato su Mixture

Per affrontare le code pesanti negative (dove il teacher assegna probabilità ~0), REOPOLD introduce un clipping del reward derivato da un'analisi di distribuzioni miste.

• Invece di permettere al reward di divergere a $-\infty$, viene impostato un limite inferiore teorico: $\log(\frac{\lambda}{1-\lambda})$.
• Questo "pavimento" (floor) previene aggiornamenti distruttivi dei parametri, filtrando i segnali dannosi senza alterare globalmente la funzione obiettivo come fanno altri metodi (es. Skew RKL).

C. Campionamento Dinamico a Livello di Token (Guidato dall'Entropia)

Il paper osserva che la maggior parte dei token ha un basso rapporto di verosimiglianza (reward vicino a zero) perché le distribuzioni di teacher e studente sono già allineate.

• REOPOLD utilizza l'entropia del token come proxy per la densità informativa.
• Vengono selezionati e aggiornati solo i token ad alta entropia (punti di incertezza critica dove le distribuzioni divergono), filtrando il rumore dei token deterministici. Questo migliora drasticamente l'efficienza del campione.

D. Strategia di Addestramento Multi-Stadio (Esplorazione-Raffinamento)

L'addestramento è diviso in due fasi dinamiche:

1. Fase di Esplorazione: Si applica un mascheramento basato sul reward per eliminare le penalità eccessive negative. Questo permette allo studente di esplorare diverse soluzioni plausibili (simile al SFT) senza collassare prematuramente.
2. Fase di Raffinamento: Si passa a un mascheramento basato sull'entropia, permettendo feedback negativi sui token critici e incerti per consolidare i percorsi di ragionamento corretti (simile al RL).

3. Contributi Chiave

1. Diagnosi dell'Instabilità: Identificazione delle code pesanti negative e dell'inefficienza dei reward vicini allo zero come cause principali del fallimento della distillazione on-policy standard.
2. Framework Unificato: Introduzione di REOPOLD, che integra clipping del reward, campionamento dinamico e addestramento multi-stadio in un unico obiettivo ottimizzabile.
3. Efficienza e Scalabilità: Dimostrazione empirica che rilassare i vincoli di imitazione rigida è essenziale per scalare le capacità di ragionamento nei modelli compatti.

4. Risultati Sperimentali

I risultati sono stati valutati su compiti di ragionamento matematico, visivo e agenziale (uso di strumenti).

• Efficienza del Campione (Sample Efficiency):
  • Su benchmark matematici (AIME-25), REOPOLD raggiunge prestazioni superiori rispetto agli approcci RL (come ProRL) con 6.7x - 12x meno campioni di addestramento.
  • Supera i baselines RKL (Relaxed KL) e SFT in tempi molto più brevi.
• Scalabilità al Test-Time (Test-Time Scaling):
  • Nei compiti di ragionamento visivo (Geometry3K, MathVerse), REOPOLD permette a modelli piccoli (es. 7B) di avvicinarsi alle prestazioni di un teacher da 32B.
  • Si ottiene un speedup nell'inferenza di ~3.3x rispetto al teacher, mantenendo un'accuratezza superiore o paragonabile.
• Robustezza:
  • REOPOLD è robusto alla scelta del modello teacher e scala bene anche su modelli studenti più grandi (7B), dove i metodi vanilla falliscono a causa dell'instabilità.
  • Nei task di ragionamento agenziale (uso di tool visivi), supera sia RKL che GRPO, dimostrando una migliore generalizzazione.

5. Significato e Impatto

Il lavoro di REOPOLD è significativo perché:

• Colma il divario tra Distillazione e RL: Fornisce una base teorica solida che collega la distillazione on-policy all'ottimizzazione della politica, applicando le migliori pratiche del RL moderno (clipping, controllo della varianza) alla distillazione.
• Abilita Modelli Compatti: Risolve il problema dell'instabilità che finora limitava l'uso di modelli piccoli per compiti di ragionamento complessi, rendendo possibile ottenere prestazioni di livello "o1" o "R1" su modelli da pochi miliardi di parametri.
• Efficienza Computazionale: Offre un percorso praticabile per scalare le capacità di ragionamento senza richiedere risorse computazionali massive, riducendo i tempi di addestramento e di inferenza.

In sintesi, REOPOLD dimostra che un approccio "rilassato" e intelligente alla distillazione, che filtra attivamente i segnali dannosi e si concentra sui punti di divergenza critica, è superiore sia alla distillazione on-policy classica che al Reinforcement Learning diretto per i modelli di ragionamento di piccole dimensioni.