SALSA: Suppenbasiertes Alignment-Lernen für eine stärkere Anpassung in RLHF

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

November 4, 2024

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

cs.AI

Zusammenfassung

Bei der Entwicklung von Large Language Models (LLMs) ist das Reinforcement Learning from Human Feedback (RLHF) entscheidend, um Modelle mit menschlichen Werten und Präferenzen in Einklang zu bringen. RLHF basiert traditionell auf der Kullback-Leibler (KL)-Divergenz zwischen der aktuellen Richtlinie und einer eingefrorenen Ausgangsrichtlinie als Referenz, die als Strafe in Richtlinienoptimierungsalgorithmen wie der Proximal Policy Optimization (PPO) hinzugefügt wird. Während diese Einschränkung verhindert, dass Modelle zu weit von der anfänglichen Überprüfung abweichen, begrenzt sie die Exploration der Belohnungslandschaft und verringert die Fähigkeit des Modells, qualitativ hochwertigere Lösungen zu entdecken. Als Ergebnis bleibt die Richtlinienoptimierung oft in einem engen Bereich des Parameterraums stecken, was zu suboptimaler Ausrichtung und Leistung führt. Dieses Papier stellt SALSA (Soup-based Alignment Learning for Stronger Adaptation) vor, einen neuartigen Ansatz, der diese Einschränkungen überwinden soll, indem ein flexibleres und besser positioniertes Referenzmodell durch Gewichtsraum-Mittelung von zwei unabhängigen überwachten feinabgestimmten (SFT) Modellen erstellt wird. Diese Modellsuppe ermöglicht eine größere Abweichung in der KL-Divergenz und die Exploration eines vielversprechenden Bereichs des Lösungsraums, ohne die Stabilität zu beeinträchtigen. Durch die Nutzung dieses robusteren Referenzmodells fördert SALSA eine bessere Exploration, erzielt höhere Belohnungen und verbessert die Robustheit des Modells, die Generalisierung außerhalb der Verteilung und die Leistung. Wir validieren die Wirksamkeit von SALSA durch umfangreiche Experimente an beliebten offenen Modellen (Llama2-7B, Mistral-7B und Gemma-2B) über verschiedene Benchmarks (MT-Bench, Arena-Hard, UltraFeedback), wo es kontinuierlich PPO übertrifft, indem es tiefere Exploration fördert und eine überlegene Ausrichtung in LLMs erreicht.

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.