WARP: On the Benefits of Weight Averaged Rewarded Policies

## Abstract

Abstract Reinforcement learning from human feedback (RLHF) aligns large language models (LLMs) by encouraging their generations to have high rewards, using a reward model trained on human preferences.
To prevent the forgetting of pre-trained knowledge, RLHF usually incorporates a KL KL \mathrm{KL} roman_KL regularization; this forces the policy to remain close to its supervised fine-tuned initialization, though it hinders the reward optimization.
To tackle the trade-off between KL KL \mathrm{KL} roman_KL and reward, in this paper we introduce a novel alignment strategy named Weight Averaged Rewarded Policies ( WARP ). WARP merges policies in the weight space at three distinct stages.
First, it uses the exponential moving average of the policy as a dynamic anchor in the KL KL \mathrm{KL} roman_KL regularization.
Second, it applies spherical interpolation to merge independently fine-tuned policies into a new enhanced one.
Third, it linearly interpolates between this merged model and the initialization, to recover features from pre-training.
This procedure is then applied iteratively, with each iteration’s final model used as an advanced initialization for the next, progressively refining the KL KL \mathrm{KL} roman_KL -reward Pareto front, achieving superior rewards at fixed KL KL \mathrm{KL} roman_KL .
Experiments with Gemma policies validate that WARP improves their quality and alignment, outperforming other open-source LLMs.

###### keywords:

## 1 Introduction

LLM alignment.
Conversational agents like Gemini [, ] and GPT-4 [], along with their open-weight counterparts like Gemma [], have demonstrated remarkable abilities in complex tasks including mathematics, coding, and tool use []. These capabilities largely emerge from pre-training on next-token prediction [, ], subsequently refined through supervised fine-tuning (SFT) [, ].
As these LLMs become more powerful, aligning them with human values becomes increasingly crucial to ensure safe deployment [, ].
To this end, reinforcement learning from human feedback (RLHF) has become the prominent strategy [, , ], first learning a reward model (RM) on human preferences, before optimizing the LLM to maximize predicted rewards.

Challenges in RLHF.
However, RLHF introduces several unresolved challenges [].
First, the limited scope of fine-tuning, often restricted to relatively small datasets, can lead to excessive specialization and catastrophic forgetting [] of the broad and diverse knowledge acquired during pre-training [, , , ]. Such *alignment tax* [] can degrade the LLM’s reasoning capabilities and performance on NLP benchmarks [, ].
Second, maximizing an imperfect RM presents several issues on its own, as the LLM can learn to exploit loopholes in the RM [, ] when it deviates significantly from its initialization []. Such *reward hacking* [, ] can produce outputs that are linguistically flawed [], excessively verbose [], or sycophantic [, ], thereby raising misalignment [, ] and safety [, ] concerns.
Finally, RLHF can reduce the diversity of generations [], potentially leading to policy collapse [, ]. Such *loss of diversity* limits use in creative or exploratory tasks and can result in the LLM systematically refusing to answer.
Overall, achieving high rewards based on an imperfect RM on a selected distribution of prompts is insufficient due to potential reward misspecification and distribution shifts upon deployment.

RL with $\mathrm{KL}$ regularization.
To address these issues, previous works constrained the reward optimization by integrating a Kullback-Leibler ($\mathrm{KL}$) regularization [, ], using the SFT initialization as the anchor.
As clarified in Section 2, this $\mathrm{KL}$ regularization forces the policy to remain close to its initialization [, ], mitigating forgetting and reward hacking []. However, employing the SFT model as the anchor may lead to reward underfitting: indeed, there is a fundamental tension between reducing $\mathrm{KL}$ and maximizing reward. Thus, different policies should be compared in terms of $\mathrm{KL}$-reward Pareto optimality as in Figure 1(b), where the $x$-axis is the $\mathrm{KL}$ and the $y$-axis is the reward as estimated by the RM, with the optimal policies located in the top-left of the plot.

Figure: (a) *WARP* with three model merging stages, applicable iteratively.
Refer to caption: x1.png

On model merging by weight averaging.
To improve the trade-off between $\mathrm{KL}$ and reward during RLHF,
we leverage the ability to merge LLMs by weight averaging (WA) [].
WA relies on the linear mode connectivity  [, ], an empirical observation revealing linear paths of high performance between models fine-tuned from a shared pre-trained initialization.
Model merging was shown to improve robustness under distribution shifts [, , ] by promoting generalization and reducing memorization [], to combine models’ abilities [, , ], to reduce forgetting in continual learning [], to enable collaborative [] and distributed [] learning at scale, without computational overheads at inference time.
Model merging is increasingly adopted within the open-source community [, ], leading to state-of-the-art models in specialized domains [] but also significant advancements on general-purpose benchmarks [, ].
In particular, while WA was initially mostly used for discriminative tasks [] such as reward modeling [], it is now becoming popular for generative tasks [, ]; its use in $\mathrm{KL}$-constrained RLHF has already shown preliminary successes in a few recent works [, , , , , ], further elaborated in Section 5.

*WARP*.
In this paper, we propose Weight Averaged Rewarded Policies (*WARP*), a simple strategy for aligning LLMs, illustrated in Figure 1(a) and detailed in Section 3. *WARP* is designed to optimize the $\mathrm{KL}$-reward Pareto front of solutions, as demonstrated in Figure 1(b).
*WARP* uses three variants of WA at three different stages of the alignment procedure, for three distinct reasons.

Experiments and discussion. In Section 4, we validate the efficacy of *WARP* for the fine-tuning of Gemma "7B" []. Finally, in Section 6, we discuss the connections between *WARP*, the distributed learning literature [, ] and iterated amplification [], illustrating how *WARP* embodies their principles to enable scaling post-training, for continuous alignment and improvement of LLMs.

## 2 Context and notations

RL for LLMs.
We consider a transformer [] LLM $f\mathopen{}\mathclose{{}\left(\cdot,\theta}\right)$ parameterized by $\theta$.
Following the foundation model paradigm [] and the principles of transfer learning [], those weights are trained via a three-stage procedure: pre-training through next token prediction, supervised fine-tuning resulting in $\theta_{\mathrm{sft}}$, and ultimately, RLHF [, ] to optimize a reward $r$ as determined by a RM trained to reflect human preferences.
In this RL stage, $\theta$ defines a policy $\pi_{\theta}(\cdot\mid{\bm{x}})$ by auto-regressively generating token sequences ${\bm{y}}$ from the prompt ${\bm{x}}$.
The primary objective is to find weights maximizing the average reward over a dataset of prompts $\mathcal{X}$: $\operatorname*{argmax}_{\theta}\mathbb{E}_{{\bm{x}}\in\mathcal{X}}\mathbb{E}_{
{\bm{y}}\sim\pi_{\theta}(\cdot\mid{\bm{x}})}\Big{[}r\mathopen{}\mathclose{{}
\left({\bm{x}},{\bm{y}}}\right)\Big{]}$.

$\mathrm{KL}$ vs. reward.
Optimizing solely for $r$ can (i) forget general abilities from pre-training [] as an alignment tax [, ], (ii) hack the reward [, ] leading to potential misalignment, or (iii) reduce the diversity of possible generations [] (confirmed in Appendix F).
To mitigate these risks, a $\mathrm{KL}$ regularization is usually integrated to balance fidelity to the initialization and high rewards:

$$ $\operatorname*{argmax}_{\theta}\mathbb{E}_{{\bm{x}}\in\mathcal{X}}\mathopen{} \mathclose{{}\left[\mathbb{E}_{{\bm{y}}\sim\pi_{\theta}(\cdot\mid{\bm{x}})}r \mathopen{}\mathclose{{}\left({\bm{x}},{\bm{y}}}\right)-\beta\mathrm{KL} \mathopen{}\mathclose{{}\left(\pi_{\theta}(\cdot\mid{\bm{x}})\|\pi_{\theta_{ \mathrm{anchor}}}(\cdot\mid{\bm{x}})}\right)}\right],$ (1) $$

where $\theta_{\mathrm{anchor}}\leftarrow\theta_{\mathrm{sft}}$ and $\beta$ is the regularization strength, with high values leading to low $\mathrm{KL}$ though also lower reward.
The reward function adjusted with this $\mathrm{KL}$ is $r({\bm{x}},{\bm{y}})-\beta\log\mathopen{}\mathclose{{}\left(\frac{\pi_{\theta}
({\bm{y}}\mid{\bm{x}})}{\pi_{\theta_{\mathrm{anchor}}}({\bm{y}}\mid{\bm{x}})}}
\right).$
Our base RL algorithm is a variant of REINFORCE []. This choices follows recent RLHF works [, , ] and the findings from [, , ] that, in terms of $\mathrm{KL}$-reward Pareto optimality, REINFORCE performs better than the more complex PPO [] and also better than various offline algorithms such as DPO [], IPO [] or RAFT [].
Practitioners then typically employ early stopping to select an optimal point on the training trajectory based on their specific use cases.

