Sull’ottimizzazione delle policy con regolarizzazione KL

Abstract degli autori

Nell’apprendimento per rinforzo (RL) asincrono per agenti basati su grandi modelli linguistici (LLM), una policy viene addestrata su traiettorie generate da un’altra: i rollout provengono da checkpoint non aggiornati e le probabilità del motore di inferenza differiscono da quelle del sistema di addestramento anche a parità di parametri. I rimedi standard consistono nel limitare i rapporti di importanza, introducendo così una distorsione nell’aggiornamento, oppure, come in GRPO, nel campionare un gruppo di risposte per ogni prompt, soluzione costosa quando gli episodi sono lunghi. Proponiamo KL-Regularized Policy Optimization (KLPO), un approccio in cui il regolarizzatore KL è ancorato al campionatore. Il passo di miglioramento regolarizzato ammette così una soluzione di Gibbs in forma chiusa, e KLPO adatta ai minimi quadrati la propria condizione di ottimalità espressa in termini di log-rapporto sulle traiettorie del campionatore stesso: la probabilità del campionatore entra quindi attraverso un log-rapporto e non sono necessari pesi di importanza. Eliminando per profilazione l’intercetta della regressione, la funzione di log-partizione, non calcolabile direttamente, viene sostituita dalla media del segnale rispetto al campionatore più una divergenza KL dal campionatore al sistema di addestramento. Per gli obiettivi di policy mirror descent a livello di token, mostriamo che il gradiente risultante può essere calcolato dai ritorni terminali senza un critico, tramite score centrati rispetto al campionatore o il residuo di una singola traiettoria, anche quando gli output degli strumenti sono stocastici. Dimostriamo inoltre che stime Monte Carlo indipendenti del termine KL mantengono questi gradienti non distorti, ricaviamo l’esatto scarto KL delle approssimazioni top-$K$ e binarie, meno costose, e mostriamo che SPPO, GPO, REBEL e BPO sono casi particolari di KLPO. Il risultato è un aggiornamento senza critico che usa un solo rollout per prompt e non richiede né un normalizzatore appreso né un gruppo di risposte.