Sintesi Tecnica: Temperature-Scaled On-Policy Self-Distillation (TS-OPSD)

1. Definizione del Problema

Il Reinforcement Learning da Ricompense Verificabili (RLVR) ha migliorato significativamente le capacità di ragionamento dei Large Language Models (LLM), in particolare nei domini matematici e logici. Tuttavia, questo progresso affronta un collo di bottiglia critico noto come collasso dell'entropia (entropy collapse). Man mano che l'ottimizzazione della policy procede, la distribuzione di output del modello diventa sempre più concentrata (acuta), portando a:

• Riduzione della Diversità dei Rollout: Il modello genera traiettorie ripetitive, limitando l'esplorazione dello spazio delle soluzioni.
• Diminuzione dei Segnali di Apprendimento: Le informazioni di gradiente efficaci si concentrano su una piccola frazione di token ad alta entropia in punti critici di ramificazione del ragionamento. Una volta che la policy collassa, questi segnali svaniscono, bloccando ulteriori miglioramenti.

I rimedi esistenti tentano di mitigare questo problema in due modi:

1. Modifica dell'Obiettivo RL: Utilizzando la regolarizzazione dell'entropia o il clipping adattivo. Questi metodi impongono vincoli esterni che possono entrare in conflitto con la massimizzazione della ricompensa.
2. Regolazione della Temperatura di Campionamento: Aumentando la temperatura durante la raccolta dei rollout. Sebbene ciò alteri la policy di comportamento (behavior policy), non cambia la distribuzione parametrizzata sottostante del modello, creando un disallineamento tra la policy di comportamento e la policy ottimizzata.

Questi approcci rimangono esterni ai parametri del modello. L'articolo postula che l'effetto esplorativo della temperatura debba essere internalizzato direttamente nei pesi del modello per fornire una migliore inizializzazione per il proseguimento del RL.

2. Metodologia: TS-OPSD

Gli autori propongono la Temperature-Scaled On-Policy Self-Distillation (TS-OPSD), un metodo leggero di "riaccensione della policy" (policy reheating) progettato per ripristinare l'entropia senza insegnanti esterni, dati privilegiati o costi di inferenza aggiuntivi.

Meccanismo Core

TS-OPSD opera come un intervento tra le fasi di RL. Partendo da un checkpoint con entropia collassata ($\theta_{RL-collapse}$), il metodo costruisce un auto-insegnante (self-teacher) e distilla una distribuzione più fluida verso il modello studente ($\theta_{reheat}$).

1. Costruzione dell'Auto-Insegnante: Dato i logit correnti dello studente $z_\theta$, una distribuzione auto-insegnante $p_T$ viene generata applicando una trasformazione softmax ad alta temperatura ($T > 1$) ai logit. Fondamentalmente, i logit dell'insegnante sono "stop-gradient" e non viene utilizzato alcun modello esterno.
   $$p_T(v | x_{<t}) = \frac{\exp(z_{\theta,v}/T)}{\sum_u \exp(z_{\theta,u}/T)}$$
2. Distillazione On-Policy: Lo studente viene addestrato per minimizzare la divergenza di Kullback-Leibler (KL) tra la sua distribuzione corrente $q_\theta$ e l'auto-insegnante ad alta temperatura $p_T$, utilizzando rollout on-policy campionati dalla sua stessa $q_\theta$.
   $$\mathcal{L} = D_{KL}(p_T \parallel q_\theta) \quad \text{(Forward KL)}$$
   $$\mathcal{L} = D_{KL}(q_\theta \parallel p_T) \quad \text{(Reverse KL)}$$

Proprietà Strutturali

• Preservazione del Ranking: Poiché $p_T$ è una trasformazione di potenza monotona di $q_\theta$ ($p_T(v) \propto q_\theta(v)^{1/T}$), il ranking dei token è preservato. L'insegnante e lo studente condividono lo stesso supporto locale e l'ordine di preferenza semantica; differisce solo la nitidezza della distribuzione.
• Internalizzazione: A differenza degli aggiustamenti standard della temperatura di decoding, TS-OPSD aggiorna i pesi del modello per codificare permanentemente il comportamento più fluido e ad alta entropia. Ciò elimina il disallineamento tra addestramento e inferenza.
• Dinamica del Gradiente:
  • Forward KL (FKL): Il gradiente assume la forma $q_\theta(v) - p_T(v)$. Questo esegue uno "smoothing strutturato", ricalibrando le alternative soppresse in base al giudizio pregresso del modello stesso.
  • Reverse KL (RKL): Sotto l'accoppiamento di auto-temperatura, minimizzare la RKL è localmente equivalente alla massimizzazione dell'entropia dello studente, con la temperatura che agisce solo come moltiplicatore scalare sulla magnitudo del gradiente.

3. Contributi Chiave