## 3 WARP

We introduce a novel alignment strategy named Weight Averaged Rewarded Policies (*WARP*), illustrated in Figure 1(a) and described in Algorithm 1 below.
*WARP* merges LLMs in the weight space to enhance the $\mathrm{KL}$-reward Pareto front of policies.
The following Sections 3.1, 3.2 and 3.3 describe the motivations behind applying three distinct variants of WA at the three different stages of *WARP*.
In particular, we summarize the key insights as observations, that will be experimentally validated in Section 4 (and in Appendices C and D), and theoretically motivated in Appendix B when possible.
Overall, *WARP* outperforms other RL alignment strategies, without any memory or inference overhead at test time.
However, training *WARP* is costly, requiring multiple RL runs at each iteration: see Section 6 for a detailed discussion on the required compute scaling.

Figure: Algorithm 1 *WARP* for $\mathrm{KL}$-reward Pareto optimal alignment

### 3.1 Stage 1: exponential moving average as a dynamic anchor in KL KL \mathrm{KL} roman_KL regularization

*EMA* anchor.
$\mathrm{KL}$-regularized methods typically use the SFT initialization as a static anchor [, ], but in RL for control tasks, it is common to regularly update the anchor [, ]. In this spirit, *WARP* uses the policy’s own exponential moving average (*EMA*) [], updated throughout the RL fine-tuning process such as, at each training step with $\mu=0.01$:

$$ $\theta_{\mathrm{ema}}\leftarrow(1-\mu)\cdot\theta_{\mathrm{ema}}+\mu\cdot \theta_{\mathrm{policy}}.$ ( EMA ) $$

Using $\theta_{\mathrm{ema}}$ as the anchor $\theta_{\mathrm{anchor}}$ in Equation 1 provides several benefits, outlined below.

###### Observation 1 ( EMA ) .

Policies trained with an exponential moving average anchor benefit from automatic annealing of the $\mathrm{KL}$ regularization and from distillation from a dynamic mean teacher [].
Empirical evidence in Section 4.1.

Benefits from *EMA*.
Unlike a static SFT anchor, the dynamic nature of an *EMA* anchor induces a gradual automatic annealing and relaxation of the $\mathrm{KL}$ regularization. Specifically, the policy is initially strongly tied to the SFT initialization, and then progressively unleashed, allowing for more aggressive gradient updates later in training, leading to higher rewards. Moreover, by progressively incorporating knowledge from the training, *EMA* acts as slow weight [, ], and thus performing better than the initialization. But, by also maintaining essential information from the initialization, *EMA* can even perform better than the final policy’s weights;
studies [, , ] (see [] for a review), and specifically [] within the context of LLMs, indicate that averaging checkpoints over steps improves internal representations and thus predictions. Then, *EMA* guides the policy by $\mathrm{KL}$ distillation [] of high-quality target predictions, akin to a mean teacher [] for self-supervised [, , , , ] learning.
This also relates to deep RL techniques where *EMA* stabilizes exploration toward a Nash equilibrium [, , , ], and approximates mirror descent [, , ].

###### Observation 1 ( EMA ) .

### 3.2 Stage 2: spherical linear interpolation of independently rewarded policies

*SLERP*.
While *EMA* helps for a single RL and a fixed compute budget, it faces limitations due to the similarity of the weights collected along a single fine-tuning [].
In this second stage, we merge $M$ weights RL fine-tuned independently (each with their own *EMA* anchor). This follows model soups from Wortsman et al. [] and its variants [, ] showing that WA improves generalization, and that task vectors [] (the difference between fine-tuned weights and their initialization) can be arithmetically manipulated by linear interpolation (*LERP*) []. Yet, this time, we use spherical linear interpolation (*SLERP*) [], illustrated in Figure 2 and defined below for $M=2$:

$$ $\operatorname{slerp}\mathopen{}\mathclose{{}\left(\theta_{\mathrm{init}}, \theta^{1},\theta^{2},\lambda}\right)=\theta_{\mathrm{init}}+\frac{\sin[(1- \lambda)\Omega]}{\sin\Omega}\cdot\delta^{1}+\frac{\sin[\lambda\Omega]}{\sin \Omega}\cdot\delta^{2},$ ( SLERP ) $$

where $\Omega$ is the angle between the two task vectors $\delta^{1}=\theta^{1}-\theta_{\mathrm{init}}$ and $\delta^{2}=\theta^{2}-\theta_{\mathrm{init}}$, and $\lambda$ the interpolation coefficient.
Critically *SLERP* is applied layer by layer, each having a different angle. In Section B.3 we clarify how *SLERP* can be used iteratively to merge $M>2$ models.
To enforce diversity across weights, we simply vary the order in which text prompts ${\bm{x}}$ are given in each run: this was empirically sufficient, though other diversity strategies could help, e.g., varying the hyperparameters or the reward objectives (as explored in Figure 18(c)).

Figure: Figure 2: *SLERP* vs. *LERP*.

Benefits from *SLERP* vs. *LERP*.
Merging task vectors, either with *SLERP* or *LERP*, combines their abililities [].
The difference is that *SLERP* preserves their norms, reaching higher rewards than the base models; this is summarized in Observation 2.
In contrast, and as summarized in Observation 3, the more standard *LERP* has less impact on reward, but has the advantage of reducing $\mathrm{KL}$; indeed, as shown in Appendix B, *LERP* tends to pull the merged model towards the initialization, especially as the angle $\Omega$ between task vectors is near-orthogonal (see Observation 3).

###### Observation 2 ( SLERP ) .

Spherical linear interpolation boosts rewards, yet slightly increases $\mathrm{KL}$. Empirical evidence in Section 4.2 and theoretical insights in Equation 4.

###### Observation 3 ( LERP ) .

Linear interpolation reduces $\mathrm{KL}$, yet has reduced impact on reward. Empirical evidence in Section C.1 and theoretical insights in Equations 10 and 3.

###### Observation 4 (Task vectors) .

Task vectors $\delta$ are close to orthogonal with $\Omega\approx 90^{\circ}$, while the full weights $\theta$ are collinear. Empirical evidence in Section C.2.

###### Observation 2 ( SLERP ) .

###### Observation 3 ( LERP ) .

###### Observation 4 (Task vectors) .

### 3.3 Stage 3: linear interpolation towards initialization

*LITI*.
In the previous stage, *SLERP* combines multiple policies into one with higher rewards and slightly higher $\mathrm{KL}$.
This third stage, inspired by WiSE-FT from Wortsman et al. [], interpolates from the merged model towards the initialization:

$$ $\theta^{\eta}\leftarrow(1-\eta)\cdot\theta_{\mathrm{init}}+\eta\cdot\theta_{ \mathrm{slerp}}.$ ( LITI ) $$

Adjusting the interpolating coefficient $\eta\in[0,1]$ trades off between some newly acquired behaviors leading to high rewards vs. general knowledge from the SFT initialization.
Specifically, large values $\eta\approx 1$ provide high rewards but also high $\mathrm{KL}$, while smaller values $\eta\approx 0$ lean towards smaller rewards and minimal $\mathrm{KL}$.
Fortunately, we observe that the reduction in $\mathrm{KL}$ is proportionally greater than the reduction in reward when decreasing $\eta$.
Then, *LITI* empirically yields Pareto fronts that are noticeably above the “diagonal”, but also above those revealed during the base RLs.

###### Observation 5 ( LITI ) .

Interpolating weights towards the initialization reveals a better Pareto front than the one revealed during RL fine-tuning. Empirical evidence in Figure 1(b) and Section 4.3, and theoretical insights in Equations 21 and 5.

Benefits from *LITI*.
Previous works tried to understand how weight interpolation can mitigate forgetting while increasing robustness and generalization.
[] argues that WiSE-FT, akin to *LITI* in supervised learning contexts, recovers generalizable features from pre-training that might be lost during fine-tuning [], consistently with WA reducing catastrophic forgetting [, ] in continual learning.
Then in the context of RL, [] argues that *LITI* increases feature diversity, efficiently balancing between generality and task specificity.
Finally, [] argues that the geometric projection of the ideal weights is located between the merged model and the initialization.

###### Observation 5 ( LITI ) .

### 3.4 Iterative WARP

Iterative training.
The model merging strategies previously described not only establish an improved Pareto front of solutions, but also set the stage for iterative improvements. Indeed, if the computational budget is sufficient, we can apply those three stages iteratively, using $\theta^{\eta}$ from previous Pareto front (usually with $\eta=0.3$, choice ablated in Section D.3) as the initialization $\theta_{\mathrm{init}}$ for the next iteration, following the model recycling [, ] strategies.
Then, the entire training procedure is made of multiple iterations, each consisting of those three stages, where the final weight from a given iteration serves as an improved initialization for the next one.

