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

## Abstract

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 ( S oup-based A lignment L earning for S tronger A daptation), 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.

## 1 Introduction

Large language models (LLMs) have revolutionized natural language processing (NLP) by demonstrating remarkable capabilities in understanding and generating human language. These models, powered by vast amounts of data and advanced neural architectures, have set new benchmarks in various NLP tasks, from machine translation to conversational agents. Despite these advancements, aligning LLMs with human values and preferences remains a significant challenge. Misalignment can lead to undesirable behaviors, including generating biased or inappropriate content, which undermines the reliability and safety of these models (Kenton et al., ; Bender et al., ; Bommasani et al., ; Gehman et al., ).

Reinforcement Learning from Human Feedback (RLHF) has become a promising technique for aligning large language models (LLMs) with human preferences. By fine-tuning LLMs based on human feedback, RLHF guides models towards more human-aligned behaviors, improving truthfulness, helpfulness, and harmlessness while maintaining the generation of high-probability, correct answers (Christiano et al., ).
Reward-based RLHF methods utilize a reward model to determine the reward for a given (prompt, response) pair. The policy model is then optimized by maximizing the average reward through reinforcement learning algorithms like Proximal Policy Optimization (PPO) (Schulman et al., ).
A crucial aspect of RLHF is the use of a reference model to compute the Kullback-Leibler (KL) divergence penalty, which prevents the fine-tuning process from deviating too far from the original model (Ziegler et al., ). This approach ensures that the policy remains close to the initial model, reducing the risk of generating nonsensical responses.

While effective, reliance on a single reference model can be limiting. The KL penalty term constrains the policy model to stay close to the initial supervised fine-tuning (SFT) model, restricting its ability to fully explore the solution space for higher-reward models. This constraint can lead to suboptimal alignment and a lack of robustness in the training process, increasing the risk of generating nonsensical outputs. Ensuring that the reference model is already positioned in a robust space can mitigate this issue, allowing for more confident exploration without compromising output quality.

To address this limitation, we propose SALSA. It integrates a “model soup” as the reference model within the RLHF framework. A model soup is constructed by performing weight-space averaging of multiple independently supervised fine-tuned models that demonstrate comparable performance. This method leverages the principle that fine-tuned models from the same pre-trained initialization often reside in a shared low-error basin in the loss landscape, enabling effective weight interpolation without compromising accuracy. As evidenced by Wortsman et al. (), this approach results in significant improvements in both in-distribution and out-of-distribution generalization. The key advantage of model soup lies in its ability to harness the complementary strengths of diverse models, reducing variance and improving robustness, while maintaining computational efficiency. This makes the “model soup” reference model a superior choice for improving the stability and reliability of the RLHF training process compared to relying on a single reference model.

In this paper, we demonstrate the effectiveness of SALSA through comprehensive experiments. We apply SALSA to Llama2-7B, Mistral-7B, and Gemma-2B and benchmark the results against standard evaluation datasets, including MT-Bench, Arena-Hard, and UltraFeedback—the latter being used for RLHF training in our experiments. Our findings reveal that weight space averaging is a straightforward yet effective approach for aligning LLMs with human preferences, and enhancing their performance on real-world-like datasets.

In particular our contributions are following:

- •
We demonstrate that the reward in the region near the model soup is inherently superior to that of the original SFT model. The improvement in reward is a newly observed phenomenon and is complementary to the improvement in accuracy of model soups. We further show having model soups as reference point of RLHF results in higher reward outcomes.
(Section 4.2).
- •
Drawing from these observations, we propose SALSA, a novel approach for implementing RLHF that utilizes the model soup as the reference model.
- •
We perform a comprehensive evaluation across diverse benchmarks and models, demonstrating that SALSA consistently outperforms PPO (Section 4.3).

The remainder of this paper is organized as follows: Section 2 reviews related work on RLHF and model averaging techniques. Section 3 details our methodology for creating and integrating the model soup in the RLHF process. Sections 4.1 through 4.4 present the experimental results and analysis of model averaging for better alignment.

Figure: (a) Llama2-7B
Refer to caption: x1.png

## 2 Related Works

### 2.1 Model Soup

Model soups build on the findings of (Neyshabur et al., ), which showed that independently fine-tuned models often lie within the same loss basin, allowing their weights to be successfully combined through interpolation. (Wortsman et al., ) extended this, demonstrating that averaging fine-tuned model weights can improve performance compared to using a single model. Unlike traditional ensembling, model soups require no extra memory or inference time, while still enhancing robustness. This method has shown consistent performance gains across NLP and image classification tasks (Wortsman et al., ; Izmailov et al., ).

### 2.2 Reinforcement Learning from Human Feedback (RLHF)

Large language models (LLMs) have achieved significant success across tasks, thanks to fine-tuning methods like Supervised Fine-Tuning (SFT) and Reinforcement Learning from Human Feedback (RLHF) (Ouyang et al., ; Touvron et al., ; Bai et al., ; Anil et al., ). RLHF has been crucial for aligning models with human preferences, enabling outputs that are more helpful and aligned with societal values (Ziegler et al., ; Stiennon et al., ; Achiam et al., ).

