Riepilogo Tecnico: Non Tutti i Disaccordi Sono Apprendibili: Insegnabilità dei Token nella Distillazione On-Policy

Enunciato del Problema

La distillazione on-policy (OPD) addestra un modello linguistico studente sui propri rollouts utilizzando una supervisione a livello di token da una politica insegnante. Sebbene i recenti metodi OPD selettivi tentino di migliorare l'efficienza dando priorità a token "salienti" — come quelli con alta entropia o alto disaccordo insegnante-studente (divergenza KL) — questi approcci si basano sull'assunzione che un alto disaccordo equivalga a un alto valore di apprendimento.

Il documento identifica un difetto critico in tale assunzione: il disaccordo KL grezzo è un proxy grossolano per l'apprendimento. Esso confonde due scenari qualitativamente diversi:

1. Disaccordo Apprendibile: L'insegnante assegna massa probabile correttiva a token già presenti nel supporto predittivo top-$K$ corrente dello studente.
2. Disaccordo Incompatibile: L'insegnante assegna massa a token ben al di fuori del supporto corrente dello studente. Sebbene ciò crei un ampio divario KL, lo studente non può facilmente assorbire questa correzione, portando a effetti di apprendimento locale deboli o negativi.

Il problema centrale è distinguere quali segnali a livello di token dell'insegnante siano effettivamente apprendibili (cioè capaci di indurre un miglioramento locale) rispetto a quelli che sono meramente salienti o incompatibili.

Metodologia

1. Diagnostica a Contesto Fisso

Per isolare il valore di apprendimento di specifici token dal rumore dei rollouts e dagli spostamenti del contesto a valle, gli autori introducono una diagnostica a contesto fisso.

• Protocollo: I prefissi generati dallo studente vengono raccolti e congelati come banca di contesti. Sia il modello studente iniziale (pre-aggiornamento) che quello addestrato vengono risvalutati rispetto alla stessa distribuzione insegnante su questi contesti fissi.
• Metrica: Il Guadagno del Token a Contesto Fisso ($G^{\text{fix}}_{i,t}$) è definito come la riduzione della divergenza KL tra insegnante e studente dopo l'addestramento, misurata sullo stesso prefisso.
• Obiettivo: Ciò permette agli autori di correlare i segnali pre-aggiornamento (come il KL grezzo) con il reale miglioramento locale, controllando la varianza di campionamento.

2. Formalizzazione dell'Insegnabilità dei Token

Utilizzando la diagnostica, gli autori decompongono il disaccordo in due componenti basate sulla compatibilità del supporto locale:

• Supporto Locale ($S^S_t(K)$): L'insieme dei token top-$K$ dello studente.
• Massa di Compatibilità ($C_t$): La massa di probabilità totale che l'insegnante assegna al supporto locale dello studente $S^S_t(K)$.
• Decomposizione:
  • Disaccordo Apprendibile ($D^L_t$): Alto disaccordo in cui la correzione dell'insegnante è allineata con il supporto dello studente (Alta $C_t$).
  • Disaccordo Incompatibile ($D^I_t$): Alto disaccordo in cui la correzione dell'insegnante cade fuori dal supporto dello studente (Bassa $C_t$).

Il Punteggio di Insegnabilità è definito come il prodotto del disaccordo normalizzato e della compatibilità:
$$s^{\text{teach}}_t = \tilde{D}_t \cdot \tilde{C}_t$$
dove $\tilde{D}_t$ è il disaccordo normalizzato e $\tilde{C}_t$ è la massa di compatibilità normalizzata.

3. OPD Consapevole dell'Insegnabilità (TA-OPD)

Sulla base della constatazione che $D^L_t$ predice il guadagno a contesto fisso meglio del KL grezzo, gli autori propongono TA-OPD, un metodo leggero di selezione delle posizioni dei token:

• Selezione: Per un dato budget (es. 5% o 10% dei token), TA-OPD seleziona le top-$N$ posizioni con i punteggi di insegnabilità più alti ($s^{\text{teach}}_t$).
• Implementazione: Richiede solo le probabilità dei token insegnante e studente (specificamente i log-prob top-$K$) disponibili durante l'OPD standard. Non richiede modelli di ricompensa, verificatori o etichette aggiuntive.
• Contrasto: A differenza dei metodi precedenti che selezionano in base all'entropia o alla divergenza grezza (TIP), TA-OPD filtra esplicitamente per correzioni "allineate al supporto".