###### Observation 6 (Iterative WARP ) .

Applying *WARP* iteratively improves results, converging to an optimal Pareto front. Empirical evidence in Sections 4.4 and 4.5.

###### Observation 6 (Iterative WARP ) .

## 4 Experiments: on the benefits of WARP

Setup.
We consider the Gemma "7B" [] LLM, which we seek to fine-tune with RLHF into a better conversational agent.
We use REINFORCE [] policy gradient to optimize the $\mathrm{KL}$-regularized reward.
The dataset $\mathcal{X}$ contains conversation prompts.
We generate on-policy samples with temperature $0.9$, batch size of $128$, Adam [] optimizer with learning rate $10^{-6}$ and warmup of $100$ steps.
*SLERP* is applied independently to the 28 layers.
Except when stated otherwise, we train for $T=9k$ steps, with $\mathrm{KL}$ strength $\beta=0.1$, *EMA* update rate $\mu=0.01$, merging $M=2$ policies uniformly $\lambda=0.5$, and *LITI* update rate $\eta=0.3$; we analyze those values in Appendix D.
We rely on a high capacity reward model, the largest available, which prevents the use of an oracle control RM as done in [, ].

Summary.
In our experiments, we analyze the $\mathrm{KL}$ to the SFT policy (reflecting the forgetting of pre-trained knowledge) and the reward (evaluating alignment to the RM).
In Section 4.1, we first show the benefits of using an *EMA* anchor; then in Section 4.2, we show that merging policies trained independently helps.
Moreover, in Section 4.3, we show that *LITI* improves the $\mathrm{KL}$-reward Pareto front; critically, repeating those three *WARP* stages can iteratively improve performances in Section 4.4.
A limitation is that our RM accurately approximates true human preferences only in low $\mathrm{KL}$ region, though can be hacked away from the SFT [].
Therefore, we finally report other metrics in Section 4.5, specifically comparing against open-source baselines such as Mixtral [], and reporting performances on standard benchmarks such as MMLU [].

### 4.1 Stage 1: exponential moving average as a dynamic anchor in KL KL \mathrm{KL} roman_KL regularization

Figure: (a) Reward vs. steps.
Refer to caption: x3.png

In Figures 3(a) and 3(b), we compare the training trajectories of different REINFORCE variants, where the changes lie in the choice of the anchor in the $\mathrm{KL}$ regularization and of the hyperparameter $\beta$ controlling its strength.
Results are computed every 100 training steps.
In our proposed version, the anchor is the *EMA* of the trained policy with $\beta=0.1$ and an *EMA* update rate $\mu=0.1$ (other values are ablated in Figure 15).
As the Pareto front for our strategy is above and to the left in Figure 3(b), this confirms the superiority of using such an adaptive anchor.
The baseline variants all use the SFT as the anchor, with different values of $\beta$.
The lack of regularization ($\beta=0.0$) leads to very fast optimization of the reward in Figure 3(a), but largely through hacking, as visible by the $\mathrm{KL}$ exploding in just a few training steps in Figure 3(b).
In contrast, higher values such as $\beta=0.1$ fail to optimize the reward as regularization is too strong, causing a quick reward saturation around $-0.62$ in Figure 3(a). Higher values such as $\beta=0.01$ can match our *EMA* anchor in low $\mathrm{KL}$ regime, but saturates around a reward of $-0.46$.
In contrast, as argued in Observation 1, the dynamic *EMA* anchor progressively moves away from the SFT initialization, causing implicit annealing of the regularization.
In conclusion, relaxing the anchor with *EMA* updates allows the efficient learning of $\mathrm{KL}$-reward Pareto-optimal policies, at any given $\mathrm{KL}$ level, for a fixed compute budget.
We refer the interested reader to additional experiments in Figure 14 from Section D.2 where we compare the trained policies with their online *EMA* version.

### 4.2 Stage 2: spherical linear interpolation of independently rewarded policies

In Figure 3(c), we plot $\lambda\rightarrow r\mathopen{}\mathclose{{}\left(\operatorname{slerp}
\mathopen{}\mathclose{{}\left(\theta_{\mathrm{init}},\theta^{1},\theta^{2},
\lambda}\right)}\right)$ showing reward convexity when interpolating policies via *SLERP*, validating Observation 2.
This mirrors the linear mode connectivity [] property across weights fine-tuned from a shared initialization, i.e., the fact that interpolated weights perform better than the initial models (recovered for $\lambda=0$ or $\lambda=1$).
Moreover, *SLERP* consistently obtains higher rewards than *LERP*;
yet, this is at slightly higher $\mathrm{KL}$, as further detailed in
Appendices B and C.1, where we analyze respectively their theoretical and empirical differences.

### 4.3 Stage 3: linear interpolation towards initialization

In Figure 4(a), we merge policies trained for $T$ steps, and then apply the *LITI* procedure. Critically, sliding the interpolating coefficient $\eta\in\{0,0.1,0.3,0.5,0.8,1.0\}$ reveals various Pareto fronts, consistently above the training trajectories obtained during the two independent RL fine-tunings.
Interestingly, longer fine-tunings improve performances, at high $\mathrm{KL}$, but also at lower $\mathrm{KL}$, simply by using a smaller $\eta$ afterwards.
Then in Figure 4(b), we report the Pareto fronts when merging up to $M=5$ weights.
We note that all Pareto fronts revealed when applying *LITI* are consistently above the ones from RL fine-tunings, validating Observation 5.
More precisely, best results are achieved by merging an higher number of policies $M$, suggesting a promising scaling direction.

Figure: (a) *LITI* of *SLERP* after $T$ steps.
Refer to caption: x6.png

### 4.4 Iterative WARP

In Figure 4(c), we apply the iterative procedure described in Section 3.4.
At each of the $I=5$ iterations we train $M=2$ policies for $T$ steps, with $T=9k$ for the first iteration, and $T=7k$ for iterations 2 and 3, and then $T=5k$ for computational reasons.
The *LITI* curves interpolate towards their own initialization (while Figure 1(b) interpolated towards the SFT initialization, see Section D.4 for a comparison).
We systematically observe that *LITI* curves are above the RL training trajectories used to obtain the inits.
Results get better at every iteration, validating Observation 6, although with reduced returns after a few iterations.

### 4.5 Comparisons and benchmarks

Side by side comparisons.
To conclude our experiments, we compare our trained policies against Mistral [] and Mixtral [] LLMs.
Each policy generates a candidate answer on an held-out collection of prompts, as in the Gemma tech report [].
Then similarly to Gemini 1.5 [], we compute side by side preference rates [] with “much better”, “better” and “slightly better” receiving scores of $\pm 1.5$, $\pm 1$, and $\pm 0.5$ respectively (and ties receiving a score of 0 0).
A positive score represents better policies.
The results in Table 1 validate the efficiency of *WARP*, as our policies are preferred over Mistral variants, and also outperform the two previous Gemma "7B" releases. However, we note that the results stagnate after the \nth3 iteration.

Table: Table 1: Side by side comparisons.

Benchmarks.
Table 2 compares *WARP* (\nth3 iter) and the latest Gemma "7B" 1.1 release [] on popular benchmarks in the zero-shot setup: MBPP [] and HumanEval [] benchmarking coding capabilities, MMLU [] assessing STEM knowledge, the GSM8K [] and MATH [] benchmarks targeting reasoning abilities, and the Big Bench Hard (BBH) [] benchmark evaluating general capabilities through questions that were deemed difficult for frontier LLMs.
*WARP* has particularly strong results on mathematics benchmarks, suggesting higher analytical capabilities.

Table: Table 2: Benchmark results.

## 5 Related work

How to merge models.
The question of how best to merge models has recently garnered significant attention, driven by the discoveries that deep models can be merged in the weight space [] instead of in the prediction space, as traditionally done in ensembling [, ]. For clarity, we collectively refer to these methods as weight averaging (WA).
The most common is *LERP*, initially used to average checkpoints collected along a single run, uniformly [, ] or with an exponential moving average (*EMA*) [], notably as a mean teacher [] for self-supervision [, , , , ].
Following the linear mode connectivity [] observation, the model soups variants [, , ] linearly interpolate from different fine-tunings; this relies on the shared pre-training, limiting divergence [] such as models remain in constrained weight regions [], which also suggests that pre-training mitigates the need to explicitly enforce trust regions in gradient updates [, ].
Subsequent works such as TIES merging [] and DARE [] reduce interferences in multi-task setups with sparse task vectors []. In contrast, we use *SLERP*, introduced in [], increasingly popular in the open-source community [] but relatively underexplored in the academic literature, with limited studies such as [].
Some tried to align weights trained from scratch [, ] or with different architectures []; yet, the methods are complex, less robust, and usually require additional training.