RLHF methods are divided into reward-based and reward-free approaches. Reward-based methods rely on training a reward model to guide policy optimization, typically with Proximal Policy Optimization (PPO) (Schulman et al., ; Gao et al., ). PPO is widely used for its balance of stability and exploration (Ouyang et al., ). Studies have also examined the importance of hyperparameter tuning and reward model quality to avoid overoptimization (Casper et al., ; Zheng et al., ). Reward-free methods, such as Direct Preference Optimization (DPO), bypass the reward model by directly optimizing human preference data (Rafailov et al., ; Liu et al., ; Yuan et al., ). While DPO simplifies training and performs well on tasks like instruction following (Touvron et al., ), it struggles with out-of-distribution data, and its generalization capabilities are limited (Yuan et al., ; Xu et al., ).

In this paper, we focus on Proximal Policy Optimization (PPO) for its demonstrated robustness in managing large models and diverse tasks. Studies such as (Xu et al., ) highlight the limitations of Direct Preference Optimization (DPO), which tends to find biased solutions when confronted with out-of-distribution responses and is highly sensitive to distribution shifts. In contrast, PPO effectively mitigates these challenges through the use of reward models and KL divergence regularization (Ouyang et al., ; Schulman et al., ). This enables PPO to consistently outperform DPO in tasks like dialogue and code generation, proving its reliability across alignment challenges (Xu et al., ).

Recent RLHF research has introduced novel approaches to improve upon traditional methods, particularly in addressing the limitations of using reference models. SimPO (Meng et al., ) eliminates the need for a reference model by optimizing the average log probability of sequences with a reward-free approach. This reduces computational costs and improves memory efficiency over traditional DPO, achieving better performance across benchmarks like AlpacaEval and Arena-Hard. SimPO underscores the limitations of static reference models in DPO and advocates for scalable, reference-free approaches. Conversely, (Gorbatovski et al., ) introduces a dynamic reference model that evolves throughout training using Trust Region methods. This dynamic approach allows the reference model to adapt alongside the policy, preventing constraints imposed by outdated checkpoints and enabling more effective generalization and alignment with human preferences.

The aforementioned papers have identified that the reference model may limit the optimization process in DPO, and they propose innovative solutions to address this issue. Similarly, we aim to solve this problem within the PPO framework. We adopt a model soup approach, which averages the weights of fine-tuned models. This method allows the optimization process in PPO to explore a broader solution space, offering greater flexibility and enhancing alignment performance.

## 3 Method

In this section we present our method, SALSA (Soup-based Alignment Learning for Stronger Adaptation). We begin with a brief overview of Proximal Policy Optimization (PPO), and model soup, followed by a detailed description of our approach.

### 3.1 RLHF

In our main experiments, we focus on reward-based RLHF, in particular Proximal Policy Optimization (PPO). The conventional framework for reward-based RLHF consists of several key stages as follows.

Supervised Fine-Tuning (SFT). The initial stage of alignment involves supervised fine-tuning, where a pre-trained language model is refined using a high-quality instruction dataset.

Reward Model. Reward-based RLHF involves training a reward model, which is typically initialized from the SFT model. In this process, the logit layer of the SFT model is replaced by a new linear layer. This linear layer takes the embedding of the last token and outputs a scalar reward. Higher rewards indicate better samples. For a given prompt $x$ and a pair of responses $y_{w}$ (chosen) and $y_{l}$ (rejected), the loss function is optimized as:

$$ $\mathcal{L}(R_{\theta})=-\log\sigma(R_{\theta}(x,y_{w})-R_{\theta}(x,y_{l})),$ (1) $$

where $R_{\theta}(\cdot,\cdot)$ is the reward model, $\theta$ denotes its parameters, $\sigma$ indicates the sigmoid function.

Policy Training. The last phase of RLHF is dedicated to training the policy model, which is initialized from a reference model, typically the SFT model. Based on a recent study (Xu et al., ), PPO performs better in distribution shifts and results in superior alignment with human preferences across challenging tasks like code generation and dialogue systems. Therefore, we selected PPO as the training algorithm. The goal is to optimize the policy model to maximize the reward for a given prompt $x$ and its generated response $y$, while also minimizing the KL divergence between the policy model and the reference model. The overall loss function for this stage is given by:

$$ $\mathcal{L_{PPO}}(\pi_{\theta})=-R(x,y)+\beta\text{KL}\left(\pi_{\theta}(y\mid x )\,\|\,{\color[rgb]{0.63671875,0.09765625,0.12109375}\definecolor[named]{ pgfstrokecolor}{rgb}{0.63671875,0.09765625,0.12109375}\boldsymbol{\pi_{\text{{ ref}}}}}(y\mid x)\right)$ (2) $$

where $\pi_{\theta}$ is the policy model, $\pi_{ref}$ is the reference model, and $R(.,.)$ is the trained reward model.

### 3.2 Model Soup

The concept of a model soup, introduced in Wortsman et al. (), aims to enhance model performance by averaging the weights of multiple pre-trained networks. This technique combines the parameters of independently trained models to produce a more robust outcome, leveraging the strengths of each. There are three primary strategies for constructing a model soup: uniform, greedy, and learned. In our approach, we employ its most basic strategy, *i.e.*, the uniform method, where the parameters of separate SFTs are averaged. Investigating other mixing strategies for RLHF and comparing them is left as future work.
Formally, we construct our soup model using two SFT models: $\pi_{ref}$, which serves as the initialization of the policy model in the PPO framework, and $\pi_{other}$, an additional SFT model trained on the same data with a different random seed. Their weights denoted as $\theta$ is averaged using a coefficient $\alpha$:

$$ $\theta_{soup}=(1-\alpha)\theta_{ref}+\alpha\theta_{other}$ (3) $$

