SALSA: Apprendimento dell'allineamento basato su zuppe per un'adattamento più forte in RLHF

SALSA: Soup-based Alignment Learning for Stronger Adaptation in RLHF

November 4, 2024

Autori: Atoosa Chegini, Hamid Kazemi, Iman Mirzadeh, Dong Yin, Maxwell Horton, Moin Nabi, Mehrdad Farajtabar, Keivan Alizadeh

cs.AI

Abstract

Nello sviluppo dei Large Language Model (LLM), il Reinforcement Learning from Human Feedback (RLHF) è cruciale per allineare i modelli con i valori umani e le preferenze. RLHF tradizionalmente si basa sulla divergenza Kullback-Leibler (KL) tra la policy attuale e una policy iniziale congelata come riferimento, che viene aggiunta come penalità negli algoritmi di ottimizzazione della policy come il Proximal Policy Optimization (PPO). Sebbene questo vincolo impedisca ai modelli di deviare troppo dal checkpoint iniziale, limita l'esplorazione del paesaggio dei premi, riducendo la capacità del modello di scoprire soluzioni di qualità superiore. Di conseguenza, l'ottimizzazione della policy è spesso intrappolata in una regione ristretta dello spazio dei parametri, portando a un allineamento e a una performance subottimali. Questo articolo presenta SALSA (Soup-based Alignment Learning for Stronger Adaptation), un approccio innovativo progettato per superare tali limitazioni creando un modello di riferimento più flessibile e meglio posizionato attraverso la media nello spazio dei pesi di due modelli supervisionati fine-tuned (SFT) indipendenti. Questa "zuppa" di modelli consente una maggiore deviazione nella divergenza KL ed esplora una regione promettente dello spazio delle soluzioni senza sacrificare la stabilità. Sfruttando questo modello di riferimento più robusto, SALSA favorisce una migliore esplorazione, raggiungendo premi più elevati e migliorando la robustezza del modello, la generalizzazione fuori distribuzione e le performance. Convalidiamo l'efficacia di SALSA attraverso ampi esperimenti su modelli aperti popolari (Llama2-7B, Mistral-7B e Gemma-2B) su vari benchmark (MT-Bench, Arena-Hard, UltraFeedback), dove supera costantemente il PPO promuovendo una maggiore esplorazione e ottenendo un allineamento superiore nei LLM.

English

In Large Language Model (LLM) development, Reinforcement Learning from Human Feedback (RLHF) is crucial for aligning models with human values and preferences. RLHF traditionally relies on the Kullback-Leibler (KL) divergence between the current policy and a frozen initial policy as a reference, which is added as a penalty in policy optimization algorithms like Proximal Policy Optimization (PPO). While this constraint prevents models from deviating too far from the initial checkpoint, it limits exploration of the reward landscape, reducing the model's ability to discover higher-quality solutions. As a result, policy optimization is often trapped in a narrow region of the parameter space, leading to suboptimal alignment and performance. This paper presents SALSA (Soup-based Alignment Learning for Stronger Adaptation), a novel approach designed to overcome these limitations by creating a more flexible and better located reference model through weight-space averaging of two independent supervised fine-tuned (SFT) models. This model soup allows for larger deviation in KL divergence and exploring a promising region of the solution space without sacrificing stability. By leveraging this more robust reference model, SALSA fosters better exploration, achieving higher rewards and improving model robustness, out-of-distribution generalization, and performance. We validate the effectiveness of SALSA through extensive experiments on popular open models (Llama2-7B, Mistral-7B, and Gemma-2B) across various benchmarks (MT-Bench, Arena-Hard, UltraFeedback), where it consistently surpasses PPO by fostering deeper exploration and achieving superior alignment in LLMs.