Benefits of model merging. WA boosts generalization by reducing variance [, ], decreasing memorization [, , ] and flattening the loss landscape [].
Additionally, merging weights combines their strengths [], which helps in multi-task setups [, ], to tackle catastrophic forgetting [, ] or to provide better initializations [], as explored in [, , ] for iterative procedures in classification tasks.
In particular, we considered using the geometric insights from Eq. 2 in []; yet, as our task vectors are nearly orthogonal $\Omega\approx 90^{\circ}$ (see Section C.2), using the update rule $\eta\rightarrow\frac{2\cos\Omega}{1+\cos\Omega}$ failed.
WA is now also used in RL setups [, , ]; for example, *WARM* [] merges reward models to boost their efficiency, robustness and reliability.
Actually, *WARP* is conceived as a response to *WARM*, demonstrating that model merging can tackle two key RLHF challenges; policy learning in *WARP* and reward design in *WARM*.
The most similar works are the following, which also explore how WA can improve policy learning.
[] proposes an iterative approach with the *EMA* as a new initialization for subsequent iterations.
[] and [] uses *EMA* as the reference, but only for direct preference optimization.
[] employs *LERP* to improve alignment in multi-objective RLHF when dealing with different objectives; similarly, [] targets multi-task setups with *LERP*.
Finally, [] and [] use model merging to reduce the alignment tax, although without incorporating *EMA* during training, without merging multiple rewarded policies and not iteratively.
Critically, none of these works focus on $\mathrm{KL}$ as a measure of forgetting, use *EMA* as the anchor in $\mathrm{KL}$, apply *SLERP* or use *LITI* as the initialization for subsequent RL iterations.
In contrast, *WARP* integrates all those elements, collectively leading to an LLM outperforming Mixtral [].

## 6 Discussion

Distributed learning for parallelization and open-source.
*WARP* addresses a crucial challenge: aligning LLMs with human values and societal norms, while preserving the capabilities that emerged from pre-training. To this end, we leverage a (perhaps surprising) ability: policies trained in parallel can combine their strengths within a single policy by weight averaging.
Then, the distributed nature of *WARP* makes it flexible and scalable, as it is easily parallelizable by enabling intermittent weight sharing across workers.
Actually, iterative *WARP* shares similarities with DiLoCo []: by analogy, the first stage performs inner optimization on multiple workers independently; the second stage merges gradients from different workers; the third stage performs SGD outer optimization with a learning rate equal to $\eta$.
More generally, *WARP* could facilitate open-source [] collaborative training of policies [], optimizing resource and supporting privacy in federated learning [] scenarios; collaborators could train and share their LLMs, while keeping their data and RMs private.
In particular, we show in Appendix E that *WARP* can handle diverse objectives, similarly to [].

Iterated amplification.
*WARP* improves LLM alignment by leveraging the principles of iterated amplification [] and progressive collaboration of multiple agents. By analogy, model merging via WA acts as an effective alternative to debate [], with agents communicating within the weight space instead of the token space, ensuring that only essential information is retained [].
Then, *WARP* refines the training signal by combining insights and exploration from diverse models, iteratively achieving higher rewards through self-distillation [], surpassing the capabilities of any single agent. If this is the way forward, then an iterative safety assessment would be required to detect and mitigate potential risks early, ensuring that the development remains aligned with safety standards.

Scaling alignment.
The *WARP* procedure increases the compute training cost by performing multiple fine-tunings at each iteration.
Yet, this should be viewed as “a feature rather than a bug”. Specifically, by preventing memorization and forgetting, we see *WARP* as a fine-tuning method that can transform additional compute allocated to alignment into enhanced capabilities and safety. This would allow scaling (the traditionally cheap) post-training alignment, in the same way pre-training has been scaled [].
Critically for large-scale deployment, the acquired knowledge is within a single (merged) model, thus without inference or memory overhead, in contrast to “more agents” approaches [, ].
Finally, although *WARP* improves policy optimization, it is important to recognize that *WARP* does not address other critical challenges in RLHF []: to mitigate the safety risks [, , ] from misalignment [, ], *WARP* should be part of a broader responsible AI framework.

## 7 Conclusion

We introduce Weight Averaged Rewarded Policies (*WARP*), a novel RLHF strategy to align LLMs with three distinct stages of model merging: exponential moving average as a dynamic anchor during RL, spherical interpolation to combine multiple policies rewarded independantly, and interpolation towards the shared initialization.
This iterative application of *WARP* improves the $\mathrm{KL}$-reward Pareto front, aligning the LLMs while protecting the knowledge from pre-training, and compares favorably against state-of-the-art baselines.
We hope *WARP* could contribute to safe and powerful AI systems by scaling alignment, and spur further exploration of the magic behind model merging.

## Appendix A Strategy illustration

In Figure 5, we propose an alternative illustration of *WARP*, where the different stages are more detailed than in Figure 1(a).
Then in Figure 6, we also refine our illustration showcasing the similarity and difference between *SLERP* and *LERP*.

Figure: Figure 5: Detailed illustration of the *WARP* strategy. From a (pre-trained and supervised fine-tuned) LLM $\theta_{\mathrm{init}}$, we launch $M=2$ fine-tunings (black arrows ). The innovation of *WARP* lies in the use of model merging by weight averaging at three different stages. First, the exponential moving averages (*EMA*, blue dashed arrows ) of the policy (collected at different training steps) serves as the anchor for the $\mathrm{KL}$ regularization (black double-headed dotted arrows ). The fine-tuned networks are weight averaged using spherical linear interpolation of task vectors (*SLERP*, yellow dashed arrows ). Third, we interpolate towards the initialization (*LITI*, red dashed arrows ). This obtained model $\theta_{\mathrm{init}}^{\prime}$ serves as an updated initialization for the next iteration, progressively refining the model’s capabilities and alignment. Overall, the final model $\theta_{\mathrm{slerp}}^{\prime}$ has high reward but also high $\mathrm{KL}$. Then, by interpolation towards the SFT init, we reveal a $\mathrm{KL}$-reward Pareto front of solutions: $\{(1-\eta)\cdot\theta_{\mathrm{sft}}+\eta\cdot\theta_{\mathrm{slerp}}^{I}\mid 0 \leq\eta\leq 1\}$.

Figure: Figure 6: Illustration of the difference between the full weights $\theta^{m}$ and their task vectors $\delta^{m}=\theta^{m}-\theta_{\mathrm{init}}$, where darker areas are of better performance. We found in Section C.2 that $\Omega\approx 90^{\circ}$ where $\Omega$ is the angle between task vectors such as $\cos\Omega=\frac{\delta^{1}\cdot\delta_{2}}{\lVert\delta^{1}\rVert\lVert\delta ^{2}\rVert}$, while $\omega$ the angle between the full weights such as $\cos\omega=\frac{\theta^{1}\cdot\theta_{2}}{\lVert\theta^{1}\rVert\lVert\theta ^{2}\rVert}$ satisfies $\omega\approx 0^{\circ}$.

## Appendix B Theoretical insights on task vectors, SLERP , LERP and LITI

Based on the insights from [] that task vectors (the differences between a fine-tuned model and its initialization) are semantically manipulable and interpretable units in the weight space, we compare *SLERP* and *LERP* merging operations by analyzing their task vectors.

Background.
Linear interpolation (*LERP*) [] is the simplest merging strategy, notably used in the model soups variants [], and defined as:

$$ $\displaystyle\operatorname{lerp}\mathopen{}\mathclose{{}\left(\theta^{1}, \theta^{2},\lambda}\right)$ $\displaystyle=(1-\lambda)\cdot\theta^{1}+\lambda\cdot\theta^{2}.$ ( LERP ) $$

Then, as illustrated in Figure 6, the task vector for *LERP* with interpolating coefficient $\lambda$ is given by: $\delta_{\mathrm{lerp}}^{\lambda}=\operatorname{lerp}\mathopen{}\mathclose{{}
\left(\theta^{1},\theta^{2},\lambda}\right)-\theta_{\mathrm{init}}=(1-\lambda)
\cdot\delta^{1}+\lambda\cdot\delta^{2}$. Similarly, we define $\delta_{\mathrm{slerp}}^{\lambda}=\operatorname{slerp}\mathopen{}\mathclose{{}
\left(\theta_{\mathrm{init}},\theta^{1},\theta^{2},\lambda}\right)-\theta_{
\mathrm{init}}$ where $\operatorname{slerp}$ is defined in Equation *SLERP*.