Contributi Chiave

1. Diagnostica a Contesto Fisso: Un metodo per misurare la riduzione KL nello stesso contesto, collegando i segnali di token pre-aggiornamento al miglioramento OPD locale senza confondere il rumore dei rollouts.
2. Formalizzazione dell'Insegnabilità dei Token: La dimostrazione che il disaccordo KL grezzo confonde segnali apprendibili e incompatibili. Il documento dimostra che il disaccordo allineato al supporto è il predittore affidabile di un apprendimento utile.
3. Algoritmo TA-OPD: Un metodo di selezione pratico e leggero che applica la perdita OPD solo alle posizioni ad alta insegnabilità, migliorando l'efficienza senza strumenti esterni.
4. Validazione Empirica: Esperimenti estesi che mostrano come selezionare in base all'insegnabilità produca prestazioni a valle migliori rispetto alla selezione basata sulla salienza, anche con budget di token significativamente ridotti.

Risultati Sperimentali

Gli autori hanno valutato TA-OPD su quattro configurazioni insegnante-studente che coinvolgono modelli Qwen3 (4B, 8B, 14B) e Qwen2.5 (3B), utilizzando benchmark come AIME, GPQA-Diamond, HumanEval e MATH-500.

• Prestazioni vs OPD Completo: In molte configurazioni, TA-OPD ha eguagliato o superato l'OPD a token completo utilizzando solo il 5% dei token. Ad esempio, nella configurazione Qwen3-4B $\to$ 1.7B, TA-OPD ha ottenuto un punteggio medio di 44.89 rispetto ai 42.37 dell'OPD Completo, nonostante un budget del 10% (superando l'OPD completo in metriche specifiche).
• Confronto con le Baseline: TA-OPD ha costantemente superato le baseline basate sull'entropia, sulla divergenza e TIP (Token Importance in OPD).
  • Nella configurazione cross-backbone (DeepSeek-R1-Distill-Qwen-14B $\to$ Qwen2.5-3B), TA-OPD ha migliorato il punteggio medio a 30.62, superando significativamente il modello Base (29.98) e l'OPD Completo (28.76).
• Sensibilità al Budget: I risultati indicano che le prestazioni OPD sono limitate dalla qualità dei token piuttosto che puramente dal budget. Aumentare il budget dei token oltre il 10% spesso non ha prodotto ulteriori guadagni e talvolta ha degradato le prestazioni, rafforzando l'idea che non tutta la supervisione aggiuntiva sia benefica.
• Robustezza: Il positivo "divario di guadagno ad alto supporto" (la differenza di prestazioni tra token ad alta e bassa insegnabilità) è rimasto coerente attraverso diversi contesti (matematica, GSM8K-COT), scale insegnante (8B, 14B) e dimensioni del supporto ($K=8, 16, 32$).

Significato e Affermazioni

Il documento riformula la distillazione on-policy selettiva non come un problema di selezione di token "salienti" o "incerti", ma come un problema di identificazione di segnali insegnanti localmente apprendibili.

• Affermazione Centrale: I guadagni derivanti dall'OPD dipendono dalla compatibilità tra la correzione dell'insegnante e il supporto predittivo corrente dello studente. L'alta divergenza KL da sola è insufficiente; l'insegnante deve assegnare massa a token che lo studente sta già considerando plausibili.
• Impatto Pratico: TA-OPD offre un modo per ridurre drasticamente il costo computazionale della distillazione (filtrando i token incompatibili) mantenendo o migliorando le prestazioni del modello. Ciò è ottenuto senza l'overhead di modelli di ricompensa o verificatori, rendendolo applicabile alle pipeline di addestramento standard.
• Limitazioni: Gli autori notano che la loro analisi si concentra sul ragionamento ad alta densità matematica e sui modelli della famiglia Qwen. Sebbene la diagnostica sia agnostica rispetto al modello, è necessaria una validazione più ampia su dialoghi, codice e backbone non-Qwen per confermare la generalizzabilità. Inoltre, il "budget dei token" si riferisce alla densità di supervisione, non necessariamente all'accelerazione del tempo reale, poiché il metodo seleziona posizioni anziché potare il calcolo del transformer.

In conclusione, il lavoro suggerisce che una distillazione efficace richiede una strategia di supervisione consapevole della compatibilità, dove lo studente impara meglio dai segnali insegnanti che sono sia correttivi che localmente raggiungibili.