Verificare prima di distillare: un filtro per il docente a livello di prompt nella distillazione on-policy

Abstract degli autori

La distillazione on-policy (OPD) accelera il post-training fornendo una supervisione densa a livello di token da un modello docente congelato sulle sequenze generate dal modello studente. Vanilla OPD applica questa supervisione in modo uniforme a tutti i prompt, senza verificare se il docente sia affidabile per ciascuno di essi. Poiché la KL inversa tende a privilegiare una modalità, un docente che sbaglia con sicurezza può indurre un aggiornamento forte ma fuorviante. Gli indicatori basati sulle distribuzioni, come l’entropia o la concordanza delle verosimiglianze tra docente e studente, misurano l’incertezza o la concordanza, ma non verificano direttamente la correttezza dei risultati. Introduciamo Teacher-Gated On-Policy Distillation (TGOPD), basato sul principio che l’affidabilità del docente debba essere verificata a livello di prompt prima di ammettere la supervisione densa. TGOPD stima l’affidabilità usando un piccolo insieme di prove del docente valutate da un verificatore e indirizza ciascun prompt esclusivamente verso l’OPD con supervisione densa se il controllo di affidabilità è superato, oppure verso GRPO basato sul verificatore in caso contrario. Con modelli studenti 4B e 35B, negli ambiti della matematica, del codice e della capacità di seguire istruzioni, TGOPD supera Vanilla OPD in tutte e sei le configurazioni a dominio singolo e ottiene medie più alte su sette benchmark a entrambe le scale nell’addestramento multidominio. Usando per stimare l’affidabilità la capacità del docente che altrimenti resterebbe inutilizzata, TGOPD riduce anche lo spreco di risorse di calcolo sul lato del docente nell’OPD asincrona, portando l’utilizzo delle GPU del nodo del docente dal 9,8% al 78,9% nell’esecuzione a dominio singolo con il modello studente 4B in cui è stato misurato.