Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Sintesi Tecnica: Colmare le Traiettorie di Ragionamento nella Distillazione On-Policy tramite Guida Near-Future

Problematica
La Distillazione On-Policy (OPD) è un pilcro della moderna fase di post-training per i Large Language Models (LLM), progettata per migliorare il ragionamento addestrando un modello studente sulle traiettorie campionate dalla propria policy sotto la supervisione di un insegnante. La OPD standard opera su un meccanismo di "campionamento della traiettoria ma apprendimento del token": identifica le deviazioni tramite token ad alta perdita (high-loss) e le ripara utilizzando la correzione della divergenza reverse-KL locale.

Gli autori identificano un disallineamento fondamentale in questo approccio, sostenendo che la supervisione a livello di token sia strutturalmente disallineata rispetto ai requisiti della correzione del ragionamento a livello di traiettoria. Questo disallineamento si manifesta in due modalità critiche di fallimento:

1. Falsi Allarmi Rumorosi: I token ad alta perdita non corrispondono sempre a una genuina divergenza di ragionamento. Un'analisi empirica rivela che circa il 30% dei token ad alta perdita ricade in un regime di bassa divergenza, rappresentando discrepanze di forma superficiale (ad esempio, preferenze stilistiche come "e" invece di "poi") piuttosto che biforcazioni logiche. Trattare questi come errori critici fornisce una supervisione rumorosa che può interferire attivamente con la correzione del ragionamento.
2. La Trappola dell'Apprendimento Token-per-Token: Anche quando vengono identificati punti di divergenza genuina, la correzione reverse-KL singolo-token è spesso insufficiente per reindirizzare l'intera traiettoria di ragionamento. I fallimenti del ragionamento evolvono tipicamente come drift distribuzionali a breve orizzonte. Correggere un singolo token localmente non garantisce che il modello si allinei alla futura traiettoria dell'insegnante; il modello spesso deriva nuovamente verso rami errati immediatamente dopo la correzione locale.

Metodologia: Trajectory-aware OPD (TOPD)
Per affrontare queste limitazioni, gli autori propongono la Trajectory-aware OPD (TOPD), che colma il divario tra la correzione locale del token e l'evoluzione globale della traiettoria incorporando informazioni sulla traiettoria near-future (vicino-futura). Il metodo consiste in due fasi primarie:

1. Rilevamento della Divergenza:

   • Per le posizioni candidate ad alta perdita, il sistema torna indietro al token precedente per stabilire un prefisso condiviso.
   • Sia l'insegnante che lo studente generano continuazioni a finestra corta (lunghezza $K=50$) a partire da questo prefisso.
   • Il sistema misura la discrepanza a livello di traiettoria utilizzando la distanza di Trasporto Ottimale (OT) tra le continuazioni dell'insegnante e dello studente.
   • I veri punti divergenti sono identificati come posizioni che esibiscono sia un'alta perdita a livello di token che un'alta divergenza OT a finestra corta. Questo filtra i mismatch superficiali benigni.

2. Iniezione del Segnale dell'Insegnante a Livello di Traiettoria:

   • Inveve di ottimizzare solo l'attuale token, TOPD estende la supervisione all'intera traiettoria a finestra corta.
   • Utilizzando il piano di trasporto OT ($\gamma$), il metodo costruisce un "target morbido consapevole della traiettoria" (trajectory-aware soft target) per ogni posizione futura dello studente. Questo target allinea i futuri token dello studente con la traiettoria dell'insegnante basandosi sulla mappatura del trasporto ottimale.
   • Lo studente viene addestrato per corrispondere a questi target morbidi allineati all'OT attraverso la finestra, imparando efficacementamente come guidare il proprio percorso di ragionamento near-future verso il percorso dell'insegnante, non solo come correggere il token corrente.
   • L'obiettivo finale combina la perdita OPD standard con un termine di perdita della traiettoria ($L_{TOPD} = L_{OPD} + \lambda L_{traj}$).

Contributi Chiave

• Identificazione del Disallineamento: Il lavoro stabilisce che la supervisione a livello di token è un "proxy grossolano" per la correzione del ragionamento a livello di traiettoria, fallendo nel distinguere tra genuine biforcazioni logiche e variazioni superficiali benigne.
• Caratterizzazione Empirica delle Modalità di Fallimento: Gli autori forniscono prove che ~30% dei token ad alta perdita sono falsi allarmi a bassa divergenza e dimostrano la "trappola dell'apprendimento token-per-token" dove le correzioni locali falliscono nel ripristinare l'intera traiettoria di ragionamento.
• Principio di Guida a Livello di Traiettoria: Il paper introduce un principio per cui l'informazione della traiettoria near-future è distribuita su più token. Utilizzando l'allineamento OT a finestra corta, TOPD inietta le discrepanze di percorso tra insegnante e studente direttamente nella funzione di perdita, consentendo allo studente di apprendere spostamenti di traiettoria piuttosto che isolati editing di token.

Risultati Sperimentali
Gli esperimenti sono stati condotti utilizzando modelli Qwen3 (distillando un insegnante da 30B in uno studente da 4B) su benchmark matematici competitivi (AIME24, AIME25, HMMT25-Feb).

• Guadagni di Performance: La OPD standard ha migliorato l'accuratezza media del modello warm-start dal 38,9% al 47,8%. TOPD ha ulteriormente aumentato questa percentuale al 52,2%.
• Benchmark Specifici: TOPD ha migliorato l'accuratezza di AIME24 dal 60,0% al 63,3% e di AIME25 dal 46,7% al 53,3%.
• Studi di Ablazione:
  • Sopprimendo i token ad alta perdita ma a bassa OT (falsi allarmi) nella OPD standard, si è ottenuto un modesto guadagno (+0,4%), confermando che questi token sono rumorosi.
  • Sopprimendo i token ad alta OT ma ad alta perdita (vere divergenze), si è verificato un calo significativo delle prestazioni, confermando la loro necessità.
  • L'analisi della correzione della traiettoria locale ha mostrato che mentre la OPD standard riduceva parzialmente la distanza OT a finestra corta (da 0,538 a 0,350), TOPD ha ottenuto una riduzione sostanziale (a 0,204), dimostrando una capacità superiore di reindirizzare le traiettorie near-future.
• Efficienza: TOPD introduce un modesto overhead computazionale, con un tempo di addestramento di 1,41× rispetto alla OPD standard, poiché l'OT viene calcolato solo su finestre selezionate ad alta perdita.

Significatività e Rivendicazioni
Il paper sostiene che i fallimenti del ragionamento nella OPD sono fondamentalmente fenomeni a livello di traiettoria piuttosto che isolati mismatch di token. Passando dall'editing di token puntuale all'allineamento globale consapevole delle transizioni, TOPD consente un allineamento del ragionamento più efficace. Gli autori pongono questo lavoro come un passo verso obiettivi di distillazione più consapevoli della traiettoria, sostenendo che una correzione del ragionamento di successo richiede che lo studente impari non solo quale token correggere, ma come navigare il percorso di ragionamento successivo per rimanere coerente con l'insegnante.

Limitazioni
Gli autori riconoscono che TOPD comporta costi computazionali aggiuntivi a causa dei calcoli OT, si concentra principalmente sulla divergenza a breve orizzonte (potendo perdere errori a lungo contesto) ed è attualmente validato principalmente su benchmark di ragionamento matematico piuttosto che su compiti di coding o a contenuto aperto. Inoltre, la distanza OT misura la prossimità della traiettoria ma potrebbe non catturare l'equivalenza semantica dove percorsi diversi portano alla stessa risposta corretta.