Riepilogo Tecnico: OISD – Auto-distillazione Interna On-Policy di Modelli Linguistici

1. Enunciato del Problema

Le recenti approcci di post-addestramento basati sull'apprendimento per rinforzo (RL) per i Modelli Linguistici su Larga Scala (LLM), in particolare quelli che utilizzano l'Ottimizzazione della Politica Relativa di Gruppo (GRPO) con ricompense verificabili, hanno avanzato significativamente il ragionamento matematico. Tuttavia, questi metodi sono prevalentemente centrati sull'output finale. Ottimizzano la politica di azione finale utilizzando ricompense sparse a livello di esito (ad esempio, la correttezza della risposta finale), trascurando in gran parte i segnali predittivi, i processi di ragionamento e le rappresentazioni interne codificate negli strati intermedi.

Sebbene la Distillazione della Conoscenza (KD) offra un meccanismo per fornire supervisione densa a livello di token, i paradigmi esistenti soffrono di limitazioni specifiche:

• Distillazione Off-Policy: Crea una discrepanza di distribuzione perché lo studente viene addestrato su traiettorie generate da un insegnante fisso, non dalle proprie generazioni.
• Distillazione On-Policy: Riduce la discrepanza ma tipicamente si affida a modelli insegnanti esterni.
• Auto-distillazione On-Policy: I metodi recenti rimuovono gli insegnanti esterni ma dipendono da informazioni privilegiate iniettate esternamente (ad esempio, prompt privilegiati, soluzioni ground-truth o feedback ausiliari) per supervisionare il modello.

Esiste un vuoto nello sfruttamento dei propri calcoli interni del modello come fonte di supervisione senza richiedere dati privilegiati esterni o multiple politiche di azione. Inoltre, ricerche precedenti indicano che le rappresentazioni del ragionamento evolvono sistematicamente attraverso gli strati (da previsioni ampie di candidati negli strati iniziali a decisioni convergenti negli strati profondi), suggerendo che gli strati intermedi potrebbero essere guidati dagli strati più profondi per migliorare la coerenza del ragionamento.

2. Metodologia: Framework OISD

Gli autori propongono OISD (Auto-distillazione Interna On-Policy), un framework che tratta lo strato finale del modello come un "insegnante interno distaccato" e uno strato intermedio selezionato come "studente interno ausiliario". Ciò avviene all'interno di un singolo passaggio in avanti durante il processo di rollout e ottimizzazione GRPO.

Architettura Core

• Singola Politica di Azione: Lo strato finale ($L$) rimane l'unica politica per generare rollouts e ottimizzare l'obiettivo GRPO.
• Insegnante-Studente Interno: Durante lo stesso passaggio in avanti, lo stato nascosto di uno strato intermedio selezionato ($l < L$) viene proiettato nello spazio dei token utilizzando la matrice di unembedding (seguendo il framework logit lens). Ciò crea una distribuzione "studente" provvisoria derivata dallo stesso contesto dell'"insegnante" (strato finale).
• Nessun Privilegio Esterno: I segnali di supervisione sono derivati interamente dagli stati interni del modello sullo stesso rollout on-policy, senza richiedere soluzioni ground-truth esterne o prompt privilegiati.

Obiettivi di Allineamento

OISD impiega due meccanismi interni di allineamento complementari per trasferire segnali predittivi dallo strato finale allo strato intermedio:

1. Allineamento dei Logit ("Come Pensare"):

   • Obiettivo: Trasferire comportamenti di ragionamento di alto livello e credenze predittive.
   • Meccanismo: La divergenza Jensen–Shannon (JS) viene calcolata tra le distribuzioni di token scalate per temperatura dello strato intermedio ($p^l_t$) e dello strato finale ($q^L_t$).
   • Pesatura: La perdita è pesata dall'vantaggio GRPO firmato ($A_t$). I token ad alto vantaggio attirano lo strato intermedio verso la distribuzione dello strato finale, mentre i token a basso vantaggio lo allontanano per evitare di rafforzare continuazioni scadenti.
   • Formula: $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. Allineamento dell'Attenzione ("Dove Guardare"):

   • Obiettivo: Imporre pattern di attenzione coerenti e comportamenti di instradamento delle evidenze.
   • Meccanismo: La divergenza JS viene calcolata tra i pesi di attenzione dello strato intermedio e dello strato finale su un insieme campionato di chiavi causali (combinando finestre locali e posizioni globali a scatti per gestire la memoria).
   • Pesatura: Simile all'allineamento dei logit, la perdita è pesata dall'vantaggio firmato. Vantaggi positivi incoraggiano lo strato intermedio a prestare attenzione alle stesse evidenze dello strato finale; vantaggi negativi riducono l'accordo con pattern a bassa ricompensa.
   • Formula: $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

