SALSA: Aprendizaje de Alineación basado en Sopas para una Adaptación más Fuerte en 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

Resumen

En el desarrollo de Modelos de Lenguaje a Gran Escala (LLM), el Aprendizaje por Refuerzo a partir de Retroalimentación Humana (RLHF) es crucial para alinear los modelos con los valores y preferencias humanas. RLHF tradicionalmente se basa en la divergencia Kullback-Leibler (KL) entre la política actual y una política inicial congelada como referencia, la cual se agrega como una penalización en algoritmos de optimización de políticas como la Optimización de Políticas Proximales (PPO). Si bien esta restricción evita que los modelos se desvíen demasiado del punto de control inicial, limita la exploración del paisaje de recompensas, reduciendo la capacidad del modelo para descubrir soluciones de mayor calidad. Como resultado, la optimización de políticas a menudo queda atrapada en una región estrecha del espacio de parámetros, lo que conduce a un alineamiento y rendimiento subóptimos. Este documento presenta SALSA (Aprendizaje de Alineamiento Basado en Sopas para una Adaptación más Fuerte), un enfoque novedoso diseñado para superar estas limitaciones mediante la creación de un modelo de referencia más flexible y mejor ubicado a través del promedio de pesos en el espacio de pesos de dos modelos supervisados finamente ajustados (SFT) independientes. Esta sopa de modelos permite una mayor desviación en la divergencia KL y explorar una región prometedora del espacio de soluciones sin sacrificar la estabilidad. Al aprovechar este modelo de referencia más robusto, SALSA fomenta una mejor exploración, logrando recompensas más altas y mejorando la robustez del modelo, la generalización fuera de distribución y el rendimiento. Validamos la efectividad de SALSA a través de experimentos extensos en modelos abiertos populares (Llama2-7B, Mistral-7B y Gemma-2B) en varios benchmarks (MT-Bench, Arena-Hard, UltraFeedback), donde consistentemente supera a PPO fomentando una exploración más profunda y logrando un alineamiento superior en 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.