In our experimental setup, as shown in Figure 4(a), setting $\alpha$ to 0.5 produces the best results. This will be further discussed in the next section, where we also explain our proposed method, SALSA.

### 3.3 SALSA

The KL term in equation 2 ensures that the model remains closely aligned with the reference model. This term is essential because optimizing solely for the reward can result in the generation of nonsensical outputs. However, it also imposes a constraint by preventing the model from deviating significantly from the initial reference model. To address this limitation, we propose replacing the KL term in 2, with the following loss function:

$$ $\mathcal{L_{SALSA}}(\pi_{\theta})=-R(x,y)+\beta\text{KL}\left(\pi_{\theta}(y \mid x)\,\|\,{\color[rgb]{0,0.32421875,0.60546875}\definecolor[named]{ pgfstrokecolor}{rgb}{0,0.32421875,0.60546875}\boldsymbol{\pi_{\text{{soup}}}}} (y\mid x)\right)$ (4) $$

As mentioned in Equation 3, $\pi_{soup}$ refers to a model soup, which is the result of averaging two independently trained supervised fine-tuned models (SFTs), including the reference model. Since the policy model $\pi_{\theta}$ is initialized from the reference model $\pi_{ref}$, substituting the KL term in equation 2 with the KL term in 4 which is the model soup of $\pi_{ref}$ and $\pi_{other}$ allows the policy model to search around averaged model, thereby enabling exploration of a broader and more promising parameter space. The primary distinction between the loss terms in equations 2 and 4 is that the soup model is used in place of the reference model. This substitution keeps our approach straightforward to implement while proving highly effective. Our experiments demonstrate improved performance, resulting in higher win rates over PPO across three models and three datasets. These consistent results indicate SALSA’s effectiveness in various settings.

## 4 Experiments

### 4.1 Experimental setups

In this section, we outline our experimental setup. We use three models: Llama2-7B (Touvron et al., ), Mistral-7B (Jiang et al., ), and Gemma-2B (Team et al., ). For Supervised Fine-Tuning (SFT), we employ the UltraChat-200k dataset (Ding et al., ). The UltraFeedback dataset (Cui et al., ) is utilized for both training the reward model and optimizing preferences. All experiments across the three stages (SFT, reward model, RLHF) are conducted using the TRL library by HuggingFace (von Werra et al., ). For both PPO and SALSA, we use the KL coefficient $\beta$ that achieves the highest win rate. Further details on the hyperparameter settings are available in Appendix A.

Evaluation Benchmarks. We evaluate the effectiveness of our method using three well-established instruction-following benchmarks: MT-Bench (Zheng et al., ), Arena-Hard v0.1 (Li et al., ), and UltraFeedback (Cui et al., ) test dataset. MT-Bench comprises 80 questions across 8 categories, while Arena-Hard is an enhanced version of MT-Bench, featuring 500 well-defined technical problem-solving questions.

In our evaluation process, we used the datasets described in the previous section to generate samples. Pairwise comparisons were then conducted using GPT-4-Turbo as the judge model, following the "LLM-as-a-judge" methodology. The prompt used for our judge model is provided in Figure 7 in the Appendix. For evaluation, we utilized the FastChat repository (Zheng et al., ). For the MT-Bench questions, which involve two rounds, we generated and evaluated outputs for each round.

### 4.2 Reward Analysis

Figure: (a) Gemma-7B
Refer to caption: x5.png

Figure: (a) Arena-Hard
Refer to caption: x8.png

We hypothesize that $\pi_{soup}$ resides in a region of the parameter space associated with generally higher rewards, suggesting that models explored in this vicinity could generate responses with increased reward values, in addition to the improved loss observed in the original model soup paper Wortsman et al. ().
To test this hypothesis, we conducted an analysis along the interpolation line between two SFT models, $\pi_{ref}$ and $\pi_{other}$, for Gemma-7B and Llama2-7B, examining how the mean reward changes for the MT-Bench dataset. These rewards were calculated on raw models prior to RLHF process. As illustrated in Figure 2(a) and 2(b), we observe that the mean reward on the dataset increases as we move towards the midpoint between these two models, and subsequently decreases beyond this point.
This pattern indicates that $\pi_{soup}$ - which represents the average of the SFTs - resides in a region of the parameter space associated with higher rewards further supporting our hypothesis.

Next, based on our observations of improved rewards from model soup combining two SFT models, we extended our experiments to explore reward behavior within the space defined by three SFT models. These models were trained on UltraChat-200k using a pretrained Llama2-7B, each initialized with different random seeds. Specifically, we evaluated rewards on a plane defined by these three models, both inside and outside this space. Figure 2(c) presents the results of these experiments. The vertices of the dotted-line triangle represent the three SFT models. As shown, moving towards the midpoint between any two SFT models consistently leads to increased rewards. A similar trend is observed as we approach the center of the triangle. Additionally, the rewards for points outside the triangle decreases. Since PPO tends to find solutions near the vertices of the triangle (due to its reliance on the initial SFT model), the corresponding rewards in these regions are lower, as shown in the figure. This suggests that PPO may struggle to guide the model towards areas associated with higher rewards, which are located more centrally between the SFT models