### B.1 Theoretical insights on the SLERP and LERP task vectors

We denote $\Omega$ the angle between the the task vectors $\delta^{1}$ and $\delta^{2}$:

$$ $\cos\Omega=\frac{\delta^{1}\cdot\delta_{2}}{\lVert\delta^{1}\rVert\lVert\delta ^{2}\rVert}.$ (2) $$

Based on the empirical observations from [], confirmed in our Figure 11(c), we introduce the following Equation 3 for simplicity.

###### Assumption 1 (Task vectors of equal norm) .

Independently fine-tuned task vectors have a same norm $l$:

$$ $\lVert\delta^{1}\rVert=\lVert\delta^{2}\rVert=l.$ (3) $$

###### Lemma 1 ( SLERP task vector) .

Under Equation 3, *SLERP* preserves the norm of the task vector:

$$ $\lVert\delta_{\mathrm{slerp}}^{\lambda}\rVert=l.$ (4) $$

###### Proof.

By definition,

$$ $\displaystyle\delta_{\mathrm{slerp}}^{\lambda}=\frac{\sin[(1-\lambda)\Omega]}{ \sin\Omega}\cdot\delta^{1}+\frac{\sin[\lambda\Omega]}{\sin\Omega}\cdot\delta^{2}$ (5) $$

Then, as $\delta^{1}\cdot\delta^{2}=l^{2}\cos\Omega$,

$$ $\displaystyle\frac{\lVert\delta_{\mathrm{slerp}}^{\lambda}\rVert^{2}}{l^{2}}$ $\displaystyle=\mathopen{}\mathclose{{}\left(\frac{\sin[(1-\lambda)\Omega]}{ \sin\Omega}}\right)^{2}+2\frac{\sin[(1-\lambda)\Omega]}{\sin\Omega}\frac{\sin[ \lambda\Omega]}{\sin\Omega}\cos(\Omega)+\mathopen{}\mathclose{{}\left(\frac{ \sin[\lambda\Omega]}{\sin\Omega}}\right)^{2}$ (6) $\displaystyle=\frac{\sin^{2}[(1-\lambda)\Omega]+2\sin[(1-\lambda)\Omega]\sin[ \lambda\Omega]\cos(\Omega)+\sin^{2}[\lambda\Omega]}{\sin^{2}\Omega}$ (7) $\displaystyle=\frac{\sin^{2}\Omega}{\sin^{2}\Omega}$ (8) $\displaystyle=1$ (9) $$

using trigonometric identities, proving Equation 4.
∎

###### Lemma 2 ( LERP task vector) .

Under Equation 3, *LERP* reduces the norm of the task vector:

$$ $\lVert\delta_{\mathrm{lerp}}^{\lambda}\rVert=l\sqrt{1-2(1-\cos\Omega)(\lambda- \lambda^{2})}.$ (10) $$

We recover that averaging weights with $\lambda=0.5$ tends to reduce the norm of the task vectors, as previously highlighted in [].

###### Proof.

By definition:

$$ $\displaystyle\delta_{\mathrm{lerp}}^{\lambda}=(1-\lambda)\cdot\delta^{1}+ \lambda\cdot\delta^{2}.$ (11) $$

Then, as $\delta^{1}\cdot\delta^{2}=l^{2}\cos\Omega$,

$$ $\displaystyle\frac{\lVert\delta_{\mathrm{slerp}}^{\lambda}\rVert^{2}}{l^{2}}$ $\displaystyle=(1-\lambda)^{2}+2\lambda(1-\lambda)\cos\Omega+\lambda^{2}$ (12) $\displaystyle=1-2\lambda(1-\cos\Omega)+2\lambda^{2}(1-\cos\Omega)$ (13) $\displaystyle=1-2(1-\cos\Omega)(\lambda-\lambda^{2}),$ (14) $$

proving Equation 10 when $0<\lambda<1$.
∎

###### Assumption 1 (Task vectors of equal norm) .

###### Lemma 1 ( SLERP task vector) .

###### Proof.

###### Lemma 2 ( LERP task vector) .

###### Proof.

### B.2 Theoretical insights on the KL KL \mathrm{KL} roman_KL

#### B.2.1 Linear regime

###### Assumption 2 (Linear regime [ 138 ] ) .

We assume that the predictions of a model $f$, with weights initialized from $\theta_{0}$ and fine-tuned into $\theta$, can be approximated by first-order Taylor expansion: $\forall{\bm{x}},$

$$ $f({\bm{x}},\theta)\approx f({\bm{x}},\theta_{0})+\mathopen{}\mathclose{{}\left (\theta-\theta_{0}}\right)\cdot\nabla_{\theta}f\mathopen{}\mathclose{{}\left({ \bm{x}},\theta_{0}}\right).$ (15) $$

Assumption 2 defines a neural tangent [] space in which the relationship between weights and functions is linear.
As previously argued in [, ], this Taylor expansion is reasonable partly because weights remain close during fine-tunings [], as confirmed in Figure 11 where they have equal norms and a cosine of one.
Yet, please note that [] highlighted some limitations.

###### Assumption 2 (Linear regime [ 138 ] ) .

#### B.2.2 KL KL \mathrm{KL} roman_KL variations for LERP

We consider $\theta^{1}$ and $\theta^{2}$ weights fine-tuned from a shared SFT initialization $\theta_{\mathrm{sft}}$. Then in the linear regime from Assumption 2, weight and prediction ensembling behaves similarly:

$$ $f\big{(}{\bm{x}},(1-\lambda)\cdot\theta^{1}+\lambda\cdot\theta^{2}\big{)} \approx(1-\lambda)\cdot f({\bm{x}},\theta^{1})+\lambda\cdot f({\bm{x}},\theta^ {2}).$ (16) $$

This similarity enables to prove the following Lemma 3.

###### Lemma 3 (LERP reduces $\mathrm{KL}$ ) .

For an interpolating coefficient $0\leq\lambda\leq 1$, denoting $\pi_{\lambda}$ the *LERP* policy from weight interpolation $(1-\lambda)\cdot\theta^{1}+\lambda\cdot\theta^{2}$, and $\hat{\pi}_{\lambda}$ the ensembling policy from prediction interpolation $(1-\lambda)\cdot\pi_{\theta^{1}}+\lambda\cdot\pi_{\theta^{2}}$, then under Assumption 2,

$$ $\mathrm{KL}(\pi_{\lambda}||\pi_{\theta_{\mathrm{sft}}})\approx\mathrm{KL}(\hat {\pi}_{\lambda}||\pi_{\theta_{\mathrm{sft}}})\leq(1-\lambda)\mathrm{KL}(\pi_{ \theta^{1}}||\pi_{\theta_{\mathrm{sft}}})+\lambda\mathrm{KL}(\pi_{\theta^{2}}| |\pi_{\theta_{\mathrm{sft}}}),$ (17) $$

i.e., the $\mathrm{KL}$ for *LERP* is lower than the interpolated $\mathrm{KL}$.

###### Proof.

The following proof applies the linear assumption and properties of the KL divergence.

Approximation of KL. The first approximate equality is a direct application of Assumption 2 to $\pi_{\lambda}$. Precisely, applying Equation 16 to the definition of $\pi_{\lambda}=\pi_{(1-\lambda)\theta_{1}+\lambda\theta_{2}}$ yields that $\pi_{\lambda}\approx\hat{\pi}_{\lambda}$.

Upper bound of the KL. The KL divergence is convex in both its arguments [], thus we directly have that

$$ $\mathrm{KL}((1-\lambda)\cdot\pi_{\theta^{1}}+\lambda\cdot\pi_{\theta^{2}}||\pi _{\theta_{\mathrm{sft}}})\leq(1-\lambda)\mathrm{KL}(\pi_{\theta^{1}}||\pi_{ \theta_{\mathrm{sft}}})+\lambda\mathrm{KL}(\pi_{\theta^{2}}||\pi_{\theta_{ \mathrm{sft}}}),$ (18) $$

which completes the proof.

∎

###### Remark 1 .

Lemma 3 shows that the LERP $\pi_{\lambda}$ is closer in $\mathrm{KL}$ to the original SFT initialization. This relates to Equation 10, where we show that the linear interpolation reduces the norm to the initialization. As the interpolation brings the weights of the models closer, it is natural that it would also bring the resulting policies closer.

###### Lemma 3 (LERP reduces KL KL \mathrm{KL} roman_KL ) .

###### Proof.

###### Remark 1 .

#### B.2.3 KL KL \mathrm{KL} roman_KL and reward variation for LITI

We now consider a given weight $\theta$ (in practice either obtained from *LERP* or *SLERP* of multiple fine-tuned weights) and its associated task vector $\delta$ = $\theta-\theta_{\mathrm{sft}}$.
In the linear regime from Assumption 2, for each $\eta\in[0,1]$, we have the following:

$$ $f\mathopen{}\mathclose{{}\left({\bm{x}},\theta_{\mathrm{sft}}+\eta\cdot\delta} \right)-f\mathopen{}\mathclose{{}\left({\bm{x}},\theta_{\mathrm{sft}}}\right) \approx\eta\cdot\mathopen{}\mathclose{{}\left(f\mathopen{}\mathclose{{}\left({ \bm{x}},\theta_{\mathrm{sft}}+\delta}\right)-f\mathopen{}\mathclose{{}\left({ \bm{x}},\theta_{\mathrm{sft}}}\right)}\right).$ (19) $$

We try to show that:

$$ $\mathrm{KL}\mathopen{}\mathclose{{}\left(\pi_{\theta_{\mathrm{sft}}+\eta\cdot \delta}\|\pi_{\theta_{\mathrm{sft}}}}\right)\leq\eta\cdot\mathrm{KL}\mathopen{ }\mathclose{{}\left(\pi_{\theta_{\mathrm{sft}}+\delta}\|\pi_{\theta_{\mathrm{ sft}}}}\right).$ (20) $$

###### Lemma 4 ( $\mathrm{KL}$ upper bound for interpolated distributions) .

For an interpolating coefficient $0\leq\eta\leq 1$, denoting $\pi_{\eta}$ the *LITI* policy from weight interpolation $\theta_{\mathrm{sft}}+\eta\cdot\delta$, and $\hat{\pi}_{\eta}$ the ensembling policy from prediction interpolation $(1-\eta)\cdot\pi_{\theta_{\mathrm{sft}}}+\eta\cdot\pi_{\theta_{\mathrm{sft}}+\delta}$, then under Assumption 2,

For each $\eta\in[0,1]$, we have that

$$ $\mathrm{KL}\mathopen{}\mathclose{{}\left(\pi_{\eta}\|\pi_{\theta_{\mathrm{sft} }}}\right)\approx\mathrm{KL}(\hat{\pi}_{\eta}\|\pi_{\theta_{\mathrm{sft}}}) \leq\eta\mathrm{KL}\mathopen{}\mathclose{{}\left(\pi_{\theta_{\mathrm{sft}}+ \delta}\|\pi_{\theta_{\mathrm{sft}}}}\right).$ (21) $$

###### Proof.

The following proof uses the same method as the one of Lemma 3. We use Assumption 2 to link the policy with the interpolation of polices, and the inequality is a result of the KL convexity.

Approximation of KL. The first approximate equality is a direct application of Assumption 2 to $\pi_{\eta}$. Precisely, applying Equation 19 to the definition of $\pi_{\eta}=\pi_{\theta_{\mathrm{sft}}+\eta\cdot\delta}$ yields that $\pi_{\eta}\approx\hat{\pi}_{\eta}$.

Upper bound of the KL. Using the fact that the KL is convex, we have

$$ $\mathrm{KL}(\eta\cdot\pi_{\theta_{\mathrm{sft}}+\delta}+(1-\eta)\cdot\pi_{ \theta_{\mathrm{sft}}}||\pi_{\theta_{\mathrm{sft}}})\leq\eta\mathrm{KL} \mathopen{}\mathclose{{}\left(\pi_{\theta_{\mathrm{sft}}+\delta}\|\pi_{\theta_ {\mathrm{sft}}}}\right).$ (22) $$

∎

###### Assumption 3 (LITI reward is above the expected reward) .

The rewards for the *LITI* interpolated weights are above the interpolated rewards:

$$ $r(\pi_{0}+\eta\cdot(\pi-\pi_{\theta_{\mathrm{sft}}}))\geq\eta r(\pi)+(1-\eta)r (\pi_{\theta_{\mathrm{sft}}}),$ (23) $$

This Assumption 3
is based on observations from Figure 8(e), and extends to a reward maximization setup the notion of linear mode connectivity [], usually defined w.r.t. the accuracy in supervised learning.

###### Lemma 5 (LITI Pareto optimality) .

Be given the convexity of the $\mathrm{KL}$ from Equation 21 and the concavity of the reward $r$ in Assumption 3, then the reward vs. $\mathrm{KL}$ front of *LITI* is above the diagonal. Illustration in Figure 7.

###### Proof.

We obtain a policy $\pi_{\theta}$ fine-tuned from $\pi_{\theta_{\mathrm{sft}}}$. The *LITI* policy for $\theta_{\eta}=(1-\eta)\cdot\theta_{\mathrm{sft}}+\eta\cdot\theta$ is noted $\pi_{\eta}$. Combining the approximation from Equation 21 and Assumption 3, we have that

$$ $r(\pi_{\eta})\geq(1-\eta)r(\pi_{\theta_{\mathrm{sft}}})+\eta r(\pi_{\theta}).$ (24) $$

And, from Equation 21, we also have that

$$ $\mathrm{KL}\mathopen{}\mathclose{{}\left(\pi_{\eta}\|\pi_{\theta_{\mathrm{sft} }}}\right)\leq\eta\mathrm{KL}\mathopen{}\mathclose{{}\left(\pi_{\theta}\|\pi_{ \theta_{\mathrm{sft}}}}\right).$ (25) $$

This means that for every *LITI* coefficient $\eta$, the *LITI* policy has a higher reward than the interpolated reward at a lower KL. Geometrically, this means that each point on the Reward-KL front from LITI is on the top left quadrant of the plane according to the corresponding point on the diagonal.
∎

Figure: Figure 7: Illustration of Lemma 5. Based on experimental observation and theoretical insights, we see that the Pareto front of the *LITI* policy is better than the identity. It highlights how Equations 24 and 25 place *LITI* policies on the $\mathrm{KL}$-reward plane.

###### Lemma 4 ( KL KL \mathrm{KL} roman_KL upper bound for interpolated distributions) .

###### Proof.

###### Assumption 3 (LITI reward is above the expected reward) .

###### Lemma 5 (LITI Pareto optimality) .

###### Proof.

### B.3 Uniformly averaging M > 2 𝑀 2 M>2 italic_M > 2 weights with SLERP

The *SLERP* merging formula from Equation *SLERP* is only defined for $M=2$ weights.
We trivially (and certainly suboptimally) generalize this to $M>2$ weights in the uniform averaging setup, thus giving an equal coefficient to each of them, i.e., $\lambda=\frac{1}{M}$.
In that setup, removing the dependency of $\theta_{\mathrm{init}}$ that is assumed shared, we generalize *SLERP* to merge $M$ weights uniformly through the iterative procedure defined below:

$$ $\operatorname{slerpm}\mathopen{}\mathclose{{}\left(\{\theta^{m}\}_{m=1}^{M}} \right)=\operatorname{slerp}\mathopen{}\mathclose{{}\left(\operatorname{slerpm }\mathopen{}\mathclose{{}\left(\{\theta^{m}\}_{m=1}^{M-1}}\right),\theta^{M}, \lambda=\frac{1}{M}}\right).$ (26) $$

Though these operations are not associative, the standard deviations in performances are small, as indicated by the shaded areas in Figure 4(b).

## Appendix C Empirical insights on task vectors, SLERP , LERP and LITI

### C.1 Empirical insights on the difference between SLERP and LERP

We now empirically investigate how those theoretical differences between *SLERP* and *LERP* affect the performance of the merged policies.

*SLERP* vs. *LERP*.
In Figure 10 we adjust the interpolating coefficient $\lambda$, highlighting distinct behaviors for *SLERP* and *LERP*.
*SLERP* consistently enhances rewards more than *LERP*, as depicted in Figures 3(c) and 8(a). However, a comprehensive evaluation must consider both $\mathrm{KL}$ and reward.
As shown in Figure 8(b), *LERP* consistently reduces $\mathrm{KL}$, corroborating with Equation 10 that *LERP* reduces the norm of updates (while *SLERP* preserves it).
When plotting these metrics together in Figure 8(c), we observe that *SLERP* and *LERP* target different regions on the Pareto front: *SLERP* achieves higher rewards at the expense of increased $\mathrm{KL}$, while the main impact of *LERP* is to lower $\mathrm{KL}$.
This is consistent with Equations 10 and 3, be given the orthogonal angles between task vectors $\Omega\approx 90^{\circ}$ (as shown in Figure 11(a)).

