🤖 AI

Trajectory-Refined Distillation

Questo articolo introduce la Trajectory-Refined Distillation (TRD), un nuovo metodo di distillazione on-policy che mitiga il problema strutturale del "prefix failure" applicando correzioni a livello di traiettoria ai rollout dello studente sotto la guida dell'insegnante, migliorando così l'accuratezza e la copertura del ragionamento attraverso vari benchmark e scale di modelli.

Autori originali: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Pubblicato 2026-06-09

📖 4 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di insegnare a uno studente (l'IA) come risolvere un problema matematico complesso. Hai un insegnante brillante (un'IA più avanzata) che conosce la risposta corretta.

Nel metodo standard descritto nel documento, chiamato On-Policy Distillation (OPD), il processo funziona così:

Lo studente prova a risolvere il problema da solo.
Mentre lo studente scrive i suoi passaggi, l'insegnante osserva e dice: "Per questa specifica parola che hai appena scritto, ecco la probabilità di ciò che viene dopo".
Lo studente cerca di adattare il proprio cervello per corrispondere alle probabilità dell'insegnante per ogni singola parola.

Il Problema: "Prefix Failure" (La Svolta Sbagliata)

Il documento identifica un grave difetto di questo metodo standard chiamato Prefix Failure.

Immagina che lo studente stia percorrendo un sentiero per risolvere il problema. Al passaggio 3, prende una strada sbagliata. Si trova ora in una strada senza uscita.

Il Dilemma dell'Insegnante: L'insegnante vede che lo studente si è perso. L'insegnante sa che il percorso corretto inizia con "Aspetta, in realtà..." e poi prende una direzione completamente diversa.
La Confusione: Ma l'insegnante è costretto a guardare la posizione attuale dello studente (la strada senza uscita). Di conseguenza, il consiglio dell'insegnante diventa un mix confuso: "Resta su questa strada senza uscita perché sei già qui" E "Gira immediatamente l'angolo".
Il Risultato: Lo studente riceve un segnale confuso. Cerca di imparare dall'insegnante, ma poiché lo studente è bloccato sul percorso sbagliato, il consiglio dell'insegnante risulta frammentato. Lo studente impara a essere "bravo a essere perso" piuttosto che imparare come tornare sulla strada giusta. Il documento lo chiama un "mix bimodale": l'insegnante sta cercando di insegnare due cose opposte contemporaneamente, e lo studente si confonde.

Le soluzioni esistenti hanno cercato di risolvere il problema semplicemente modificando il "voto" (la funzione di perdita o loss function) per parole specifiche, ma non hanno risolto il fatto che lo studente fosse ancora bloccato sul percorso sbagliato.

La Soluzione: Trajectory-Refined Distillation (TRD)

Gli autori propongono un nuovo metodo chiamato Trajectory-Refined Distillation (TRD). Invece di limitarsi a valutare la bozza disordinata dello studente parola per parola, cambia il flusso di lavoro:

La Bozza: Lo studente fa comunque un primo tentativo sul problema (la "raw rollout").
La Riscrittura: Prima della valutazione, l'insegnante prende quella bozza disordinata e la riscrive. L'insegnante corregge le svolte errate, corregge la logica e produce una versione pulita e perfetta della soluzione basata sul punto di partenza dello studente.
La Lezione: Lo studente impara quindi da questa versione pulita e corretta, non dalla bozza disordinata originale.

L'Analogia:

Il Vecchio Modo: Scrivi una storia con un buco nella trama a metà del racconto. Il tuo editor legge la tua storia e dice: "Per la parola 'gatto' nel paragrafo 3, avresti dovuto scrivere 'cane'". Ma tu sei ancora bloccato nel paragrafo 3 cercando di capire perché hai scritto "gatto" in primo luogo. È confusionario.
Il Modo TRD: Scrivi la storia con il buco nella trama. Il tuo editor prende la tua bozza, corregge il buco nella trama e riscrive l'intero segmento centrale in modo che la storia scorra perfettamente. Poi, studi la versione riscritta dall'editor per imparare a scrivere meglio le storie la prossima volta.

Perché Questo Funziona Meglio

Risolve la Causa Radice: Correggendo la "svolta errata" prima che lo studente provi a imparare da essa, il consiglio dell'insegnante non è più confuso. È un percorso unico e chiaro.
Esplora Nuovi Percorsi: Anche se lo studente trova la risposta giusta al primo colpo, l'insegnante potrebbe mostrargli un modo diverso, più breve o più intelligente per risolvere il problema. Questo insegna allo studente a essere più flessibile.
Efficienza: Il documento ha scoperto che i percorsi "riscritti" erano spesso molto più brevi e chiari rispetto ai tentativi originali e prolissi dello studente, il che ha reso l'addestramento effettivamente più veloce.