Third, we hypothesize that SALSA’s ability to explore a better parameter space, enables the model to discover regions with higher rewards while maintaining the generation of sensible responses. This expanded exploration ultimately results in SALSA’s superior performance over PPO.
To validate this hypothesis, we compared the reward distributions for responses generated by SALSA and PPO across the MT-Bench and Arena-Hard datasets. Figure 3 illustrates the reward distributions for both methods. The plot reveals that the reward distribution for SALSA is shifted towards higher values compared to PPO. This rightward shift is consistent across both datasets, indicating that SALSA consistently generates responses associated with higher rewards. Furthermore, the mean reward for SALSA is higher than that of PPO in both datasets, further supporting our hypothesis.

Furthermore our findings indicate that employing a model soup as the reference model permits greater deviation in KL divergence. This increased flexibility allows the policy to investigate a more extensive area within the solution space. More info in this regard can be found in Appendix A.4.

### 4.3 Main Results

**Table 1: Comparison of Adjusted Win Rates across Models and Datasets**
| Dataset | Model | SALSA vs PPO | SALSA vs SFT |
| --- | --- | --- | --- |
| MT-Bench | Llama2-7B | 52.50 | 52.50 |
| Mistral-7B | 51.89 | 55.94 |  |
|  | Gemma-2B | 57.19 | 56.88 |
| Arena-Hard | Llama2-7B | 54.01 | 54.70 |
| Mistral-7B | 54.40 | 55.70 |  |
| Gemma-2B | 53.7 | 53.8 |  |
| UltraFeedback | Llama2-7B | 50.75 | 52.23 |
| Mistral-7B | 52.40 | 51.93 |  |
| Gemma-2B | 50.68 | 54.32 |  |

Figure 1 visualizes the win rate comparison between SALSA and PPO, where SALSA achieves notable win rates of 54.01% for the Llama2-7B model and 54.40% for the Mistral-7B model on the challenging Arena-Hard dataset. Table 1 provides a comparative analysis of SALSA against the original PPO and SFT models for Llama2-7B, Mistral-7B, and Gemma-2B. The results indicate that SALSA consistently outperforms both PPO and SFT, demonstrating its effectiveness. Based on the detailed results in tables 2, 3, and 4 in the Appendixse results, PPO and SFT often perform similarly on MT-Bench and Arena-Hard, likely because UltraFeedback and UltraChat, which are utilized for SFT, Reward Modeling, and RLHF, are considered out-of-distribution for these benchmarks. As a result, a basic version of PPO does not significantly outperform SFT. However, SALSA’s robustness to out-of-distribution data (derived from weight averaging and model soup techniques) delivers improvements of up to 57% and 54% on these datasets.
These results underscore SALSA’s effectiveness in enhancing out-of-distribution robustness in RLHF training while maintaining competitive performance for in-distribution, and emphasize SALSA’s superior exploration capabilities and improved reward optimization, resulting in better task alignment and overall performance. Additional detailed results are available in tables 2, 3, and 4 in the Appendix. Also a qualitative comparison of responses generated by SALSA and PPO is presented in Figure 8 in Appendix.

### 4.4 Ablation Study

Figure: (a) Win rates of SALSA over PPO for varying $\alpha$
Refer to caption: x10.png

We explored using alternative reference points along the line between $\pi_{\text{ref}}$ and $\pi_{\text{other}}$ by varying the $\alpha$ value in Equation 3. For each $\alpha$, we adjusted the KL divergence to achieve the optimal win rate over PPO.
This systematic adjustment allowed us to observe how the win rate varies with $\alpha$, as shown in Figure 4(a) which reveals a clear trend: the adjusted win rate increases as $\alpha$ approaches 0.5, peaking at this midpoint before declining at higher values. Notably, using $\pi_{other}$ alone does not enhance performance, yielding an adjusted win rate of only 43.07% over PPO. This outcome arises because, although the model explores a wide range between $\pi_{ref}$ and $\pi_{other}$, it ultimately converges on a lower-reward region near $\pi_{other}$. These findings support our hypotheses on reward dynamics (Section 4.2) and the effects of KL divergence (Section A.4).

To explore alternatives to SALSA, we experimented with a different loss function, $\mathcal{L_{MKL}}$, shown in equation 5, which regularizes the policy by averaging the KL divergences between the policy $\pi_{\theta}$ and two SFT models, $\pi_{\text{ref}}$ and $\pi_{\text{other}}$.

$$ $\mathcal{L_{MKL}}(\pi_{\theta})=-R(x,y)+\frac{\beta}{2}\left[\text{KL}\left( \pi_{\theta},\boldsymbol{\pi_{\text{ref}}}\right)+\text{KL}\left(\pi_{\theta}, \pi_{\text{other}}\right)\right]$ (5) $$

While SALSA employs model soup (an ensemble average of two models) as a reference point, this alternative approach calculates individual divergences from each model separately. Our empirical results, as demonstrated in Figure 4(b) and Table 5, indicate that this method does not outperform PPO. These findings highlight the significance of the averaging methodology between SFT models: using the averaged model as a single reference point for KL divergence proves more effective than computing the average of two separate KL divergences with distinct reference points.

Finally, we conducted an ablation study to assess the impact of using soups with more than two SFT models. As illustrated in Figure 2(c), the reward is higher at the midpoints between SFTs compared to the vertices. Moreover, the reward near the center of the triangle is at its peak, higher than other regions. This suggests that incorporating more SFTs into the soup is likely to enhance SALSA’s performance. Figure 5 supports this hypothesis, showing that the win rate increases as the number of SFTs in the soup model grows. Although increasing the number of SFTs for constructing the soup yields better performance, we limit the number to two due to computational constraints. Investigating soups of more than three elements and finding the optimal one is left as future work.