Combining *SLERP* and *LERP* with *LITI*.
We also compare the behaviours of *SLERP* and *LERP* when we apply *LITI*, as we adjust the interpolating coefficient $\eta$.
Figure 8(d) and Figure 8(e) validate that $\mathrm{KL}$ is convexe with regard to $\eta$ while the reward is concave with regard to $\eta$, for different values of $M$.
This is also highlighted in Figure 9(a), which reproduces the results from Figure 4(b) (and maintaining the same axis limits), replacing *SLERP* by *LERP*: this leads to critical changes in the Pareto fronts. Inded, increasing $M$ now tends to decrease $\mathrm{KL}$ for *LERP*, while it used to increase reward with *SLERP*.
In Figure 9(b), we explore the extrapolation strategies from [], using $0\leq\eta\leq 2$ to compare the full extrapolated fronts from *LERP* and *SLERP*. While both perform similarly on low $\mathrm{KL}$, our results suggest that *SLERP* perform better in high $\mathrm{KL}$ regions.

Conclusion.
*SLERP* demonstrates some key advantages. In particular, it reveals the full Pareto front of solutions, while *LERP* only exposes a portion; extrapolation Figure 9(b) with $\eta>1$ can partially mitigate this but as our experiments suggest, *LERP* curves consistently lag behind *SLERP* curves in high-reward regions.
Moreover, from a practical perspective, *SLERP* scales the choice of $\eta$ effectively, where $1$ represents full updates and a fixed value of $0.3$ always corresponds to the same operational region, optimizing for high reward and $\mathrm{KL}$.

Figure: (a) Reward vs. $\lambda$.
Refer to caption: x9.png

### C.2 Empirical insights on the role of task vectors

We now explore the effectiveness of applying *SLERP* on task vectors $\delta$ vs. full weights $\theta$, as illustrated in Figure 6.
To this end, in Figure 11 we draw inspiration from [] and plot the angles $\Omega$ and $\omega$ and norms of $\delta$ and $\theta$.

Angles of task vectors $\Omega\approx 90^{\circ}$.
Figure 11(a) shows that the task vectors are typically orthogonal ($\Omega\approx 90^{\circ}$), highlighting the diverse trajectories of the different RL fine-tunings.
This is in contrast with [] for supervised fine-tunings, where $\Omega$ typically range between $40^{\circ}$ and $80^{\circ}$.
We suspect that this is related to the underlying differences between reinforcement and supervised learning; in RL the policies are trained on their own generations, creating more orthogonal task vectors, whereas in supervised learning the LLM try to imitate the groundtruth labels, leading to more similar task vectors.
The orthogonality of our task vectors prevents the use of the update rule $\eta\rightarrow\frac{2\cos\Omega}{1+\cos\Omega}$ suggested from Eq. 2 in [], as it would lead to $\eta\approx 0$, deleting any potential update.

Angles of full weights $\omega\approx 0^{\circ}$.
In contrast, Figure 11(b) show that full weights remain collinear ($\omega\approx 0^{\circ}$).
This explains the empirical results from Figure 9(c), where applying *SLERP* directly to full weights results in behaviors similar to *LERP*.
Indeed, as the angles $\omega\approx 0^{\circ}$, spherical interpolation effect is minimal because $\sin(x)\approx x+\mathcal{O}(x^{3})$, and thus $\frac{\sin[\lambda\omega]}{\sin\omega}\approx\frac{\lambda\omega}{\omega}\approx\lambda$.

Norms consistency.
Figure 11(c) confirms the consistency in the norms of different task vectors, supporting our Equation 3. This uniformity is aligned with previous research [].
As a side note, this consistency extends to full weights $\theta$, confirming that fine-tuning typically results in minimal changes to the overall weight [].

Figure: (a) $\frac{\delta^{1}\cdot\delta^{2}}{\lVert\delta^{1}\rVert\lVert\delta^{2}\rVert}$.
Refer to caption: x17.png

## Appendix D Empirical investigation of several design choices

We include several experiments showcasing the robustness of *WARP* to different design choices, while further demonstrating its superiority in terms of $\mathrm{KL}$-reward Pareto optimality.
Specifically,
Section D.1 analyzes the performances along training at different steps $T$;
Section D.2 provides results with different values for the hyperparameters $\mu$ and $\beta$;
Section D.3 shows the impact of the update rate $\eta$ to provide an improved initialization for the \nth2 iteration of *WARP*;
finally, Section D.4 shows that in iterative *WARP*, interpolating towards the episode initialization or the SFT initialization both perform similarly.

### D.1 Analyzing the number of training steps

Figure: (a) Reward along training.
Refer to caption: x21.png

Figure: Figure 13: *LITI* with $M=1$ at different number of training steps $T$. The reward gain is significantly reduced compared to Figure 4(a) where we first merged $M=2$ policies before applying *LITI*. We also try to perform moving average (MA) [, ] before applying *LITI*, averaging multiple checkpoints collected along a single RL fine-tuning at steps $\{6k,7k,8k,9k\}$; this does not improve performances, suggesting the need to merge weights from independent fine-tunings to have enough diversity [].
Refer to caption: x23.png

### D.2 Analyzing the values of μ 𝜇 \mu italic_μ and β 𝛽 \beta italic_β

Figure: (a) Reward vs. $\mathrm{KL}$.
Refer to caption: x24.png

Figure: (a) Reward vs. $\mathrm{KL}$.
Refer to caption: x27.png

### D.3 Analyzing the values of η 𝜂 \eta italic_η

Figure: Figure 16: Experiments ablating the *LITI* update rate $\eta$. As we initiate the \nth2 iteration of *WARP*, selecting an appropriate value for $\eta$ is key, as it determines the starting point $\theta^{\eta}$ and functions similarly to an outer learning rate (see Section 6). We usually set $\eta=0.3$, but now provide results with an increased $\eta=0.5$, starting the \nth2 iteration from a more “advanced” position on the previous Pareto front. We run and average $M=2$ fine-tunings from each of those two initializations for $T=7k$ steps, before applying *LITI*. Our results indicate that a higher $\eta$ (0.5) performs better in regions of high $\mathrm{KL}$, whereas a lower $\eta$ (0.3) helps in regions with $\mathrm{KL}$ below 65. This suggests that the optimal choice for $\eta$ is compute-dependent; a lower $\eta$ is appropriate if further iterations can explore high $\mathrm{KL}$ regions, whereas a limited compute budget might benefit from a higher $\eta$. This resembles the learning rate trade-off in optimization, where lower rates improve results but require more training steps. As a final note, we can also use different $\eta$ for the different fine-tunings; notably, we observe that merging all those $M=4$ RLs perform better (though it doubles the compute).
Refer to caption: x29.png

### D.4 Interpolate towards the initialization? or towards the SFT?

Figure: Figure 17: Experiments ablating the initialization in *LITI*. We compare linear interpolating either towards the episode-specific initialization (i.e., the $\theta^{\eta}$ selected from previous iteration) or towards the SFT, which was the initialization of the \nth1 episode. The two resulting Pareto fronts are similar. However, in our iterative experiments we interpolate towards the episode-specific initialization as it allows maintaining a constant $\eta$ at each *WARP* iteration, enabling a smooth progression towards the high $\mathrm{KL}$ regions.
Refer to caption: x30.png

## Appendix E Addressing length bias in WARP

Figure: (a) Length bias in iterative *WARP*.
Refer to caption: x31.png

Problem: length bias.
We investigate a potential length bias in *WARP* []. LLMs after RLHF tend to be unnecessarily verbose because RMs often prefer longer generations to shorter ones, leading to this form of reward hacking. We confirm such a phenomenon in Figure 18(a), where the length of the generated text increases with higher $\mathrm{KL}$ values. This trend is even more pronounced in iterative *WARP*, where the \nth3 iteration generates longer sentences than the \nth1 iteration at the same $\mathrm{KL}$.

Mitigation strategy: length penalty.
To mitigate this length bias, we integrate a length penalty (LP) into the reward: $-0.0005\times\text{len}(y)$, following []. From the SFT init, we launch one RL fine-tuning run with this LP, highlighted with red stars in Figure 18(b). This LP leads to significantly shorter outputs as training occurs and $\mathrm{KL}$ increases, in contrast to policies trained without LP.

*SLERP* with different configurations.
Figure 18(b) also displays the lengths of generations from a *SLERP* merging of two policies, one trained with the LP and the other without.
Critically, merging policies from diverse training configurations not only mitigates the length bias but also improves the Pareto front, as illustrated in Figure 18(c). This improvement is likely due to the increased diversity in weights and predictive mechanisms across policies, which seems beneficial for generalization, as shown in supervised learning [].

Conclusion.
Those experiments highlight the possibility to fix the length bias, and also the benefits of merging policies trained with diverse rewards, supporting previous suggestions from [].

## Appendix F Diversity in predictions

Figure: Figure 19: Confirming diversity loss in RLHF. The $x$-axis is the $\mathrm{KL}$ compared to the SFT initialization; the $y$-axis is the similarity across two generations from a given policy when decoding with temperature $0.9$.
Refer to caption: x34.png