I Risultati

I ricercatori hanno testato questo metodo su difficili competizioni matematiche (come AIME e HMMT) e sfide di programmazione. Hanno scoperto che:

Il nuovo metodo (TRD) ha costantemente superato i vecchi metodi.
È stato particolarmente efficace nell'aiutare l'IA a risolvere i problemi più difficili dove solitamente rimaneva bloccata.
Ha funzionato sia che l'insegnante fosse un'IA separata e più grande, sia che fosse la stessa IA ad agire come il proprio "insegnante intelligente" (usando una tecnica chiamata "informazione privilegiata" o privileged information).

In breve, il documento sostiene che per insegnare efficacemente a un'IA, non dovresti solo valutare i suoi errori parola per parola mentre li sta commettendo. Invece, dovresti lasciare che l'insegnante corregga prima l'intero percorso, e poi permettere allo studente di imparare dal viaggio corretto.

Riepilogo Tecnico: Trajectory-Refined Distillation (TRD)

1. Definizione del Problema: Il Fallimento del Prefisso nella Distillazione On-Policy

Il documento identifica un limite strutturale fondamentale nella Distillazione On-Policy (OPD) e nella Auto-Distillazione On-Policy (OPSD) per i Large Language Models (LLM), denominato prefix failure (fallimento del prefisso).

Nella standard OPD/OPSD, un modello studente genera una traiettoria (rollout) $y_o$ , e un modello insegnante fornisce una supervisione densa, per ogni token, basata su quel particolare rollout. Gli autori sostengono che quando il rollout dello studente segue un percorso di ragionamento errato (un "prefisso fallito"), il segnale di supervisione dell'insegnante diventa strutturalmente compromesso:

Miscela Bimodale dell'Insegnante: Quando un prefisso è errato, la distribuzione dell'insegnante diventa una miscela di due modalità: una che continua il prefisso fallito (per coerenza della sequenza) e un'altra che devia verso la soluzione corretta.
Gradienti Frammentati: Le funzioni di perdita per token standard (divergenza KL) valutano la supervisione lungo la traiettoria fallita fissa. Di conseguenza, gli aggiornamenti del gradiente sono frammentati. L'insegnante raccomanda ripetutamente un token di correzione (ad esempio, "Aspetta" o "In realtà") condizionato dal contesto errato (il prefisso fallito), invece di condizionarsi sul percorso di correzione che si sta sviluppando.
Inefficacia dei Correttivi a Livello di Token: I rimedi esistenti (ad esempio, clipping della perdita, troncamento top-K, pesatura) operano a livello di perdita del singolo token. Essi tentano di selezionare token "informativi" o sopprimere quelli "rumorosi", ma non riescono ad affrontare la causa radice: il prefisso fallito stesso rimane strutturalmente intatto. Le coppie di supervisione fornite allo studente divergono dal percorso di correzione ideale, impedendo al modello di imparare efficacemente come recuperare dagli errori.

2. Metodologia: Trajectory-Refined Distillation (TRD)

Per mitigare il fallimento del prefisso, gli autori propongono la Trajectory-Refined Distillation (TRD), un metodo che sposta l'intervento dal livello del token al livello della traiettoria.

Meccanismo Centrale

Invece di utilizzare direttamente il rollout grezzo dello studente $y_o$ per la distillazione, la TRD introduce un passaggio di raffinamento prima del calcolo della perdita:

Campionamento: Il modello studente $\pi_\theta$ genera un rollout on-policy grezzo $y_o \sim \pi_\theta(\cdot | x)$ .
Raffinamento: Il modello insegnante $\pi_T$ $π_{T}$ (o lo studente stesso in OPSD, condizionato da informazioni privilegiate come la verità di base $y^*$ $y^{*}$ ) genera una traiettoria raffinata $y_r$ $y_{r}$ .
- In OPD: $y_r \sim \pi_T(\cdot | x, y_o)$ .
- In OPSD: $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$ .
- Fondamentalmente, l'insegnante viene istruito a rivedere il rollout dello studente guidato dalla soluzione di riferimento, "riscrivendo" efficacemente le parti errate del prefisso pur mantenendo i pattern di ragionamento dello studente dove sono validi.
Distillazione: La traiettoria raffinata $y_r$ viene utilizzata come target di supervisione per la standard perdita di matching della distribuzione (Forward KL con vocabolario completo).

Giustificazione Teorica

La TRD affronta il problema del "gradiente frammentato" garantendo che il contesto di supervisione cresca lungo il percorso di correzione piuttosto che lungo il prefisso fallito.

