SAKI: supervisione del modello docente instradata tramite accoppiamento massimale per la distillazione on-policy

Abstract degli autori

La distillazione on-policy (OPD) riduce la discrepanza tra gli stati incontrati durante l’addestramento e quelli incontrati durante il test addestrando un modello studente sulle traiettorie che genera, ma gli studenti più deboli possono raggiungere prefissi non allineati con il modello docente, nei quali la supervisione è meno rappresentativa. Presentiamo SAKI (Supervision Allocation with KL-constrained Interpolation), che combina un rollout guidato dal docente e vincolato dalla KL con l’accoppiamento massimale, e riutilizza gli eventi di accettazione e correzione effettivamente verificatisi per indirizzare la supervisione a livello di token. Nelle posizioni accettate resta la supervisione basata sulla KL inversa per il token campionato, mentre le posizioni corrette ricevono una supervisione diretta sul token a cui il docente assegna la probabilità più alta. Con l’accoppiamento massimale, la probabilità di correzione è esattamente TV(p_t, q_t): lo stesso raggio della trust region controlla quindi la deviazione del rollout e pone un limite superiore alla frequenza degli interventi e della supervisione specializzata. Implementiamo inoltre un verificatore speculativo integrato nel motore che preserva la distribuzione delle traiettorie exact-q e le proprietà dell’accoppiamento, aumentando di 4,22 volte la velocità dei rollout a parità di carico di lavoro. Su sette benchmark di ragionamento matematico, SAKI migliora la baseline guidata dal docente a parità di condizioni in Mean@8 e Pass@8, sia per gli studenti da 1,7B sia per quelli da 0,6B. I controlli sulla collocazione e l’analisi con prefissi fissi forniscono ulteriori elementi a sostegno dell’uso delle correzioni per attivare una supervisione adattiva ai conflitti.