Finally, we investigate the loss in diversity across generations when aligning LLMs, as reported in [].
This could have negative consequences for creative or exploratory tasks, or even lead to policy collapse [, ].
In Figure 19 we plot the BLEURT similarity [] across generations, during REINFORCE, or in *LITI* (as we interpolate back towards the SFT initialization).
We observe that $\mathrm{KL}$ is strongly positively correlated with similarity across generations, confirming that RLHF induces a loss of diversity across generations.
This experiment confirms that optimizing the Pareto optimality between reward and $\mathrm{KL}$ enables to trade-off between alignment and other benefits from pre-training, such as diversity in generations.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Why should we have the Linear Interpolation Towards Initialization (LITI) step? This step takes an interpolation
from the merged model towards the initialization as the next iteration's initialization. Why do we need that? Can't that be implemented by adjusting the \beta in KL regularization with updated base policy and using only one iteration?
2. In Algorithm 1, is line 13 and the ouput using the same \eta?
soundness: 3
strengths: - The paper is clearly written and easy to understand, using exponential moving average of the base policy in RLHF should be less conservative than a fixing target.
- The empirical analysis is robust, featuring comparisons against state-of-the-art models and demonstrating WARP's effectiveness in balancing KL regularization and reward optimization.
- The paper considers a wide variety of tasks to demonstrate potential improvements.
summary: The paper improves the trade-off between KL and reward during RLHF by proposing Weight Averaged Rewarded Policies (WARP). WARP is designed to optimize the KL-reward Pareto front of solutions, by using three variants of WA at three different stages of the alignment procedure: Exponential Moving Average (EMA), Spherical Linear intERPolation of task vectors (SLERP) and Linear Interpolation Towards Initialization (LITI).
weaknesses: - Application of existing techniques to RLHF: The novelty is limited, the techniques used in the paper are almost proposed from the deep learning or reinforcement learning literature. For example, EMA anchor used ideas from trust-region updated reinforcement learning algorithms like TRPO [1]; SLERP uses idea from model merging by weight averaging [2,3]. LITI uses idea from WiSE-FT [4]. Given previous work that already applies similar ideas to reward modeling RLHF [5], the novelty is further weakened. A rigorous theoretical framework to justify why applying these techniques to the policy is better than reward is missing.
- Ablation studies: The paper introduces three different procedures, it's better to provide a detailed ablation studies to each procedure to verify no one procedure is useless.
- Training WARP is costly, requiring multiple iterations and multiple RL runs at each iteration. Also, more hyperparameters are introduced, making the algorithm more complex.

[1] Trust region policy optimization.

[2] Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.

[3] Diverse weight averaging for out-of-distribution generalization.

[4] Robust fine-tuning of zero-shot models.

[5] WARM: On the Benefits of Weight Averaged Reward Models.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Some experimental details are unclear:

- In Table 1, the reviewer would like to know who is responsible for evaluating the side preference rate. Is it conducted by GPT-4 or by human labelers? Additionally, could you report the response length? If judged by GPT-4, the win rate may be more easily biased by the response length.

- The reviewer would like to know what kind of dataset is used for RLHF. It is quite surprising to see improved performance in the benchmarks shown in Table 2. The results in the table seem to suggest there is no alignment tax.
- The main numerical evaluation in this paper is the reward. Since only Table 2 appears to show the benefit of KL (I assume), could the authors provide evaluation results of benchmarks in Table 2 at different points along the KL-reward trade-off curve, or for configurations with different hyperparameters?
soundness: 3
strengths: - This paper is well-written and easy to follow. Extensive related work is provided to help readers understand the research problem.
- The proposed methods seem to be simple yet effective in balancing the trade-off between KL divergence and reward.
- Extensive numerical results are provided to justifiy the proposed method.
summary: This paper studies fine-tuning language models to align with human preferences. It addresses the trade-off between KL divergence and reward, which may help mitigate the alignment tax issue. Specifically, it introduces Weighted Averaged Reward Policies (WARP), merging policies in the weight space using techniques such as exponential moving average and spherical linear interpolation. Experiments on Gemma models demonstrate the effectiveness of the proposed method.
weaknesses: - The main drawback of this paper is that techniques like exponential moving average and SLERP are well-known, which renders the technical novelty insufficient.
- Despite extensive numerical evaluation, it is unclear why model averaging helps balance the trade-off between KL divergence and reward. According to the reviewer’s understanding, achieving this trade-off effectively requires the policy to explore intelligently and update gradually during initialization. However, it is unclear how this relates to the proposed technique or if there are other possible mechanisms and explanations.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: See weakness part above.
soundness: 3
strengths: 1. This work has a good presentation and is easy to understand. 

2. The method is shown clearly. The intuition is given and the ablation study for each stage is given.

3. Many experimental results are given to support the method.
summary: This work gives a novel alignment strategy WARP. This method first uses EMA anchor to train multiple policies with RLHF, then use spherical averaging to average policies, and finally use Linear interpolation to combine with the init model. This process can be repeated. Many experiments have been done to show the effectiveness of this method.
weaknesses: 1. I think it is better to add some more introduction for weight averaging. Since I am not quite familiar with this, I may wonder if this is a method  conducting averaging by each parameter or some other techniques.

2. I hope to get some intuition for why WARP can reach a Pareto front, which seems non-trivial for me.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can the authors provide insights and more detailed explanations on why interpolating weights towards the initialization can attain a better Pareto front than the one revealed during RLHF?
soundness: 3
strengths: 1. **Innovative WA-based Approach to Alignment**: WARP introduces a layered model-merging strategy that addresses KL-reward trade-offs innovatively, leveraging EMA and SLERP in a structured way to boost model alignment with human preferences. Furthermore, the iterative application of WARP stages can progressively refine the KL-Reward Pareto front. 
2. **Empirical Validation**: The authors conducted comprehensive experiments to verify the effectiveness of the proposed four algorithmic components: EMA as a dynamic anchor, SLERP, LITI and iterative WRAP. Besides, comparison experiments on the Gemma 7B model show that WARP consistently outperforms open-source models, particularly in alignment and benchmark performance, providing strong evidence for its efficacy.
3. **Scalability and Adaptability**: The WARP process is compatible with distributed learning setups, enabling scalability through parallelizable fine-tuning iterations and potentially making it viable for broader applications in open-source and collaborative environments.
summary: The paper proposes a novel alignment technique based on weight average (WA)  for reinforcement learning from human feedback (RLHF) applied to large language models (LLMs). The proposed WARP approach is designed to enhance the balance between reward optimization and the retention of pre-trained knowledge, tackling issues such as catastrophic forgetting and reward hacking often faced in RLHF. The WARP technique integrates three model-merging stages: using the Exponential Moving Average (EMA) as a dynamic anchor in KL regularization, Spherical Linear Interpolation (SLERP) for merging fine-tuned policies, and Linear Interpolation Towards Initialization (LITI) to regain pre-trained features. These steps iteratively refine the KL-reward Pareto front, progressively enhancing model performance. Experimental results on the Gemma 7B model indicate that WARP attains a better Pareto front and achieves superior alignment and generalization across benchmarks, outperforming existing methods.
weaknesses: 1. **Increased Computational Cost**: The biggest weakness of WRAP is the computation aspect. The iterative nature of WARP requires multiple RL runs and merging stages, which could be computationally demanding, potentially limiting its practicality for resource-constrained applications.

Meta-review: The paper proposes a new RLHF algorithm that merge policies in the weight space to address the trade-off between KL and reward. This is a borderline paper without a strong champion. The disagreements  from the reviewers are 

1. Novelty: Many techniques used here have already been proposed in RL, but the authors argue that they have not been used in the RLHF context with LLM.

2. Complexity of the proposed method. The proposed algorithm has 3 steps. Reviewers have raises issues on the missing theoretical rigorousness on these steps are designed and support on why they indeed work as the authors claim intuitively. There're also debates on whether the ablation experiments are sufficient to make up for the lack of theory. 

3. Increased computed cost. The proposed require training multiple policies, which cost more than some other RLHF finetuning scheme. Nonetheless, the authors suggest this is a feature of scaling post-training. 

For point 1, I think identifying the effectiveness in the LLM context warrants its own novelty. For point 2, however, I agree with the reviewers that not enough evidence is provided to support the intuition claims. There're not full theory that directly analyze WARP's performance but just lemmas for different parts at a high level. More explanation is needed also to explain why these 3 steps are designed to combine in the current way. To fully  support these claim empirically without theory, I think experimental results on multiple LLMs, datasets and labeling schemes are needed. For point 3, the authors should provide more experimental results to support the effectiveness of the scaling argument, to justify indeed the proposed way is a better way to scale the post-training compute. 

Overall, I think there is some room for improvement from the current paper to warrant a clear acceptance.

Decision: Reject