SALSA: 강화 학습에서 더 강력한 적응을 위한 수프 기반 정렬 학습

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(Reinforcement Learning from Human Feedback)는 현재 정책과 초기 정책 사이의 Kullback-Leibler(KL) 발산을 참조로 삼는 전통적인 방법을 사용하여 정책 최적화 알고리즘인 Proximal Policy Optimization(PPO)에서 패널티로 추가됩니다. 이 제약은 모델이 초기 체크포인트에서 너무 멀리 벗어나지 못하도록 방지하지만, 보상 지형의 탐색을 제한하여 모델이 더 높은 품질의 솔루션을 발견하는 능력을 감소시킵니다. 결과적으로, 정책 최적화는 매개 변수 공간의 좁은 영역에 갇히게 되어 최적이 아닌 정렬과 성능을 유발합니다. 본 논문에서는 SALSA(Soup-based Alignment Learning for Stronger Adaptation)를 제안하여 이러한 제한을 극복하기 위한 혁신적인 방법을 소개합니다. 이 방법은 두 개의 독립적인 지도 학습된(SFT) 모델의 가중치 공간 평균화를 통해 더 유연하고 적절한 참조 모델을 생성합니다. 이 모델 수프는 KL 발산에서 더 큰 편차를 허용하고 안정성을 희생하지 않고 솔루션 공간의 유망한 영역을 탐색할 수 있게 합니다. 이 더 견고한 참조 모델을 활용함으로써 SALSA는 더 나은 탐색을 유도하여 높은 보상을 달성하고 모델의 견고성, 분포 외 일반화, 성능을 향상시킵니다. 우리는 널리 사용되는 오픈 모델(Llama2-7B, Mistral-7B, Gemma-2B)에 대한 다양한 벤치마크(MT-Bench, Arena-Hard, UltraFeedback)에서 깊은 탐색을 유도하고 LLM에서 우수한 정렬을 달성함으로써 PPO를 일관되게 능가하는 SALSA의 효과를 포괄적인 실험을 통해 검증합니다.

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.