Riepilogo Tecnico: TRACE - Distillazione Dove Conta tramite Allineamento Self On-Policy Instradato a Livello di Token

1. Enunciato del Problema

Il documento affronta le limitazioni della Self-Distillazione On-Policy (Self-OPD) quando applicata all'Apprendimento per Rinforzo con Ricompense Verificabili (RLVR) a orizzonte lungo, in particolare nei compiti di ragionamento matematico. Sebbene i metodi Self-OPD (ad esempio SDPO, SRPO) densifichino il segnale di ricompensa sparso utilizzando un'istantanea della politica come insegnante per guidare lo studente, gli autori identificano una modalità di fallimento critica nell'addestramento a orizzonte lungo:

• Disallineamento della Granularità: I metodi esistenti applicano la funzione di perdita per la Distillazione della Conoscenza (KL) su tutti i token in una risposta. Tuttavia, l'analisi mostra che circa il 97% della massa di probabilità nelle traiettorie corrette è già condivisa tra studente e insegnante, il che significa che la maggior parte dei gradienti KL è ridondante.
• La "Tassa di Distillazione Uniforme": Applicare il KL a tutti i token costringe il modello a spendere gradienti su posizioni che non necessitano di correzione. Ciò porta a:
  • Fuga di Informazioni Privilegiate: Lo studente assorbe correlazioni presenti nel contesto privilegiato (ad esempio, tracce verificate) che sono assenti durante l'inferenza, causando un degrado Fuori Distribuzione (OOD).
  • Aumento dell'Entropia e Crollo del Ragionamento: Gli autori osservano un "crollo a tre sintomi" nell'addestramento a orizzonte lungo: le lunghezze delle risposte si accorciano drasticamente, l'entropia dell'attore per token aumenta (spesso >4 volte la baseline) e l'accuratezza di validazione diminuisce significativamente.
• La Domanda Fondamentale: L'unità corretta di distillazione non è l'intera risposta, ma il piccolo sottoinsieme di token di ragionamento decisivi dove la distribuzione dello studente ha effettivamente bisogno di correzione.

2. Metodologia: TRACE

Gli autori propongono TRACE (Token-Routed Alignment for Critical rEasoning), un metodo che instrada i segnali di distillazione solo verso specifici span di token invece di applicarli uniformemente.

Componenti Principali

1. Self-OPD Instradato a Livello di Token: Invece di una perdita KL globale, TRACE utilizza un annotatore privilegiato ($\pi_A$) per identificare span specifici in una traiettoria:

   • Span Chiave ($K_y$): Passaggi di ragionamento decisivi nelle traiettorie corrette (Ricompensa=1).
   • Span di Errore ($E_y$): Errori di ragionamento alla radice nelle traiettorie errate (Ricompensa=0).
   • Non-Span ($N_y$): Tutti i token rimanenti.
   • Vincolo: La copertura totale degli span è limitata a $\alpha = 0.25$ (25% dei token) per limitare l'esposizione alle informazioni privilegiate.

2. Instradamento Diagnostico Grossolano:

   • L'annotatore emette un etichetta di tipo grossolana (ad esempio, case_split, missed_case) ma non il testo dello span o il contenuto completo della traiettoria.
   • L'Insegnante ($\pi_T$) riceve il problema e questa etichetta grossolana come prefisso diagnostico privato. Calcola i logit causalmente sul prefisso generato dallo studente ma non vede le posizioni specifiche degli span né il contenuto.
   • Lo Studente ($\pi_S$) viene addestrato utilizzando una perdita instradata in cui il termine KL è applicato solo agli span mascherati.

3. Instradamento di Divergenza Asimmetrica:

   • KL Inverso (FKL) sugli Span Chiave: Applicato alle traiettorie corrette. Il FKL ($KL(\pi_T \parallel \pi_S)$) fornisce un "sollevamento" ai token che lo studente sotto-alloca, indipendentemente dalla massa di probabilità corrente dello studente.
   • KL Inverso (RKL) sugli Span di Errore (Opzionale): Applicato alle traiettorie errate. Il RKL ($KL(\pi_S \parallel \pi_T)$) sopprime i token su cui lo studente è eccessivamente sicuro ma che l'insegnante sconsiglia.
   • Nessun KL sui Non-Span: Questi token si affidano esclusivamente alla perdita standard GRPO (Group Relative Policy Optimization).

4. Programma di Esposizione Finita:

   • Il peso del KL ($\lambda_k$) viene annealato a zero dopo un breve warm-up (ad esempio, 40 step).
   • Ciò garantisce che l'esposizione cumulativa ai gradienti privilegiati rimanga finita sull'orizzonte di addestramento, prevenendo l'accumulo a lungo termine di correlazioni che degradano l'OOD.

3. Analisi Teorica

Il documento fornisce una giustificazione teorica per la strategia di instradamento:

• Pressione Puntuale Asimmetrica (Corollario 2):
  • Il FKL è efficace per i token sotto-allocati (dove la probabilità dello studente $p \ll$ probabilità dell'insegnante $q$) perché la pressione del gradiente scala con il divario $(q-p)$ e non svanisce quando $p \to 0$.
  • Il RKL è efficace per i token sicuri-ma-sbagliati (dove $p \gg q$) perché il gradiente scala con l'eccessiva sicurezza dello studente ($p$).
• Limite di Esposizione Cumulativa (Proposizione 3): Combinando la mascheratura degli span (limitando la banda a $\alpha$) e il decadimento del KL (limitando la durata), TRACE vincola l'esposizione cumulativa ai gradienti privilegiati a $O(\alpha \Lambda^2)$, prevenendo la crescita illimitata osservata nelle baseline KL su tutti i token persistenti.
• Segnale vs Fuga: Il metodo massimizza il segnale di allineamento sugli span chiave minimizzando al contempo il rischio di fugare informazioni privilegiate che causano correlazioni spurie.

4. Contributi Chiave

1. Diagnosi: Identificazione di una "tassa di distillazione uniforme" nella Self-OPD su tutti i token, collegando l'aumento dell'entropia, la soppressione epistemica e la fuga di gradienti privilegiati a un disallineamento della granularità.
2. Metodo e Teoria: Introduzione di TRACE, che instrada $\{FKL, RKL, \emptyset\}$ in base alla classe del token (Chiave, Errore, Non-span) e fa decadere il canale privilegiato. La teoria spiega perché il FKL è dominante per gli span chiave nei regimi con base forte e perché l'angolo ottimale si sposta verso il RKL per gli span di errore nei regimi con base più debole.
3. Evidenza Empirica:
   • Prestazioni: Su Qwen3-8B, TRACE migliora la media su 5 benchmark di 2,76 punti percentuali rispetto a GRPO.
   • Stabilità: A differenza di SDPO/SRPO, TRACE preserva le prestazioni OOD del modello di base su GPQA-Diamond, dove le baseline degradano.
   • Auto-Annotazione: Il metodo funziona con un auto-annotatore online (la politica dello studente stessa), recuperando circa il 69% del guadagno derivante da un potente annotatore API esterno, dimostrando che i guadagni provengono dal meccanismo di instradamento, non solo dalla capacità esterna.
   • Adattamento Cross-Scala: Su una base più debole (Qwen3-1.7B), l'angolo di instradamento dominante si inverte naturalmente da FKL-su-Chiave a RKL-su-Errore, corrispondendo alle previsioni teoriche.

5. Risultati

• Qwen3-8B (Base Forte): TRACE (FKL sugli Span Chiave) raggiunge un'accuratezza media del 81,51% su 5 benchmark (MATH-500, AIME 24/25, AMC 23, GPQA-Diamond), superando GRPO (78,75%) e la baseline Self-OPD precedente più forte RLSD (78,99%). Crucialmente, mantiene il punteggio GPQA-Diamond del modello di base (58,33% vs 58,27% base), mentre GRPO e le baseline su tutti i token crollano significativamente.
• Qwen3-1.7B (Base Debole): TRACE (RKL sugli Span di Errore) raggiunge una media del 60,16%, l'unico metodo addestrato a superare il modello di base (58,77%). Le baseline su tutti i token (SDPO/SRPO) crollano, scendendo di 11-19 punti sotto GRPO.
• Studi di Ablazione:
  • Localizzazione della Maschera: Sostituire le maschere degli span specifici di TRACE con maschere casuali del 25% o KL su tutti i token comporta cali di prestazioni, confermando che quali token sono selezionati conta più della quantità.
  • Qualità dell'Annotatore: Utilizzare lo studente in fase di addestramento attivo come proprio annotatore produce +1,90 punti, dimostrando che il metodo non dipende dall'importare capacità di ragionamento "super-umane" esterne.

6. Significato e Affermazioni

Il documento afferma che TRACE offre una soluzione geometrica alla perdita per l'instabilità della Self-OPD a orizzonte lungo, piuttosto che una correzione specifica per il dominio. Il suo significato risiede in:

• Disaccoppiamento della Distillazione dall'Instabilità: Limitando il segnale "privilegiato" agli span critici e facendolo decadere, TRACE evita l'aumento dell'entropia e il crollo del ragionamento osservati in SDPO e SRPO.
• Instradamento Adattivo: Dimostra che la direzione ottimale di distillazione (FKL vs RKL) non è un iperparametro fisso ma dipende dal profilo di errore del modello di base (sotto-allocazione vs eccessiva sicurezza).
• Fattibilità Pratica: Dimostra che è possibile ottenere distillazione di alta qualità senza insegnanti esterni o rubriche costose per query, affidandosi invece a un'annotazione diagnostica leggera e grossolana che può essere auto-generata.

Gli autori sottolineano che TRACE è un "framework instradato per angolo" piuttosto che una ricetta fissa, dove l'azione di instradamento specifica (FKL vs RKL) dovrebbe essere scelta in base alla modalità di errore residua del modello di base.