Figure: (a) Arena-Hard
Refer to caption: x12.png

## 5 Conclusion

This paper presents SALSA, a novel method for improving alignment in large language models by leveraging a model soup as a reference in the RLHF framework. By utilizing weight-space averaging of fine-tuned models as the reference, SALSA facilitates more effective exploration during policy optimization, leading to stronger performance in in-distribution and more resilience in out-of-distribution regimes. We showed that model soup resides in a higher reward region even before the PPO process, enabling SALSA to search for higher potential model. Furthermore we showed using model soup as a reference model allows for larger deviation in KL enabling search in a larger region. Experimental results across multiple benchmarks consistently show that SALSA outperforms PPO, yielding higher win rates, increased average rewards, and improved alignment with human preferences. We further extended our work to show averaging over more SFT models results even in higher win rates and robustness.

There are many avenues to extend this work:
applying model soups to other forms of learning from human feedback like DPO is a very interesting future work. Systematically exploring other forms of ensembling different models as reference, and model averaging with a non-uniform or adaptive weights is another valuable line of work. Finding out remedies for KL-Hacks when using SALSA is another direction for theoretical and emperical research.

## Acknowledgements

We would like to thank seyed mohsen dezfouli, fartash faghri, hooman shahrokhi for the valuable discussions.

## Appendix A Hyper-parameter setting

### A.1 Supervised fine-tuning setup

We use a learning rate of $2.5e-5$, warmup ratio of 0.03, and batch size of 8 for Llama2-7B and Mistral-7B, and 16 for Gemma-2B. For each of these models, we train two supervised fine-tuned models with different random seeds. Flash attention is applied to conserve resources, and the maximum sequence length is set to 2048 tokens for all models.

### A.2 Reward model training

The reward model training uses a learning rate of $2e-5$, with a batch size of 8. We apply a weight decay of 0.001 and use the AdamW optimizer. The learning rate schedule follows a linear decay policy.

### A.3 RLHF setup

For the RLHF stage, we use different learning rates: $2e-6$ for Llama2-7B and Mistral-7B, and $1e-7$ for the Gemma-2B model. The batch size is set to 8, with 2 epochs of training. We set the initial KL divergence coefficient to 0.2 for the PPO experiments and 0.01 for the SALSA experiments.

### A.4 KL divergence

The model soup resides on the line between two SFT models. As previous work by Wortsman et al. [] has shown, each point along the line between two fine-tuned models can improve the accuracy compared to the individual fine-tuned models. With this intuition, we hypothesize that there is a broader search space around the model soup that can still be beneficial for RLHF. Essentially, instead of searching for a model that does not diverge from a single SFT model, we can find models that do not diverge from the line between two SFT models. This allows for higher KL-divergence from the model soup. To verify this hypothesis, we tried using a small KL-divergence coefficient of $\beta=0.01$ for PPO as well (in comparison to $\beta=0.2$ in the optimal setting). Figure 6 illustrates that while PPO converged to a state producing gibberish output, SALSA, using the same $\beta$ value, achieved its optimal win rate.

Furthermore, high KL coefficients cannot be applied when using SALSA because the response length tends to converge to zero. Specifically, if the response length $y$ for a given prompt $x$ is zero, the KL divergence between the trained policy and the frozen policy (whether reference or soup model) also becomes zero. As a result, the policy can gravitates toward regions where the response length is zero, a phenomenon we refer to as the KL-Hack. Investigating this issue is an interesting future research direction.

Figure: Figure 6: A gibberish response from PPO when KL coefficient is small (0.01)

## Appendix B Prompt Used for Judgment by GPT4-Turbo

Figure 7 illustrates the prompt utilized to employ GPT4-Turbo as judge in our experiments.

Figure: Figure 7: Instruction prompt for AI response evaluation

## Appendix C An example of generated answers by SALSA and PPO

Figure 8 displays responses generated by SALSA and PPO for a question from the UltraFeedback dataset.

Figure: Figure 8: Qualitative comparison of PPO and SALSA.

## Appendix D Win Rates

This section provides a comprehensive breakdown of win rates across all our experiments. We have included the details of win rates for all the experiments for llama2-7b in table 2, for gemma-2b in table 4 and for Mistral-7b in table 3. For each dataset we have compared the win rates of SALSA over PPO, SALSA over SFT and PPO over SFT. We have included win rates and adjusted win rates. SALSA consistently outperforms PPO and SFT.

