Sintesi Tecnica: Apprendimento per Rinforzo con Feedback Umano Distribuzionalmente Robusto

Definizione del Problema

L'Apprendimento per Rinforzo con Feedback Umano (RLHF) è diventato una tecnica standard per allineare i Large Language Models (LLM) alle intenzioni umane. Tuttavia, le attuali pipeline RLHF soffrono di una mancanza di robustezza. Quando la distribuzione dei prompt incontrati durante l'inferenza (compiti downstream) differisce significativamente dalla distribuzione del dataset di preferenza utilizzato durante il fine-tuning, le prestazioni del modello spesso degradano. Questo problema è esacerbato da fenomeni come l'over-optimization (sovra-ottimizzazione), dove le policy migliorano le ricompense proxy a scapito delle vere funzioni di ricompensa, e dai cambiamenti distribuzionali causati da ambienti variabili o da derive nelle preferenze umane. Gli approcci precedenti alla robustezza si sono concentrati su specifici modelli di corruzione avversaria, spostamenti del modello di preferenza o metodi euristici come il model merging, che spesso falliscono in presenza di lievi spostamenti distribuzionali.

Metodologia

Gli autori propongono un framework di Ottimizzazione Distribuzionalmente Robusta (DRO) per l'RLHF. L'obiettivo centrale è garantire che i modelli perfezionati (sia i modelli di ricompensa che le policy) mantengano le prestazioni anche quando la distribuzione dei prompt si sposta entro una distanza limitata dalla distribuzione di addestramento. Gli autori formulano questo concetto introducendo un problema di ottimizzazione minimax in cui il learner minimizza la perdita rispetto alla peggiore distribuzione all'interno di una palla di $\phi$-divergenza (specificamente la distanza di Variazione Totale) attorno alla distribuzione empirica di addestramento.

Il documento sviluppa varianti robuste per due pipeline principali di RLHF:

1. Stima della Ricompensa Robusta:

   • Formulazione: Il modello di ricompensa $r$ è addestrato per minimizzare la peggiore attesa della perdita su distribuzioni $D$ tali che $d_{TV}(D, D_{src}) \leq \rho$.
   • Algoritmo: Viene proposto un algoritmo di discesa del gradiente su minibatch (Algoritmo 1). In ogni iterazione, l'algoritmo campiona un minibatch e calcola i pesi ottimali $q^*_t$ che corrispondono allo spostamento della distribuzione nel caso peggiore entro la distanza $\rho$. Il gradiente viene quindi ricalibrato tramite questi pesi ottimali prima di aggiornare i parametri della ricompensa.
   • Convergenza: Sotto l'assunto di una classe di funzioni di ricompensa lineare, gli autori dimostrano che l'algoritmo converge a una soluzione $\epsilon$-ottimale con $T = O(1/\epsilon^2)$ iterazioni e $n = \tilde{O}(1/\epsilon^2)$ campioni per iterazione.

2. Ottimizzazione della Policy Robusta:

   • Formulazione: La policy $\pi$ è ottimizzata per massimizzare la funzione di valore regolarizzata KL nel caso peggiore.
   • Algoritmo: Gli autori propongono un metodo di Natural Policy Gradient (NPG) robusto (Algoritmo 2). Similmente alla stima della ricompensa, esso calcola i pesi del caso peggiore per il minibatch. Fondamentalmente, aggiorna la policy utilizzando una matrice di informazione di Fisher pesata derivata da questi pesi del caso peggiore, anziché dalla standard matrice di Fisher empirica.
   • Convergenza: Assumendo una classe di policy log-lineari e le standard condizioni di concentrabilità, il metodo raggiunge un tasso di convergenza lineare ($T = O(\log(1/\epsilon))$), ma richiede una complessità di campionamento più elevata di $n = \tilde{O}(1/\epsilon^4)$ per iterazione per garantire la convergenza dell'ultimo iterato.

3. Ottimizzazione Diretta delle Preferenze (DPO) Robusta:

   • Formulazione: Una versione distribuzionalmente robusta di DPO è formulata per ottimizzare direttamente la policy senza un esplicito modello di ricompensa.
   • Algoritmo: L'Algoritmo 3 applica la discesa del gradiente stocastica su minibatch con gradienti ricalibrati sulla base dello spostamento della distribuzione nel caso peggiore.
   • Convergenza: L'analisi mostra una complessità di iterazione di $T = O(1/\epsilon^2)$ e una complessità di campionamento di $n = \tilde{O}(1/\epsilon^2)$, simile alla fase di stima della ricompensa robusta.