Recupero del Gradiente Ideale: Supervisionando lungo $y_r$ , il contesto per ogni token include il prefisso corretto, permettendo all'insegnante di fornire un segnale di correzione coerente e multi-step ( $g_{ideal}$ ) invece di un segnale ripetitivo e mal condizionato dal contesto ( $g_{frag}$ ).
Esplorazione: Anche quando $y_o$ è corretto, la TRD espone lo studente ad alternative derivazioni valide suggerite dall'insegnante, ampliando la distribuzione dei percorsi di ragionamento corretti che lo studente apprende.

Algoritmo

Il processo è riassunto nell'Algoritmo 1:

Campiona $y_o \sim \pi_\theta(\cdot | x)$ .
Genera $y_r \sim \pi_T(\cdot | x, y_o)$ .
Aggiorna $\theta$ utilizzando la divergenza KL tra $\pi_\theta$ e la distribuzione dell'insegnante su $y_r$ .

3. Contributi Chiave

Identificazione del Fallimento del Prefisso: Il documento formalizza il "prefix failure" come una causa strutturale di degradazione nella distillazione on-policy, spiegando perché gli interventi a livello di token falliscono nel recuperare dagli errori di ragionamento.
Correzione a Livello di Traiettoria: La TRD è proposta come la prima progettazione di ottimizzazione a livello di traiettoria che mira esplicitamente al recupero dei prefissi falliti rispettando il vincolo di supporto on-policy, senza un carico computazionale proibitivo.
Miglioramento in Doppio Regime: La TRD migliora le prestazioni in due regimi:
- Correzione dell'Errore: Mitiga il fallimento del prefisso correggendo i prefissi problematici prima della distillazione.
- Esplorazione: Amplia la copertura del ragionamento esponendo gli studenti ad alternative derivazioni valide anche quando il rollout originale era corretto.
Applicabilità all'Auto-Distillazione: Il metodo si estende naturalmente alla OPSD, dove lo studente agisce come proprio insegnante condizionato da informazioni privilegiate, eliminando la necessità di un modello insegnante separato.

4. Risultati Sperimentali

Gli autori hanno valutato la TRD attraverso cinque benchmark di matematica competitiva (AIME24/25, HMMT25, BeyondAIME, AMOBench) e benchmark di generazione di codice (HumanEval+, MBPP+, LiveCodeBench) utilizzando modelli Qwen3 (1.7B, 4B, 8B).

Risultati Chiave

Prestazioni Superiori: La TRD supera costantemente i precedenti baseline (Forward/Reverse KL, con/senza clipping/troncamento) sia nelle impostazioni OPD che OPSD.
Guadagni Significativi su Task Difficili: I miglioramenti sono più pronunciati sui benchmark più difficili.
- AMOBench: La TRD ha ottenuto un miglioramento relativo di circa il 50% nel Pass@16 per Qwen3-8B sotto l'impostazione OPSD rispetto al modello base.
- AIME24/25: La TRD ha mostrato guadagni costanti, particolarmente per i modelli più piccoli (ad esempio, +4.6% su AIME24 per Qwen3-1.7B).
Stabilità: A differenza di alcuni baseline che degradano su certi benchmark (ad esempio, la Reverse KL che cala le prestazioni su AIME24), la TRD non è mai scesa al di sotto delle prestazioni del modello base.
Efficienza: Nell'impostazione OPSD, le traiettorie raffinate $y_r$ erano significativamente più brevi (lunghezza mediana compressa di circa 9 volte) rispetto ai rollout grezzi $y_o$ , riducendo il tempo di training (wall-clock time) di circa il 60% nonostante il passaggio di campionamento extra.
Generazione di Codice: La TRD ha eguagliato o superato i baseline su HumanEval+ e MBPP+, sebbene i guadagni su LiveCodeBench siano stati limitati, suggerendo che la capacità di raffinamento dell'insegnante possa essere limitata dalla difficoltà del compito.

5. Significato e Affermazioni

Il documento afferma che la TRD rappresenta un cambiamento necessario nella strategia di post-training per gli LLM. Gli autori sostengono che:

Gli interventi a livello di token sono insufficienti: La natura strutturale del fallimento del prefisso richiede correzioni a livello di traiettoria.
Il supporto on-policy è preservato: La TRD ottiene queste correzioni senza violare il vincolo on-policy (a differenza dell'uso di una politica esperta fissa), rendendola compatibile con gli standard di apprendimento efficienti dal punto di vista computazionale.
L'auto-miglioramento è fattibile: Il successo della TRD nell'impostazione OPSD dimostra che i modelli possono raffinare efficacemente i propri percorsi di ragionamento utilizzando informazioni privilegiate (verità di base) senza richiedere un modello insegnante separato e più grande, offrendo una via leggera per l'auto-miglioramento.

Il lavoro suggerisce che i futuri metodi di distillazione debbano tenere conto dell'integrità strutturale della traiettoria di ragionamento, non solo dell'allineamento statistico dei singoli token.