Il modello insegnante è una direzione, non una destinazione: estrapolare i residui delle rappresentazioni indotti dall’RL nella distillazione on-policy

Abstract degli autori

La distillazione on-policy (OPD) addestra un modello studente a riprodurre le distribuzioni del token successivo del modello insegnante sulle traiettorie generate dallo studente stesso e ha prodotto notevoli miglioramenti empirici. Alcune varianti generalizzate consentono allo studente di superare l’insegnante estrapolando una ricompensa implicita nello spazio degli output. La testa del modello linguistico, tuttavia, attenua questa variazione in modo anisotropo: gran parte della variazione codificata negli stati nascosti dell’insegnante raggiunge i logit con una frazione ridotta del suo peso, mentre i rapporti tra le log-probabilità dei token campionati, su cui si basa l’estrapolazione nello spazio degli output, introducono rumore che l’estrapolazione amplifica, rendendo instabile l’addestramento. Osserviamo che l’apprendimento per rinforzo (RL) sposta le rappresentazioni interne di un modello rispetto al suo checkpoint di base e che la direzione di questo spostamento può essere misurata a ogni livello. Sulla base di questa osservazione, proponiamo RIDE (RL-Induced Direction Extrapolation), che estrapola la variazione indotta dall’RL direttamente nello spazio delle rappresentazioni: a ogni livello e per ogni posizione del token, RIDE calcola la differenza tra l’insegnante e il suo checkpoint precedente all’RL e, tramite regressione, avvicina gli stati nascosti dello studente a valori obiettivo situati oltre l’insegnante lungo questa differenza. Data una traiettoria campionata, questa regressione equivale a massimizzare una ricompensa direzionale lineare definita dalla differenza, con una penalità quadratica centrata sull’insegnante: ciò rende esplicito il modo in cui l’obiettivo sposta lo studente nella direzione indotta dall’RL, limitandone al contempo lo scostamento dall’insegnante. Su quattro coppie modello di base/insegnante addestrato con RL, che comprendono scale, architetture e linee di preaddestramento diverse, RIDE si avvicina all’insegnante addestrato con RL o lo supera in ogni coppia ed è l’unico metodo a farlo anche in media. Inoltre, supera sistematicamente l’estrapolazione nello spazio degli output, che peggiora le prestazioni dello studente ogni volta che l’insegnante è vicino al proprio modello di base. Pagina del progetto: https://github.com/xixixixixxxx/RIDE.