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

TL;DR

SALSA introduces a weight-space averaged reference model for RLHF, enabling larger policy deviations and improved exploration, leading to better alignment, robustness, and out-of-distribution performance in large language models.

Contribution

The paper proposes SALSA, a novel weight-space averaging method for creating a flexible reference model that enhances exploration in RLHF, surpassing traditional KL-based constraints.

Findings

SALSA outperforms PPO on multiple benchmarks.

Models trained with SALSA show improved robustness and generalization.

SALSA enables larger policy deviations without sacrificing stability.

Abstract

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.