Obiettivo di Addestramento

La perdita totale combina la perdita GRPO standard con le due perdite di allineamento interno:
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
dove $\lambda_{\text{think}}$ e $\lambda_{\text{attn}}$ controllano la forza della distillazione.

3. Contributi Chiave

1. Nuovo Paradigma: Introduzione dell'Auto-distillazione Interna On-Policy, un paradigma di ragionamento RL che sfrutta i propri calcoli interni del modello per la supervisione mantenendo una singola politica di azione.
2. Framework OISD: Un'implementazione specifica che utilizza lo strato finale come insegnante distaccato e uno strato intermedio come studente, trasferendo segnali predittivi on-policy senza informazioni privilegiate esterne.
3. Meccanismi di Allineamento Duale: Sviluppo di due obiettivi complementari:
   • Allineamento dei Logit: Supervisiona il processo di ragionamento ("come pensare").
   • Allineamento dell'Attenzione: Supervisiona l'instradamento delle evidenze ("dove guardare").
4. Validazione Empirica: Dimostrazione di miglioramenti sostanziali e coerenti rispetto a forti baseline di ragionamento RL (inclusi GRPO e BuPO) su quattro benchmark di ragionamento matematico.

4. Risultati Sperimentali

Gli autori hanno valutato OISD sui modelli Qwen3-4B, Qwen3-8B e OctoThinker su quattro benchmark: AMC23, MATH500, AIME24 e AIME25.

• Guadagni di Prestazione:
  • Su Qwen3-4B, OISD ha ottenuto un punteggio medio di 64.75, superando GRPO (55.08) di 9.67 punti e BuPO (58.51) di 6.24 punti.
  • I miglioramenti sono stati particolarmente significativi sui benchmark AIME (ad esempio, +14.31 punti su GRPO su AIME24).
  • Su Qwen3-8B, OISD ha ottenuto la migliore media complessiva di 66.72, superando GRPO (64.23) e BuPO (66.36).
• Analisi Pass@K: OISD ha costantemente ottenuto tassi Pass@K più alti per tutti i valori di $K$ rispetto a GRPO, indicando un miglioramento nella qualità complessiva dell'insieme di soluzioni campionate, non solo nell'accuratezza media.
• Studi di Ablazione:
  • Sia l'allineamento dei Logit che quello dell'Attenzione hanno contribuito positivamente. L'allineamento dei Logit ha mostrato il guadagno più grande su AIME24, mentre l'allineamento dell'Attenzione ha performato leggermente meglio in media.
  • La combinazione di entrambi ha prodotto i migliori risultati, confermando la loro natura complementare.
• Confronto con OPSD: OISD (64.75) ha superato significativamente OPSD (56.63), un metodo all'avanguardia di auto-distillazione on-policy che si affida a soluzioni ground-truth privilegiate. Ciò evidenzia l'efficacia della supervisione puramente interna di OISD.
• Analisi Qualitativa:
  • Dinamiche di Addestramento: OISD ha mostrato un miglioramento costante della ricompensa e un iniziale aumento dell'entropia (suggerendo un'esplorazione più ampia) prima di stabilizzarsi.
  • Visualizzazione dell'Attenzione: Gli strati intermedi in OISD si sono allineati gradualmente all'attenzione dello strato finale sui token di evidenze critici (ad esempio, punti finali, supremo) durante l'addestramento.
  • Visualizzazione Logit-Lens: OISD ha mostrato un allineamento più precoce tra previsioni intermedie e finali su token critici per il ragionamento rispetto a GRPO, che spesso produceva previsioni intermedie frammentate.

5. Significato e Affermazioni

Il documento afferma che OISD avvia l'indagine sull'auto-distillazione interna on-policy. Il suo significato principale risiede in:

• Eliminazione della Dipendenza Esterna: Fornisce supervisione interna densa senza richiedere insegnanti esterni, prompt privilegiati o soluzioni ground-truth durante l'addestramento.
• Coerenza della Politica: Preserva una singola politica di azione sia per la generazione dei rollouts che per l'ottimizzazione, evitando i problemi di discrepanza di distribuzione dei metodi off-policy.
• Ragionamento Potenziato: Allineando le rappresentazioni intermedie con il ragionamento e i pattern di attenzione dello strato finale, OISD incoraggia il modello a imparare rappresentazioni intermedie più forti e globalmente coerenti.

Gli autori notano limitazioni, inclusa la focalizzazione sul ragionamento matematico con risposte verificabili e la mancanza di uno studio approfondito sulla selezione ottimale degli strati. Suggeriscono che il lavoro futuro dovrebbe esplorare la generalizzazione a compiti di ragionamento più complessi e la supervisione interna multistrato.