Sintesi Tecnica: Verso un'Internalizzazione Robusta nella Distillazione On-Policy

1. Definizione del Problema

Recenti progressi nella distillazione on-policy (OPD) hanno permesso ai modelli student a internalizzare il contesto privilegiato (ad esempio, prompt di sistema, indizi sul compito o descrizioni dello stato del gioco) nei propri parametri, consentendo loro di performare efficacemente anche senza il contesto al momento dell'inferenza. Tuttavia, questo articolo identifica un fenomeno precedentemente non studiato denominato degradazione indotta dal contesto (context-induced degradation).

Contrariamente all'intuzione secondo cui la conoscenza internalizzata dovrebbe rendere il contesto ridondante (un'operazione idempotente), gli autori osservano che la reintroduzione del contesto privilegiato a uno studente distillato spesso degrada le sue prestazioni e infla la lunghezza delle risposte, anche su istanze che il modello risolve correttamente senza contesto. Ciò accade perché la standard OPD ottimizza esclusivamente la fedeltà privilegiata (corrispondenza con l'output del docente condizionato dal contesto), ma non impone vincoli sul comportamento dello studente quando il contesto è presente. Di conseguenza, lo studente non riesce a raggiungere la rimovibilità del contesto (context removability): la proprietà per cui l'output del modello rimane stabile indipendentemente dal fatto che il contesto privilegiato venga fornito o meno.

Gli autori categorizzano le interazioni post-distillazione del contesto in tre regimi:

• Regime A: Lo studente performa bene senza contesto, ma la reintroduzione del contesto causa cali significativi di accuratezza e alti "tassi di danno" (trasformazione di predizioni corrette in errate). Questo è il regime più prevalente nei loro esperimenti.
• Regime B: Il contesto agisce ancora come un supporto utile, indicando un'internalizzazione incompleta.
• Regime C: Lo studente ha internalizzato con successo il contesto, e le prestazioni sono stabili con o senza di esso.

2. Metodologia: No-Context Anchoring (NCA)

Per affrontare la degradazione indotta dal contesto, gli autori propongono il No-Context Anchoring (NCA), un regolarizzatore di coerenza leggero. L'idea centrale è imporre che l'output dello studente condizionato dal contesto ($q_c$) rimanga coerente con il suo output senza contesto ($q_x$), trattando quest'ultimo come un'ancora.

Funzione Obiettivo

La standard OPD minimizza la divergenza KL tra l'output dello studente senza contesto e la distribuzione del docente condizionata dal contesto:
$$\mathcal{L}_{OPD} = \mathbb{E}_{x \sim D} [\text{KL}(q_x \parallel f(y \mid x, c))]$$

La NCA aumenta questo termine con un termine di coerenza:
$$\mathcal{L} = \mathcal{L}_{OPD} + \beta \cdot \mathcal{L}_{NCA}$$
dove $\mathcal{L}_{NCA}$ è definita come:
$$\mathcal{L}_{NCA} \triangleq \text{KL}(\text{sg}[g_\theta(\cdot \mid x)] \parallel g_\theta(\cdot \mid x, c))$$

Scelte Chiave di Design

1. Ancoraggio Stop-Gradient: La vista senza contesto $q_x$ viene passata attraverso un operatore stop-gradient (sg). Ciò assicura che $q_x$ riceva gradienti solo dalla perdita primaria OPD, mentre $q_c$ riceva gradienti solo dalla perdita NCA. Questo elimina la competizione di gradienti tra i due obiettivi.
2. Divergenza KL Forward: Gli autori scelgono $\text{KL}(q_x \parallel q_c)$ (forward KL) invece della reverse KL. Questa scelta è "mode-covering", impedendo a $q_c$ di far collassare i modi di $q_x$. Inoltre, l'aspettativa ricade sotto $q_x$, permettendo il riutilizzo dei campioni on-policy generati per la OPD senza richiedere un nuovo rollout separato da $q_c$.
3. Efficienza Computazionale: La NCA richiede solo un passaggio forward aggiuntivo per ogni step di addestramento (per calcolare $q_c$). I logit per $q_x$ sono staccati dal grafo computazionale e le sequenze on-policy vengono riutilizzate, rendendo l'overhead trascurabile.

3. Configurazione Sperimentale

Gli autori hanno valutato la NCA attraverso 12 configurazioni che spaziano tra:

• Domini: QA con system-prompt (MedMCQA, dataset di sicurezza) e giochi testuali (Sokoban, FrozenLake).
• Modelli: Vari modelli Llama-3 e Qwen (da 1.7B a 8B di parametri), inclusi scenari di auto-distillazione e distillazione cross-model.
• Metriche: Accuratezza/Win Rate con e senza contesto, Tasso di Danno (probabilità che il contesto ribalti una risposta corretta), tassi di Both-Correct/Won, e variazioni relative della lunghezza della risposta.

4. Risultati Chiave

In tutte le 12 configurazioni, la NCA ha dimostrato i seguenti effetti:

• Mitigazione della Degradazione Indotta dal Contesto: La NCA ha ridotto il danno indotto dal contesto in 11 configurazioni su 12. Nei setting del Regime A (dove la degradazione era severa), la NCA ha ridotto significativamente il calo di accuratezza all'atto della reintroduzione del contesto.
• Miglioramento della Stabilità: La NCA ha costantemente aumentato i tassi "BothCorrect" o "BothWon", indicando che il modello mantiene un comportamento corretto sotto entrambe le viste.
• Eliminazione dell'Inflazione della Lunghezza della Risposta: La standard OPD spesso causava la generazione di risposte significativamente più lunghe quando il contesto veniva reintrodotto (fino a un aumento della lunghezza del +39.5%). La NCA ha eliminato efficacementamente questa inflazione nella maggior parte dei casi, riducendo spesso la variazione di lunghezza vicino allo zero.
• Miglioramento delle Prestazioni Senza Contesto: In 7 configurazioni su 12, la NCA ha migliorato l'obiettivo primario di addestramento (accuratezza senza contesto) dello 0.6 - 3.6 punti percentuali. Gli autori attribuiscono ciò al fatto che il vincolo di coerenza agisce come un regolarizzatore che impedisce agli aggiornamenti dei parametri di accomodare comportamenti conflittuali tra le viste.
• Casi Limite: In un setting del Regime B (FrozenLake, cross-model), la NCA ha leggermente aumentato il danno, suggerendo che quando il contesto trasporta ancora informazioni uniche e non ancora internalizzate, forzare una stretta rimovibilità può scambiare parte delle prestazioni condizionate dal contesto con la stabilità.

5. Analisi Meccanicistica

Il paper fornisce un caso studio meccanicistico sul compito Sokoban per spiegare come la NCA ottenga la rimovibilità del contesto:

• Livello di Rappresentazione: L'analisi degli stati nascosti ha rivelato che, sebbene la standard OPD riduca la divergenza tra le viste contesto e senza contesto, rimane un divario evidente negli strati medi e tardivi. La NCA, tuttavia, ha mantenuto una similarità del coseno quasi perfetta ($\ge 0.997$) attraverso tutti gli strati, confermando che la rimovibilità del contesto è ottenuta a livello di rappresentazione.
• Allineamento dei Gradienti: Nella standard OPD, i gradienti per l'obiettivo di distillazione e la vista condizionata dal contesto erano quasi ortogonali. La NCA introduce una nuova direzione del gradiente che bilancia attivamente le forze del leva del contesto per la distillazione e del rimanere invarianti ad esso.
• Aggiornamenti dei Parametri: Sia la OPD che la NCA concentrano gli aggiornamenti dei parametri negli strati superficiali (particolarmente $k\_proj$ e $q\_proj$). Tuttavia, gli aggiornamenti specifici guidati dalla NCA sono distinti da quelli della OPD, suggerendo che la rimovibilità del contesto è ottenuta principalmente attraverso modifiche alle rappresentazioni della conoscenza negli strati superficiali che la sola OPD non produce.

6. Significato e Rivendicazioni

Gli autori sostengono che il loro lavoro sposti il paradigma di valutazione per la distillazione del contesto privilegiato. Argomentano che un'internalizzazione robusta richiede non solo il matching del comportamento del docente (fedeltà privilegiata) ma anche la garanzia della rimovibilità del contesto.

• Novità: Questo è il primo studio che identifica e affronta il fenomeno per cui la reintroduzione del contesto internalizzato degrada le prestazioni del modello.
• Praticità: Il metodo proposto (NCA) è semplice, richiede solo un passaggio forward e un regolarizzatore di coerenza standard, pur risolvendo efficacemente il problema della degradazione senza sacrificare gli obiettivi primari di distillazione.
• Generalizzabilità: Il metodo migliora le prestazioni attraverso diverse famiglie di modelli, compiti e impostazioni di distillazione (sia self che cross-model).

Il paper conclude esortando la comunità a valutare i modelli distillati non solo per la loro prestazione senza contesto, ma anche per la loro robustezza alla reintroduzione del contesto, assicurando che la conoscenza internalizzata non diventi un peso quando il contesto viene accidentalmente o intenzionalmente fornito.