Sintesi Tecnica: Combattere il Decadimento della Fedeltà della Supervisione nella Distillazione On-Policy

Problema: Decadimento della Fedeltà della Supervisione (SFD)

Il paper identifica un collo di bottiglia critico nella Distillazione On-Policy (OPD), un paradigma in cui un modello studente apprende dalle proprie traiettorie generate utilizzando il feedback a livello di token da parte di un modello insegnante. Sebbene l'OPD sia stata efficace nel trasferire capacità di ragionamento, gli autori osservano che le sue prestazioni degradano significativamente all'aumentare della lunghezza della generazione (ad esempio, le prestazioni si stabilizzano intorno ai 16k token e declinano a 39k token).

Gli autori attribuiscono questo fenomeno al Decadimento della Fedeltà della Supervisione (Supervision Fidelity Decay - SFD). Man mano che lo studente genera sequenze più lunghe, i suoi output si allontanano progressivamente dalla distribuzione di addestramento dell'insegnante (Out-of-Distribution o OOD). Questo drift causa una minore confidenza e una minore capacità discriminativa della distribuzione del prossimo token dell'insegnante.

• Meccanismo: Nella distillazione reverse-KL standard, il segnale di apprendimento (vantaggio) dipende dalla differenza tra la log-probabilità dello studente e la log-probabilità dell'insegnante. Poiché la distribuzione dell'insegnante diventa diffusa (alta entropia) a causa dell'SFD, la varianza nelle log-probabilità dell'insegnante collassa.
• Conseguenza: Il contributo discriminativo dell'insegnante al gradiente svanisce. Il segnale di apprendimento si riduce effettivamente a un segnale "solo-studente" che rinforza i modi esistenti dello studente senza correzioni. Ciò crea un ciclo di feedback auto-rinforzante in cui il drift si accumula, restringendo l'orizzonte di apprendimento effettivo e stabilendo un limite alla lunghezza del ragionamento.

Metodologia: Lookahead Group Reward (LGR)

Per affrontare l'SFD, gli autori propongono il Lookahead Group Reward (LGR), un metodo che mantiene attivamente la qualità della supervisione ottimizzando per la futura confidenza dell'insegnante piuttosto che per l'attuale corrispondenza locale.

Intuizione Centrale

Anche quando la supervisione locale dell'insegnante alla posizione $t$ è collassata (ovvero, la distribuzione è vicina all'uniformità), diverse scelte di token a $t$ possono portare a livelli di confidenza dell'insegnante significativamente differenti alla posizione successiva ($t+1$). Il paper dimostra teoricamente che la discriminabilità un passo avanti può sopravvivere anche quando la supervisione locale fallisce. Selezionando i token che massimizzano la confidenza dell'insegnante al passo successivo, lo studente può rallentare il tasso di drift.

Design dell'Algoritmo

1. Confidence Reward: Ad ogni posizione $t$, il metodo valuta la probabilità del picco del prossimo token dell'insegnante ($\max_v \pi_T(v|x_{<t}, x_t)$) per i top-$K$ candidati dello studente. Questo valore funge da segnale di ricompensa che indica quanto bene un candidato preservi la capacità dell'insegnante di fornire una supervisione di alta qualità nel passo successivo.
2. Group Normalization: Per gestire la varianza nei livelli assoluti di confidenza tra diversi task e posizioni, le ricompense grezze vengono normalizzate all'interno del gruppo di candidati top-$K$ (similmente a GRPO). Ciò mantiene il ranking relativo dei candidati eliminando al contempo il rumore di scala assoluta.
3. Attivazione Triggerata dall'Entropia: Calcolare la ricompensa lookahead per tutti i token è computazionalmente costoso. Gli autori osservano che l'SFD e la necessità di una guida lookahead sono più critiche nelle posizioni ad alta entropia (dove lo studente è incerto). Pertanto, la riccompensa di confidenza viene calcolata solo per le posizioni in cui l'entropia della generazione dello studente supera una soglia $\tau$. Ciò riduce l'overhead computazionale di 4-7 volte.
4. Meccanismo di Tree-Attention: Per valutare efficientemente $K$ candidati per molteplici posizioni senza rieseguire il prefill dell'insegnante per ogni candidato, gli autori impiegano un meccanismo di tree-attention. Questo costruisce una sequenza estesa in cui il ramo principale è la traiettoria dello studente e i token candidati sono aggiunti come rami. Una maschera di attenzione sparsa permette all'insegnante di elaborare tutti i candidati in un singolo passaggio forward, ottenendo un'accelerazione teorica fino a 1000$\times$ rispetto a un'implementazione ingenua.