1. Framework di Policy Reheating: Introduce TS-OPSD, un metodo per internalizzare l'effetto esplorativo della temperatura nei parametri del modello, fungendo da intervento post-collasso.
2. Insight Teorico: Dimostra che, sotto l'accoppiamento di auto-temperatura, il disallineamento insegnante-studente è minimizzato (preservazione del ranking) e la RKL si riduce a un aggiornamento diretto di massimizzazione dell'entropia.
3. Validazione Empirica: Dimostra che il riscaldamento (reheating) di un checkpoint con entropia collassata fornisce un'inizializzazione superiore per il proseguimento del RL rispetto al normale addestramento continuato o agli aggiustamenti della temperatura a livello di rollout.
4. Analisi delle Dinamiche: Caratterizza il fenomeno del "salto di entropia" (entropy jump), dove il recupero dell'entropia avviene in due fasi (smoothing lento seguito da un salto discontinuo e netto), e analizza la reversibilità di questi aggiornamenti dei parametri.

za. Risultati Sperimentali

Gli esperimenti sono stati condotti su Qwen3-4B-Base e Qwen3-8B-Base utilizzando il dataset DAPO-Math-17K e valutati su MATH-500, AIME24/25/26 e OlymMATH.

• Guadagni di Performance: TS-OPSD (specificamente con Forward KL, PR-FKL) ha superato costantemente sia il baseline standard DAPO che il Rollout Reheating (DAPO w/ RR).
  • Su Qwen3-8B-Base, PR-FKL ha raggiunto un Avg pass@8 di 54.74, rispetto a 51.60 per il DAPO standard e 52.66 per il Rollout Reheating.
  • I guadagni sono stati più pronunciati sui benchmark più difficili (AIME, OlymMATH).
• Recupero dell'Entropia: Dopo il reheating, i modelli hanno ripreso il RL da un'inizializzazione ad alta entropia e hanno raggiunto ricompense finali superiori rispetto al DAPO standard, che rimaneva soppresso.
• Confronto con il Rollout Reheating: Sebbene l'aumento della temperatura di campionamento durante i rollout (RR) abbia fornito guadagni non trascurabili, il Policy Reheating (PR) ha prodotto risultati superiori. Ciò suggerisce che il beneficio derivi dalla distribuzione della policy più fluida in sé, che funge da migliore fondamento per l'ottimizzazione, piuttosto che solo da un campionamento più rumoroso.
• Preservazione delle Capacità: Le analisi hanno mostrato che TS-OPSD esegue una ridistribuzione della probabilità che preserva il supporto:
  • Rappresentazioni Intermedie: I punteggi lineari CKA tra il modello base e quello reheated sono rimasti vicini a 1.0, indicando che le rappresentazioni interne sono rimaste largamente invariate.
  • Capacità di Ragionamento: L'accuratezza nei task su MATH-500 e GPQA è rimasta stabile durante tutto il processo di distillazione, confermando che il recupero dell'entropia non ha degradato le capacità di ragionamento.
  • Set di Candidati: La sovrapposizione dei top-10 token e la ritenzione della massa di probabilità erano vicine a 1.0, indicando che le preferenze semantiche del modello sono state preservate mentre i margini sono stati ammorbiditi.

5. Significato e Rivendicazioni

L'articolo sostiene che il ripristino dell'entropia può essere scollegato dall'obiettivo di RL e inserito come un'operazione modulare tra le fasi, quando una policy diventa troppo deterministica.

• Intervento Modulare: TS-OPSD offre un meccanismo semplice e leggero per estendere il RL orientato al ragionamento oltre il collo di bottiglia del collasso dell'entropia senza alterare l'algoritmo di RL principale.
• Stabilità in Ambienti con Scarsità di Dati: Negli esperimenti con dati ridotti (2048 prompt), l'addestramento RL standard è diventato instabile (picchi di gradiente, degradazione della ricompensa) a causa del severo collasso dell'entropia. TS-OPSD ha agito come un meccanismo stabilizzante, permettendo un addestramento più lungo e stabile.
• Efficienza Computazionale: L'intervento di reheating aggiunge circa il 15–17% al tempo totale di addestramento RL. Quando normalizzato per dimensione del batch, il costo per campione è comparabile o addirittura inferiore ai normali step di RL.
• Limitazioni: Gli autori notano che gli esperimenti sono limitati a specifiche scale di modelli (4B/8B) e che i criteri di arresto per il reheating (soglia di entropia) sono attualmente empirici piuttosto che basati su principi. Suggeriscono che il lavoro futuro possa esplorare programmi a cascata alternando OPSD e RL.

In sintione, TS-OPSD dimostra che "internalizzare" la temperatura tramite l'auto-distillazione ripristina efficacementamente la capacità esplorativa delle policy collassate, consentendo al RL di continuare a raggiungere livelli di performance più elevati pur preservando le strutture di ragionamento acquisite durante l'addestramento.