Non tutti i token meritano di essere distillati: supervisione selettiva per Direct-OPD

Abstract degli autori

La Direct On-Policy Distillation (Direct-OPD) trasferisce i miglioramenti della politica indotti dall’apprendimento per rinforzo da un modello piccolo a un modello studente più grande, usando come supervisione dettagliata sui rollout dello studente stesso il log-rapporto, calcolato a livello di token, tra i checkpoint precedenti e successivi all’apprendimento per rinforzo. Questo trasferimento premia il cambiamento della politica in ogni stato, ma il log-rapporto misura soltanto una variazione relativa: può restare invariato anche quando tende a zero la massa di probabilità che entrambi i checkpoint assegnano ai token candidati dello studente. Con una costruzione esatta, mostriamo che la ricompensa di Direct-OPD e il relativo aggiornamento possono rimanere invariati mentre la divergenza di Jensen-Shannon (JSD) e la divergenza KL in entrambe le direzioni tra i checkpoint tendono a zero insieme a questa massa. Osserviamo inoltre che una JSD piccola pone un limite a quanto possa essere cambiato il comportamento del modello docente. Sulla base di questa analisi, proponiamo Selective Supervision for Direct-OPD (S$^2$D-OPD), che ordina gli stati campionati dallo studente in base alla loro JSD tra il modello docente e quello di riferimento ed esclude dalla supervisione Direct-OPD gli stati a bassa divergenza, conservando solo il 10% degli stati per ciascuna risposta con la divergenza più alta. Con due coppie di modelli docenti e quattro modelli studenti da 1,7 a 8 miliardi di parametri, S$^2$D-OPD migliora rispetto alla Direct-OPD con supervisione su tutti gli stati l’accuratezza su dati tenuti da parte nei benchmark AIME e HMMT in sette delle otto configurazioni, e la eguaglia nell’ottava, senza ulteriori passaggi forward.