Trust-Region Behavior Blending for On-Policy Distillation

Sintesi Tecnica: Trust-Region Behavior Blending per la Distillazione On-Policy

Problema

La distillazione on-policy (OPD) affronta il disallineamento dei prefissi (prefix mismatch) inerente alla distillazione offline, addestrando un modello studente su prefissi campionati dalla sua attuale policy mentre si cerca di eguagliare un modello docente più forte. Tuttavia, questo approccio introduce una fragilità critica nelle fasi iniziali dell'addestramento: i modelli studenti deboli generano spesso prefissi di bassa qualità. Quando la distribuzione delle traiettorie dello studente è scarsa, la supervisione del docente viene applicata a prefissi di bassa qualità o "inutilizzabili", limitando l'efficacia del segnale di apprendimento.

I tentativi precedenti di stabilizzare questo processo spesso prevedono lo spostamento della raccolta in modalità off-policy (ad esempio, iniettando token del docente o cambiando le distribuzioni target). Al contrario, le pure estensioni (rollout) dello studente preservano la distribuzione ma soffrono di una scarsa qualità iniziale. La sfida centrale è migliorare la qualità dei primi rollout senza spostare permanentemente il processo di raccolta in modalità off-policy o alterare l'obiettivo fondamentale di distillazione per singolo prefisso.

Metodologia: Trust-Region Behavior Blending (TRB)

Gli autori propongono il Trust-Region Behavior Blending (TRB), un metodo di warmup progettato per controllare la policy di comportamento durante la fase iniziale di raccolta dei rollout. TRB opera sul principio di trovare una policy di comportamento che sia il più vicino possibile al docente, ma rigorosamente vincolata entro una "regione di fiducia" (trust region) centrata sulla attuale policy dello studente.

1. Ottimizzazione per Singolo Prefisso
Ad ogni prefisso di generazione $h$, TRB definisce una policy di comportamento $\mu^*(\cdot | h)$ risolvendo un problema di ottimizzazione vincolata. L'obiettivo è minimizzare la divergenza di Kullback-Leibler (KL) rispetto alla policy del docente $\pi_T$, soggetta a un vincolo di divergenza KL rispetto alla policy dello studente $\pi_S$:

$$\mu^*(\cdot | h) = \arg \min_{\mu} D_{KL}(\mu \parallel \pi_T) \quad \text{s.t.} \quad D_{KL}(\mu \parallel \pi_S) \le \epsilon$$

Qui, $\epsilon$ rappresenta la deviazione locale consentita dallo studente. Questa formulazione assicura che la distribuzione di campionamento sia attratta verso i token supportati dal docente, ma rimanga confinata entro una regione di fiducia centrata sullo studente.

2. Soluzione in Forma Chiusa
L'ottimizzazione produce una soluzione in forma chiusa in cui la policy di comportamento è un'interpolazione geometrica tra le policy dello studente e del docente:
$$\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta} \pi_T(a | h)^\beta}{Z_\beta(h)}$$
Il coefficiente $\beta^*(h)$ è il valore massimo ammissibile in $[0, 1]$ tale che il vincolo KL sia soddisfatto. Se il docente è già all'interno della regione di fiducia, $\mu^*$ diventa il docente; se $\epsilon=0$, ritorna allo studente.

3. Warmup con Annealing
TRB viene applicato solo durante una fase iniziale di warmup. Il budget KL $\epsilon$ viene ridotto tramite annealing lineare verso lo zero su un orizzonte $K$:
$$\epsilon_k = \epsilon_0 \left(1 - \frac{k}{K}\right)$$
Ciò consente una forte guida del docente all'inizio dell'addestramento per migliorare la qualità del prefisso, tornando gradualmente alla raccolta dei rollout basata puramente sullo studente man mano che lo studente migliora. Fondamentalmente, la perdita di distillazione reverse-KL per singolo prefisso rimane invariata; viene modificata solo la policy di raccolta dei dati.

Contributi Chiave

• Una Nuova Strategia di Warmup: TRB introduce un metodo per controllare il comportamento dei rollout iniziali ottimizzando una policy guidata dal docente sotto un esplicito vincolo KL centrato sullo studente, anziché limitarsi a iniettare token o alterare la funzione di perdita.
• Preservazione delle Proprietà On-Policy: A differenza dei metodi che mescolano persistentemente le policy di docente e studente o cambiano la distribuzione target, TRB assicura che l'addestramento ritorni infine ai rollout puramente on-policy, preservando la distribuzione target per la fase principale di addestramento.
• Formulazione Teorica: Il documento fornisce una soluzione in forma chiusa per la policy di comportamento che bilancia la guida del docente con la stabilità dello studente, giustificata dalla monotonicità della divergenza KL rispetto al coefficiente di interpolazione.

Risultati Sperimentali

Gli autori hanno valutato TRB in due contesti di distillazione del ragionamento matematico:

1. Qwen3-1.7B-Base distillato da Qwen3-8B.
2. Qwen3-0.6B-Base distillato da Qwen3-4B.

Performance:

• TRB ha ottenuto la migliore performance media tra tutti i metodi confrontati in entrambi i contesti, superando vanilla OPD, il blending a $\epsilon$ fisso, Veto, SKD (iniezione intercalata del docente), il warmup della temperatura e il warmup SFT.
• Nel setting Qwen3-1.7B, TRB ha raggiunto un pass@1 medio del 33,2%, rispetto al 32,3% di vanilla OPD.
• Nel setting Qwen3-0.6B, TRB ha raggiunto il 44,4%, rispetto al 44,0% di vanilla OPD.

Analisi:

• Dinamiche di Addestramento Iniziale: L'analisi delle traiettorie (Figura 2) mostra che, sebbene alcuni metodi (come SKD) crescano più rapidamente inizialmente, non necessariamente mantengono la performance finale più elevata. TRB fornisce un "avvio rapido" che porta a risultati finali superiori.
• Entropia e Qualità del Prefisso: Durante la fase di warmup, TRB riduce l'entropia media dei token del docente sui prefissi visitati, indicando una guida di qualità superiore. I test Step-0 (Figura 4) dimostrano che i prefissi generati sotto TRB producono tassi di successo più elevati quando continuati dal docente o dallo studente rispetto ai prefissi di vanilla OPD, suggerendo che TRB sposta con successo la distribuzione degli stati iniziali verso traiettorie più apprendibili.

Significato e Rivendicazioni

Il documento afferma che TRB affronta efficacemente la fragilità della distillazione on-policy iniziale migliorando la qualità della distribuzione iniziale dei dati senza deviare permanentemente dall'obiettivo on-policy.

• Ambito Limitato: Gli autori dichiarano esplicitamente che il loro studio è limitato a due contesti di ragionamento matematico con coppie Qwen3-Base e un protocollo di valutazione basato sulla correttezza. Non affermano che gli schemi di warmup specifici si trasferiscano invariati ad altri domini o gap docente-studente.
• Compromesso sui Costi: Il documento riconosce che TRB aumenta i costi durante la fase di warmup a causa della necessità di decoding online del docente e della co-presenza studente-docente. Tuttavia, sottolineano che questi costi sono temporanei, poiché l'addestramento ritorna al profilo di runtime standard e più efficiente di OPD dopo l'orizzonte di warmup.
• Intuizione Centrale: I risultati suggeriscono che il comportamento off-policy guidato dal docente è più benefico quando i prefissi visitati dallo studente sono significativamente disallineati con il docente. Una volta risolto questo disallineamento, la guida off-policy continua può offrire rendimenti decrescenti, supportando la strategia di applicare l'annealing alla policy di comportamento per tornare allo studente.