**Table 2: Comparison of SALSA, PPO, and SFT across Various Datasets for Llama2-7B Model**
| Dataset | Comparison | Win | Loss | Tie | Win Rate | Adj. Win Rate |
| --- | --- | --- | --- | --- | --- | --- |
| Llama2-7B Model |  |  |  |  |  |  |
| MT-Bench | SALSA | $29$ | $21$ | $110$ | $18.12$ | $52.50$ |
| PPO | $21$ | $29$ | $110$ | $13.12$ | $47.50$ |  |
| SALSA | $36$ | $28$ | $96$ | $22.50$ | $52.50$ |  |
| SFT | $28$ | $36$ | $96$ | $17.50$ | $47.50$ |  |
| PPO | $23$ | $21$ | $116$ | $14.37$ | $50.63$ |  |
| SFT | $21$ | $23$ | $116$ | $13.12$ | $49.37$ |  |
| Arena-Hard | SALSA | $102$ | $62$ | $335$ | $20.44$ | $54.01$ |
| PPO | $62$ | $102$ | $335$ | $12.42$ | $45.99$ |  |
| SALSA | $91$ | $44$ | $365$ | $18.20$ | $54.70$ |  |
| SFT | $44$ | $91$ | $365$ | $08.80$ | $45.30$ |  |
| PPO | $58$ | $52$ | $387$ | $11.67$ | $50.60$ |  |
| SFT | $52$ | $58$ | $387$ | $10.46$ | $49.40$ |  |
| UltraFeedback | SALSA | $477$ | $447$ | $1075$ | $23.86$ | $50.75$ |
| PPO | $447$ | $477$ | $1075$ | $22.36$ | $49.25$ |  |
| SALSA | $504$ | $415$ | $1080$ | $25.21$ | $52.23$ |  |
| SFT | $415$ | $504$ | $1080$ | $20.76$ | $47.77$ |  |
| PPO | $443$ | $417$ | $1136$ | $22.19$ | $50.65$ |  |
| SFT | $417$ | $443$ | $1136$ | $20.89$ | $49.35$ |  |

**Table 3: Comparison of SALSA, PPO, and SFT across Various Datasets for Mistral-7B model**
| Dataset | Comparison | Win | Loss | Tie | Win Rate | Adj. Win Rate |
| --- | --- | --- | --- | --- | --- | --- |
| Mistral-7B Model |  |  |  |  |  |  |
| MT-Bench | SALSA | $30$ | $24$ | $105$ | $18.87$ | $51.89$ |
| PPO | $24$ | $30$ | $105$ | $15.09$ | $48.11$ |  |
| SALSA | $42$ | $23$ | $95$ | $26.25$ | $55.94$ |  |
| SFT | $23$ | $42$ | $95$ | $14.37$ | $44.06$ |  |
| PPO | $28$ | $23$ | $109$ | $17.50$ | $51.56$ |  |
| SFT | $23$ | $28$ | $109$ | $14.37$ | $48.44$ |  |
| Arena-Hard | SALSA | $109$ | $65$ | $326$ | $21.80$ | $54.40$ |
| PPO | $65$ | $109$ | $326$ | $13.00$ | $45.60$ |  |
| SALSA | $126$ | $69$ | $305$ | $25.20$ | $55.70$ |  |
| SFT | $69$ | $126$ | $305$ | $13.80$ | $44.30$ |  |
| PPO | $71$ | $67$ | $362$ | $14.20$ | $50.40$ |  |
| SFT | $67$ | $71$ | $362$ | $13.40$ | $49.60$ |  |
| UltraFeedback | SALSA | $497$ | $401$ | $1102$ | $24.85$ | $52.40$ |
| PPO | $401$ | $497$ | $1102$ | $20.05$ | $47.60$ |  |
| SALSA | $465$ | $388$ | $1146$ | $23.26$ | $51.93$ |  |
| SFT | $388$ | $465$ | $1146$ | $19.41$ | $48.07$ |  |
| PPO | $455$ | $450$ | $1095$ | $22.75$ | $50.65$ |  |
| SFT | $450$ | $455$ | $1095$ | $20.89$ | $49.35$ |  |

**Table 4: Comparison of SALSA, PPO, and SFT across Various Datasets for Gemma-2B model**
| Dataset | Comparison | Win | Loss | Tie | Win Rate | Adj. Win Rate |
| --- | --- | --- | --- | --- | --- | --- |
| Gemma-2B Model |  |  |  |  |  |  |
| MT-Bench | SALSA | $37$ | $14$ | $109$ | $23.13$ | $57.19$ |
| PPO | $14$ | $37$ | $109$ | $8.75$ | $42.81$ |  |
| SALSA | $40$ | $18$ | $102$ | $25.0$ | $56.88$ |  |
| SFT | $18$ | $40$ | $102$ | $11.25$ | $43.12$ |  |
| PPO | $14$ | $8$ | $138$ | $8.75$ | $51.88$ |  |
| SFT | $8$ | $14$ | $138$ | $5.0$ | $48.12$ |  |
| Arena-Hard | SALSA | $154$ | $117$ | $229$ | $30.8$ | $53.7$ |
| PPO | $117$ | $154$ | $229$ | $23.4$ | $46.3$ |  |
| SALSA | $154$ | $116$ | $230$ | $30.8$ | $53.8$ |  |
| SFT | $116$ | $154$ | $230$ | $23.2$ | $46.2$ |  |
| PPO | $24$ | $19$ | $451$ | $4.86$ | $50.51$ |  |
| SFT | $19$ | $24$ | $451$ | $3.84$ | $49.49$ |  |
| UltraFeedback | SALSA | $174$ | $147$ | $1660$ | $8.78$ | $50.68$ |
| PPO | $147$ | $174$ | $1660$ | $7.42$ | $49.32$ |  |
| SALSA | $300$ | $129$ | $1548$ | $15.17$ | $54.32$ |  |
| SFT | $129$ | $300$ | $1548$ | $6.53$ | $45.68$ |  |
| PPO | $276$ | $129$ | $1580$ | $13.90$ | $53.70$ |  |
| SFT | $129$ | $276$ | $1580$ | $6.50$ | $46.30$ |  |

## Appendix E Multiple KL

Table 5 gives additional info of win MKL win rate over PPO. We compared MKL win rate over PPO over MT-Bench and Arena-Hard and in both cases MKL doesn’t outperform PPO. This means simply trying to use multiple models is not gonna be effective and doing weight averaging in SALSA is crucial for effectiveness.

