OPMD: On-Policy Mechanism Distillation oltre la supervisione dell’output

Abstract degli autori

La distillazione on-policy (OPD) migliora i modelli linguistici sottoponendo un modello studente a supervisione su prefissi campionati dalla sua stessa policy, ma la supervisione resta limitata allo spazio degli output. Di conseguenza, l’OPD trasferisce ciò che un modello insegnante predice senza sottoporre a supervisione la computazione intermedia che produce quelle predizioni. Distillare semplicemente le rappresentazioni intermedie, per esempio tramite il matching delle rappresentazioni, non basta: le rappresentazioni dell’insegnante e dello studente possono differire sul piano geometrico e, anche quando gli stati nascosti sono ben allineati, non è detto che siano utilizzabili nella computazione successiva dello studente. Presentiamo On-Policy Mechanism Distillation (OPMD), un metodo che trasferisce la computazione intermedia derivata dall’insegnante tenendo esplicitamente conto della sua compatibilità con lo studente. OPMD costruisce dapprima un modello insegnante compatibile con lo studente, che trasforma le rappresentazioni dell’insegnante in stati intermedi utilizzabili dalla computazione successiva dello studente; poi interiorizza questi stati a partire da traiettorie on-policy. OPMD ottiene la più alta accuratezza media tra le baseline di distillazione valutate in tutte e tre le configurazioni insegnante-studente e si generalizza alla distillazione da più insegnanti e al fine-tuning supervisionato tra famiglie di modelli. Ulteriori analisi mostrano che la supervisione dei meccanismi produce un ragionamento più efficace e modifica in modo sostanziale la computazione intermedia e le strategie di soluzione dello studente: risultati che indicano un trasferimento di informazioni superiore a quello ottenuto con la sola supervisione dell’output. Nel complesso, questi risultati mostrano che la supervisione a livello dei meccanismi è un modo efficace per portare la distillazione on-policy oltre la supervisione dell’output, consentendo agli studenti di apprendere meccanismi interni che possono utilizzare e interiorizzare.