Contributi Chiave

• Framework Teorico: Il documento stabilisce le prime formulazioni distribuzionalmente robuste sia per la fase di stima della ricompensa che per quella di ottimizzazione della policy nell'RLHF, nonché per il DPO.
• Algoritmi: Introduce algoritmi pratici basati su minibatch per la Stima della Ricompensa Robusta, il NPG Robusto e il DPO Robusto. Questi algoritmi sono progettati per essere minimamente invasivi rispetto alle pipeline esistenti, richiedendo solo il calcolo dei fattori di ricalibrazione (risolti tramite ottimizzazione convessa) e la ricalibrazione del gradiente.
• Garanzie di Convergenza: Gli autori forniscono prove teoriche per la convergenza di questi algoritmi. Stabiliscono limiti sulla complessità di campionamento e di iterazione per modelli di ricompensa lineari e policy log-lineari, dimostrando che la robustezza può essere ottenuta con garanzie dimostrabili.
• Validazione Empirica: I metodi sono valutati su compiti out-of-distribution (OOD) utilizzando il dataset Unified-Feedback per l'addestramento e RewardBench, HHH-Alignment e MT-Bench per la valutazione.

Risultati Sperimentali

Gli autori hanno valutato i loro metodi utilizzando i modelli base Gemma-2B-it e Mistral-7B-Instruct-v0.2.

• Modelli di Ricompensa: Sul benchmark RewardBench, i modelli di ricompensa distribuzionalmente robusti (DR RM) hanno mostrato un miglioramento dell'accuratezza media rispetto ai modelli standard Bradley-Terry (BTL), in particolare quando il parametro di robustezza $\rho$ era impostato a 0.1 o 0.2. Notevolmente, le prestazioni sul sottoinsieme Reasoning (Ragionamento) sono migliorate significativamente (ad esempio, +23.4% su un sottoinsieme con $\rho=0.1$). Tuttavia, le prestazioni sono degradate quando $\rho$ era impostato troppo alto (ad esempio, $\rho \geq 0.6$), indicando un compromesso tra robustezza e adattamento alla distribuzione di addestramento.
• Ottimizzazione della Policy (PPO): Utilizzando i modelli di ricompensa robusti per addestrare le policy tramite PPO, le varianti robuste (DR PPO) hanno generalmente eguagliato o superato le prestazioni del PPO standard sia sui dataset in-distribution (Unified-Feedback) che su quelli OOD (HHH-Alignment, MT-Bench). Un miglioramento marcato è stato osservato nei compiti di ragionamento, con un aumento dell'accuratezza del +10.0% per $\rho=0.1$.
• Ottimizzazione Diretta delle Preferenze (DPO): Il DPO distribuionalmente robusto (DR DPO) ha costantemente migliorato le prestazioni sui dataset OOD rispetto al DPO standard. Per HHH-Alignment, la migliore prestazione è stata ottenuta con $\rho=0.8$, e per MT-Bench con $\rho=0.5$. Il DR DPO ha generalmente superato il DR PPO nei compiti OOD.

Significato e Rivendicazioni

Il documento sostiene che l'ottimizzazione distribuzionalmente robusta fornisca un approccio fondato per migliorare la generalizzazione degli LLM verso compiti non visti e spostamenti distribuzionali. Gli autori sottolineano che i loro metodi richiedono alterazioni minime alle pipeline RLHF esistenti ed efficienti (ad esempio, quelle che utilizzano le librerie Hugging Face), offrendo al contempo garanzie teoriche di convergenza.

Un risultato chiave evidenziato dagli autori è il miglioramento inaspettato nelle capacità di ragionamento sui dataset OOD quando viene imposta la robustezza distribuzionale. Gli autori osservano che non è immediatamente chiaro perché la DRO debba influenzare positivamente il ragionamento, suggerendo che questo fenomeno meriti ulteriori indagini sulle connessioni tra ragionamento e robustezza distribuzionale.

Questo lavoro si posiziona come un passo verso l'allineamento effettivo degli LLM con le intenzioni umane in ambienti reali dinamici, dove le distribuzioni dei dati non sono statiche, andando oltre i limiti del classico RLHF che spesso fatica a generalizzare oltre i specifici esempi di preferenza visti durante l'addestramento.