SALSA: Обучение выравниванию на основе супа для более сильной адаптации в RLHF

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

November 4, 2024

Авторы: Atoosa Chegini, Hamid Kazemi, Iman Mirzadeh, Dong Yin, Maxwell Horton, Moin Nabi, Mehrdad Farajtabar, Keivan Alizadeh

cs.AI

Аннотация

При разработке больших языковых моделей (LLM) обучение с подкреплением на основе обратной связи от человека (RLHF) играет решающую роль в согласовании моделей с человеческими ценностями и предпочтениями. RLHF традиционно опирается на дивергенцию Кульбака-Лейблера (KL) между текущей политикой и замороженной начальной политикой в качестве эталона, который добавляется в виде штрафа в алгоритмы оптимизации политики, такие как Проксимальная оптимизация политики (PPO). Хотя это ограничение предотвращает отклонение моделей слишком далеко от начальной точки, оно ограничивает исследование пространства вознаграждений, уменьшая способность модели обнаруживать более качественные решения. В результате оптимизация политики часто застревает в узкой области пространства параметров, что приводит к неоптимальному согласованию и производительности. В данной статье представлен метод SALSA (Soup-based Alignment Learning for Stronger Adaptation), новый подход, разработанный для преодоления этих ограничений путем создания более гибкой и лучше расположенной эталонной модели путем усреднения весового пространства двух независимых моделей с применением надзорного обучения с тонкой настройкой (SFT). Эта модельная смесь позволяет более значительное отклонение в дивергенции KL и исследование многообещающей области пространства решений без ущерба стабильности. Используя эту более надежную эталонную модель, SALSA способствует более глубокому исследованию, достижению более высоких вознаграждений и улучшению устойчивости модели, обобщению вне распределения и производительности. Мы подтверждаем эффективность SALSA через обширные эксперименты на популярных открытых моделях (Llama2-7B, Mistral-7B и Gemma-2B) по различным бенчмаркам (MT-Bench, Arena-Hard, UltraFeedback), где он последовательно превосходит PPO, способствуя более глубокому исследованию и достижению превосходного согласования в 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.