SALSA: Aprendizado de Alinhamento Baseado em Sopa para uma Adaptação mais Forte em RLHF

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

November 4, 2024

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

cs.AI

Resumo

No desenvolvimento de Modelos de Linguagem de Grande Escala (LLM), o Aprendizado por Reforço a partir do Feedback Humano (RLHF) é crucial para alinhar os modelos com os valores e preferências humanas. O RLHF tradicionalmente se baseia na divergência Kullback-Leibler (KL) entre a política atual e uma política inicial congelada como referência, que é adicionada como uma penalidade em algoritmos de otimização de política como a Otimização de Política Proximal (PPO). Embora essa restrição evite que os modelos se desviem muito do ponto de verificação inicial, ela limita a exploração da paisagem de recompensas, reduzindo a capacidade do modelo de descobrir soluções de maior qualidade. Como resultado, a otimização de política frequentemente fica presa em uma região estreita do espaço de parâmetros, levando a um alinhamento e desempenho subótimos. Este artigo apresenta o SALSA (Aprendizado de Alinhamento à Base de Sopa para uma Adaptação Mais Forte), uma abordagem inovadora projetada para superar essas limitações, criando um modelo de referência mais flexível e melhor localizado por meio da média de espaço de pesos de dois modelos independentes finamente ajustados supervisionados (SFT). Essa sopa de modelos permite uma maior divergência na divergência KL e explorar uma região promissora do espaço de soluções sem sacrificar a estabilidade. Ao aproveitar esse modelo de referência mais robusto, o SALSA promove uma melhor exploração, alcançando recompensas mais altas e melhorando a robustez do modelo, generalização fora da distribuição e desempenho. Validamos a eficácia do SALSA por meio de experimentos extensivos em modelos abertos populares (Llama2-7B, Mistral-7B e Gemma-2B) em vários benchmarks (MT-Bench, Arena-Hard, UltraFeedback), onde consistentemente supera o PPO ao promover uma exploração mais profunda e alcançar um alinhamento superior em LLMs.

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.