La funzione di perdita finale combina lo standard vantaggio reverse-KL con la ricompensa di confidenza lookahead, pesata da un iperparametro $\gamma$.

Contributi Chiave

1. Identificazione dell'SFD: Il paper caratterizza formalmente il Decadimento della Fedeltà della Supervisione come un modo di fallimento fondamentale della distillazione on-policy reverse-KL, dimostrando che esso fa collassare il gradiente in un ciclo di drift auto-rinforzante che limita la lunghezza del ragionamento.
2. Lookahead Group Reward (LGR): Gli autori propongono un rimedio principato che guarda un passo avanti. Invece di cercare di corrispondere a una distribuzione locale degradata, LGR seleziona i token che massimizzano la futura confidenza dell'insegnante, guidando efficacementamente lo studente verso stati in cui la supervisione rimane affidabile.
3. Implementazione Efficiente: Il design di un meccanismo di tree-attention triggerato dall'entropia rende pratico il processo di lookahead computazionalmente intensivo, consentendo l'addestramento su contesti lunghi (fino a 39k token) senza costi proibitivi.

Risultati Sperimentali

Il metodo è stato valutato su sei benchmark che coprono il ragionamento matematico (AIME-24, AIME-25, AIME-26, HMMT-25, HMMT-26) e la generazione di codice (LiveCodeBench v6).

• Guadagni di Performance: Per un modello studente da 7B, LGR ha migliorato il punteggio mean@8 di 2.57 punti rispetto alla OPD standard (9k token) in media.
• Scaling del Lungo Contesto: I benefici di LGR scalano con la lunghezza della generazione. Sul benchmark AIME-26 a 39k token, LGR ha ottenuto un guadagno di +4.92 punti rispetto alla OPD. Al contrario, le prestazioni della OPD standard sono degradate significativamente a questa lunghezza.
• Confronto con i Baseline: LGR ha superato altri metodi di distillazione, inclusi JSD, REOPOLD e varianti della OPD standard. In particolare, i metodi che si limitano a stabilizzare l'addestramento (come REOPOLD) senza affrontare la degradazione della qualità della supervisione non sono riusciti a eguagliare le prestazioni di LGR.
• Studi di Ablazione:
  • Max-p vs. Entropia: L'uso della probabilità massima dell'insegnante (Max-p) come metrica di confidenza ha superato l'uso di entropia o perplexity, allineandosi con il risultato teorico secondo cui la Max-p è un proxy diretto della prossimità al manifold in-distribution dell'insegnante.
  • Sensibilità del Peso: Il peso della ricompensa di confidenza $\gamma$ ha mostrato prestazioni ottimali nell'intervallo da 1.0 a 1.5. Pesi eccessivamente alti hanno portato al reward hacking (sequenze ripetitive).

Significato e Rivendicazioni

Il paper sostiene che LGR fornisce una soluzione strutturalmente diversa al problema del ragionamento a lungo orizzonte nella distillazione. A differenza delle strategie che evitano passivamente le regioni di supervisione debole (ad esempio, la troncatura) o che si limitano a stabilizzare l'ottimizzazione, LGR guida attivamente lo studente verso traiettorie in cui l'insegnante mantiene un'alta fedeltà della supervisione.

Gli autori sottolineano che il loro lavoro affronta la causa radice del decadimento delle prestazioni nelle catene di ragionamento lunghe: l'effetto cumulativo del drift dello studente sulla confidenza dell'insegnante. Operazionalizzando l'intuizione che "una maggiore confidenza nella posizione successiva implica un segnale futuro più discriminativo", LGR rompe il circolo vizioso del decadimento della supervisione, permettendo ai modelli di scalare efficacemente verso catene di ragionamento molto più lunghe (fino a 39k token) di quanto precedentemente possibile con la distillazione on-policy standard.