**Table 5: Comparison of MKL and PPO across MT-Bench and Arena-Hard for Mistral-7B.**
| Dataset | Comparison | Win | Loss | Tie | Win Rate | Adj. Win Rate |
| --- | --- | --- | --- | --- | --- | --- |
| MT-Bench | MKL | $69$ | $85$ | $340$ | $13.97$ | $48.38$ |
| PPO | $85$ | $69$ | $340$ | $17.21$ | 51.62 |  |
| Arena-Hard | MKL | $20$ | $29$ | $111$ | $12.50$ | $47.19$ |
| PPO | $29$ | $20$ | $111$ | $18.13$ | 52.81 |  |

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Why does the mixed model work the best at the midpoint? It seems that both SFT models are trained randomly, and it seems strange that always the midpoint of them performs the best.

2. For the KL hack phenomenon, is it useful to fix the length of the generated response of the model?

3. The value of the optimal $\beta$ also changes drastically with $\alpha$. When $\alpha = 0.2$ the optimal $\beta$ is 0.2. However, when $\alpha=0.01$, the optimal $\beta$ decreases to $0.01$ quickly. It is also strange that when you take the midpoint model, the KL divergence term seems unnecessary and the trained model won't produce gibberish responses. Could the authors explain more about this phenomenon?
soundness: 2
strengths: The idea of this paper is clear and straightforward. The experiment results also seem promising.
summary: The paper studies RLHF and proposes SALSA, which leverages a "model soup" approach, averaging weights from two supervised fine-tuned (SFT) models to create a more flexible reference model. They claim that this allows broader exploration of the solution space without sacrificing stability. In the experiments, the authors use models like Llama2-7B, Mistral-7B, and Gemma-2B, along with benchmarks (MT-Bench, Arena-Hard, UltraFeedback, demonstrating that SALSA consistently outperformsPPO by achieving higher rewards and win rates.
weaknesses: The authors didn't provide sufficient explanations about the occurred phenomenon (see the Questions part), which makes their method not convincing enough. It would be great if the authors can provide some theoretical analysis, even for a simple case study.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. It is well known that the neural network is non-convex, which means that applying weighted averaging on several models could lead to a poor performance. Is there always good performance when applying model soup for the reference model? Or how do we avoid this when merging models?
soundness: 3
strengths: 1. The experiment is well done. It demonstrates the advantage of model soup for reference model in RLHF from different aspects. It explores the different weighting ratio when averaging two or three models. When the ratio is averaged, the aligned model by PPO algorithm has better performance.
summary: This paper considers using the technology of model soup by weighted averaging two reference models which is fine-tuned on the same pre-trained model to improve the reference model in RLHF. It demonstrates that using model soup to improve the reference model will lead to better aligned model by PPO algorithm.
weaknesses: 1. If the weighted averaged model has better performance than initial model, it is obvious that the aligned model by the PPO algorithm has better performance than others.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Could the authors plot results showing also KL wrt fixed anchor policy? (as mentioned in Weaknesses)
2. In the abstract and several places in the paper, the author claimed that model soups lead to "improving model robustness, out-of-distribution generalization". How are model robustness and OOD generalization entailed by the experiments?
soundness: 2
strengths: 1. The paper illustrates a simple phenomenon, that applying model soups to the anchor policy for KL regularization can lead to improved win-rates in many benchmarks.
2. The paper includes some ablations, including on how the policies are souped (finding that uniform souping is best). Also, the authors show that multiple KL terms to each individual policy does not lead to any benefits, whereas a single KL term to a souped policy does.
summary: The paper proposes a simple modification to the KL-constrained objective used in RL finetuning. Instead of regularizing to the same fixed policy $\pi_{ref}$, the authors propose to regularize to a souped policy $\pi_{soups}$ which is the model soups of multiple $\pi_{ref}$ (each SFT-ed on the same dataset with different seed or hyper-parameters). The authors claim this is a more robust ref model, and demonstrate results on instruction following benchmarks (MT-Bench, Arena-Hard, UltraFeedback) where the proposed approach improves over the classic $\pi_{ref}$ regularization.
weaknesses: 1. I think some claims made in the paper are not rigorously justified. For example, Figure 2 shows that reward increases with model souping, but it is also important to plot the KL wrt a fixed anchor policy (whether pi_ref or pi_soups or both). In other words, it may be possible to get a higher reward by simply having higher KL and this is not ruled out in the experiments.
2. Related to (1), simply showing higher win-rates in Table 1, without plotting the KL as well, is also not very convincing, since these benchmarks are easy to game. It would be much more convincing if SALSA leads to higher rewards/win-rate while maintaining the same or lower KL compared to pi_ref.
3. The paper does not discuss any important related works that apply model soups to RLHF, for example:
[1] "WARP: On the Benefits of Weight Averaged Rewarded Policies", https://arxiv.org/abs/2406.16768
[2] "WARM: On the Benefits of Weight Averaged Reward Models", https://arxiv.org/abs/2401.12187
[3] "Conditional Language Policy: A general framework for steerable multi-objective finetuning", https://arxiv.org/abs/2407.15762

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - Is your model soup similar to a model ensemble? I’m curious as to why there is the separate terminology.
soundness: 3
strengths: - interesting and nice train of work! 
- I appreciate the novel look at RLHF.
summary: The authors propose an alternative to RLHF which instantiates two models and then takes the average of them at each training step, replacing the KL divergence term in the classical RLHF setup.
weaknesses: - PPO at this point has fallen out of favor as a canonical RLHF architecture (see the Back to Basics paper by Cohere AI). Furthermore, PPO operates using a trust region policy, so it bounds the steps that a learning update could anyway. In order for me to recommend your paper for publication, I'd want to see you repeat this experiment with REINFORCE, Actor-Critic, and perhaps another RL algorithm that makes up the backbone of your RLHF pipeline.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Here are some questions and suggestions for the authors:

Dataset Generality:Could you provide more details on the choice of datasets? How do you ensure that the results are generalizable across different domains?

Have you considered testing the model on additional datasets to validate its performance more broadly?

Training Curves:Can you include training curves in your supplementary materials? These would help in understanding the convergence behavior and stability of your model.What insights can you provide about the model's learning dynamics during training?

Comparison with Other Methods:Why did you choose to compare only with PPO? Are there specific reasons for not including other related methods in your analysis?Would you consider conducting experiments with other alignment or optimization methods to provide a more comprehensive evaluation?

KL-Hack Issue:Can you elaborate on the KL-Hack phenomenon and potential strategies to mitigate it?Have you explored any preliminary solutions or adjustments to address this issue?

Future Work and Extensions: What are your plans for extending this research? Are there specific areas you believe are promising for further exploration?How do you envision incorporating more than three SFT models into the model soup, considering computational constraints?

These questions aim to clarify certain aspects of your work and suggest areas for potential improvement or further research.
soundness: 3
strengths: The SALSA method introduces the novel concept of model soup to address alignment issues in large language models. By employing weight-space averaging of two independently fine-tuned models, this innovation overcomes the limitations of KL divergence constraints in traditional RLHF methods, opening new research directions.
The paper is well-structured and logically coherent, offering detailed and easily understandable descriptions of the SALSA method. By contrasting traditional methods with SALSA, the authors effectively convey the core ideas and advantages of the method, enabling readers to clearly grasp its contributions.
summary: The paper introduces SALSA (Soup-based Alignment Learning for Stronger Adaptation), a novel method in the field of large language models, aimed at enhancing model alignment through Reinforcement Learning from Human Feedback (RLHF). Unlike traditional methods that rely on a fixed initial policy as a reference, SALSA employs a model soup created by weight-space averaging of two independently fine-tuned models. This approach allows for greater flexibility and exploration during policy optimization, resulting in improved model performance and robustness.
weaknesses: 1 Limited Dataset Generality: The experiments primarily use the UltraChat-200k and UltraFeedback datasets, which may not be comprehensive enough to reflect the model's performance across different domains or applications. As a result, the findings might not be generalizable to broader real-world scenarios.

2 Lack of Training Curve Presentation: The article does not provide training curves, making it difficult for readers to assess the model's convergence and stability during training. Training curves can help understand the learning dynamics and potential issues in the optimization process.

3 Lack of Comparison with Other Related Works: Apart from PPO, the article does not compare the SALSA method with other existing alignment or optimization methods. This limits a comprehensive understanding of SALSA's relative advantages and its competitiveness in a broader context.

Reviewers will consider raising the score if the author can effectively address their concerns.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. How does the SALSA perform when compared to DPO?
2. Can using the soup model as a reference model also bring benefit to DPO?
3. In the implementation of SALSA, how is the model initialized at the beginning of the RLHF process? If the soup model is used for initialization, a further analysis about the impact of using soup model as the RL initialization model would be expected.
4. Line-321: "high KL coefficients cannot be applied when using SALSA because the response length tends to converge to zero. Specifically, if the response length y for a given prompt x is zero, the KL divergence between the trained policy and the frozen policy also becomes zero." This statement is incorrect since when the response length y is zero, the EOS token would contribute a non-zero value for KL divergence and thus the KL divergence between the trained policy and the frozen policy would not be zero.
soundness: 3
strengths: 1. The proposed technique is clear and effective according to the experiment results.
2. The proposed technique is evaluated with different base models.
3. The authors provide sufficient analysis of the impact of model merging in the final reward, which well motivates the proposed technique.
summary: This work investigates improving the reference model in RLHF through model merging. The authors propose to use the average parameter of several SFT models that are initialized from the same base model but trained with different seeds as the reference model in the PPO algorithm. Their results demonstrate that by using a "soup" model as the reference model, PPO obtains a significant performance improvement in three dialogue benchmarks.
weaknesses: 1. As important algorithms in RLHF, reward-free methods such as DPO are not discussed in this paper.
2. The proposed technique is not fully explored in the context of RLHF. For example, the soup model can also serve to be the reference model in reward-free RLHF methods like DPO.

Meta-review: This paper proposes a model-soup technique which uses the average weight to construct the KL center in RLHF. The majority of the reviewers vote for rejection. The main complains are missing sufficient, rigorous explanation of the phenomena and the current empirical finding has been shown only for PPO  and on the selected dataset. The authors provide empirical analyses and argue that the same can be applied to DPO, though the answers are not satisfactory. Some reviewers also raise similarity between this work and other weight-averaging approach to RLHF, though I agree with the authors the algorithms are different. 

Given the reviews, I decide to not recommend publication. I think the idea here is quite interesting and potentially points out big phenomenon. But I agree with the reviewers that more evidences are needed theoretically or empirically (more datasets, more algorithms).

Decision: Reject