Language Model Alignment with Elastic Reset

Summary

This paper proposes a new approach for fine-tuning language models (LMs) with RL in order to achieve a good trade-off between maximizing reward and minimizing drift from the initial model, which is typically undesirable since it results in losing some capabilities while acquiring others. The approach consists in resetting the model to an exponentially moving average (EMA) of itself, while also resetting the EMA model to the initial one. The method is simple and demonstrates good performance on three different settings with varying degrees of complexity, including RLHF on top of LLaMA-7B for a QA task.

Strengths

1. The proposed approach is simple and effective. 2. The evaluation is quite thorough considering multiple settings and thus demonstrating the generality of the approach. 3. The paper is well written and the problem setting is clearly stated. 4. The problem setting is of high importance to the community. The authors point out a major limitation with existing RLHF methods and propose a method for improving upon it.

Weaknesses

1. The main limitation of this study is the evaluation protocol used to assess the results. The model's performance is evaluated by the same reward model used to train it. However, it is well known that LMs fine-tuned with RL(HF) are prone to overoptimizing their reward (models), overfitting and actually performing badly when evaluated by humans (which is a more robust metric and ultimately what we care about for many applications). While I understand that performing human evaluations can be expensive, it is very difficult to assess the validity of these results otherwise. It could be the case that Elastic Reset is a more powerful optimization approach that overoptimizes the reward better than PPO i.e. can obtain high reward during training but this performance doesn't actually transfer well to unseen prompts when evaluated by humans. At the very least, I suggest using a different reward model for evaluation such as a different base model (of similar size) trained on the same data or the same model trained on a different dataset such as the summarization data from [1] or the HHH dataset from [2]. You could also hold out part of your data and train a separate reward model with a different base on it in order to bring it more in-domain. 2. Can you include experiments with varying reset intervals for the all tasks? It's important to know how sensitive the model is to this parameter and better understand if similar / same values work across different tasks. Do you have any suggestions or insights for selecting this hyperparameter for new tasks? 3. It would also be interesting to see how the results change if only some of the parameters are reset e.g. the last few layers as it typically done when fine-tuning LLMs. References: [1]. Stiennon et al. 2020, Learning to summarize from human feedback. [2]. Bai et al. 2022, Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback.

Questions

1. Why don't you show results with the same KL coefficients for PPO and Elastic Reset in Figure 5? Do smaller values of KL not work at all for PPO? What about larger values for Elastic Reset? 2. You mention that Elastic Reset may also work without the KL penalty but you never test this hypothesis. Can you include an experiment with Elastic Reset and KL coefficient of 0? For a fair comparison, it would also be interesting to compare with PPO without the KL penalty, which I assume will perform much worse. 3. Figure 6 is rather noisy, can you include a smoother version in the appendix to better see how the two methods compare with each other? 4. In Figures 5 and 6, the light blue line can barely be seen on printed paper, can you please use a different / stronger color?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes, the authors clearly state the limitations of their study at he end of the paper.

Summary

This paper aims to address the problem of language drift issue during RLHF (reinforcement learning with human feedback), which is also known as alignment tax and reward hacking. The problem is that during RLHF process, the model can "overfit" to the given suboptimal rewards while forgetting some important skills such as linguistic capabilities. The authors proposed a method named Elastic Reset that periodically reset the online model with an exponentially moving average (EMA) of its previous checkpoints. The authors use a few tasks to showcase the effectiveness of the proposed Elastic Reset method on top of GPT-2 and LLaMA-7B, and shows that it is better than vanilla PPO with a reduced KL penalty.

Strengths

The method is rather simple and easy to implement. It just needs to save the parameters of previous checkpoints during training, so that every $n$ steps, one can compute an EMA and merge it to the current model. The empirical results suggest that the proposed method is effective on the selected tasks and datasets, outperforming the simple PPO and NLPO baseline methods.

Weaknesses

- The method itself is not particularly novel. Using Reset mechanism and EMA of model parameters to mitigate overfitting is rather common, although it might be a novel application in RLHF. - The "drift" problem is not very clearly defined and formulated. The motivation is not well justified. Many descriptions are references while there is no concrete experiments and case studies to show the problem of drifting clearly. - I believe the key issue that this paper wants to address is essentially the same to many continual learning problems -- learning new knowledge while not forgetting the acquired skills. Therefore, many CL methods such as experience replay, regularization, EWC (Elastic weight consolidation), should also be applicable. But none of them is mentioned in the paper. The authors focused too much on the RLHF literature, using newly invented terms, however, ignored that the key challenge can be formulated with existing problem setup and can be addressed by existing techniques. - The selected tasks and datasets are quite narrowed. The experiments on GPT-2 and even smaller models are also not that convincing in that RLHF is rarely used on such LMs. I suggest authors can replace those experiments with more commonly used datasets to support the claims.

Questions

- Line 219, why do you "greatly reduce the KL coefficient $\beta$"? From the description, it seems that you want to amplify the drift issues, but it seems choosing different $\beta$ can significantly influence both PPO and Elastic Reset. How do you decide choosing which $\beta$ to compare your method and other baselines? Do you have to decide this before you see the real test data?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors use Section 8 to describe the limitations.

Summary

This paper proposes Elastic Reset, a simple technique for countering language drift and reward model overfitting when optimizing a language model policy against some communicative reward via reinforcement learning, as is done in RLHF. The idea of elastic reset is to periodically reset the trained model to an exponentialy-weighted moving average (EMA) between the initial pretrained model and the online model trained since the last reset. The authors demonstrate on a variety of RLHF tasks from relatively simple (pivot translation) to relatively substnatial (StackExchange QA) that Elastic Reset seems to be a simple way to mitigate language drift which outperforms other approaches, including the commonly-used KL penalty in RLHF. They propose a nice way of interpreting the tradeoff by using a "pareto frontier graph", which plots methods' downstream task reward and measure of language drift on x/y axes, similarly to ROC curves, and demonstrate that for some given compute budget, Elastic Reset dominates existing baselines like REINFORCE and PPO with KL penalties. Overall there's a lot I like about this paper, but there are some issues I have with the experimental setup that prevent me from unconditionally recommending the paper for acceptance. If my concerns are clarified or resolved, I am willing to update my score and look forward to the author response.

Strengths

- A simple technique that appears to give gains across a variety of RLHF benchmarks at varying scales (from translation to stackexchange Q/A). I am impressed by the breadth of experiments in this paper. It seems this could be one of those simple tricks that practitioners find useful and widely deploy in future RLHF pipelines (but time will tell whether the results are robust enough). - Good comparison to other sensible baselines, with some sensible ablations on the IMDB mock sentiment task, but some areas for improvement here (see Weaknesses) - I like the Pareto figures, reminiscent of ROC curves, which demonstrate the tradeoff between task performance and language drift and how to identify an optimal method under a given practitioner's constraints. I agree with authors that this is the right way to think about language drift, though I have some points of confusion (see firsrt Weakness)

Weaknesses

## Unclear how robust elastic reset under different compute constraints - The IMDB and StackLLaMA experiments demonstrate something subtly worrying in my view: they show that there are portions during training where elastic reset *doesn't* help over comparable baselines. For example, in Figure 4 a/b, training prior to the first two resets demonstrate a "spiking" behavior where the model overfits to the reward and has higher language drift than the other two methods. It is only after the 2nd reset, untnil the end of training, that the reset causes task reward to rise higher than existing methods *but*. - This feels "lucky" to me, and it is hard to measure how robust elastic reset is under different compute constraints. For example, lets say we only have enough compute to train the model up to (but not after) the 2nd reset, i.e. epoch ~33 in Figure 4. It seems like in this case we would *not* prefer to use the elastic reset model, since it seems to still be in the regime of overfitting, and it is only after the 2nd reset that things start to look better. - A similar concern exists for StackLLaMA where there is not a significant difference between Elastic reset and PPO, until around ~500-600 epochs, when there suddenly is (and perhaps there's even a regime in 400-500 where Elastic Reset is overfitting). - To address this concern, it seems like the pareto graphs, and in general the performance of elastic reset, need some notion of compute budget to more accurately assess when it is appropriate to use elastic reset. For example, what does the pareto frontier graph look like if we only consider training up to about epoch 33 in Figure 4? (i.e. before the 2nd reset)? Is elastic reset still the preferred choice? Would changing the number of resets change anything? (Perhaps I'm misinterpreting the pareto graphs here). - A more detailed analysis of how many resets are needed and how robust elastic reset is to the timing of rests would partially alleviate these concerns (see next point) ## Could use more carefully controlled baselines - The baselines could be clearer. If my understanding is correct, the *only* way in which elastic reset diverges from traditional RLHF is by periodically resetting the model to be the EMA of the past n model steps. As authors say, this is a strength of the method, but the comparisons tend not to directly compare to a model with/without elastic reset. For example, considering Elastic reset with 3 resets during training, the sensible baseline is to compare to the exact same method (e.g. REINFORCE), with the same hyperparameters, just with 0 resets during training. But the baselines in the paper seem to always have a slight confounding factor, e.g. for the pivot translation task L152, a KL penalty for Elastic Reset is added on top of the REINFORCE baseline that does not seem to be present in vanilla REINFORCE; same for L219 on top of PPO (the beta parameter is reduced, and the decision to let the PPO model "drift more" is not well justified); only the StackLLAMA experiment in Section 7 appears to stay consistent (L280), but the improvement of elastic reset here is not as clear. - More generally it would be better if, instead of a "1 vs many" comparison where a single implementation of elastic reset is compared to a single implementation of PPO and REINFORCE baselines, a "paired" comparison strategy was adopted, where for each method and specification of hyperparameters for each baseline, authors measure the effect of elastic reset on top (as authors say, it seems simple to just add this onto any arbitrary alignment technique, even not necessarily RLHF). Does Elastic Reset improve consistently over other RLHF methods, keeping hyperparameters consistent? Of course, it doesn't have to improve all the time, but knowing when it helps and when it doesn't (because maybe the baseline method by itself keeps language drift under control) would be nice. - Again, Figure 5c is lacking some context. The Elastic Reset figure is built on PPO with beta = 0.01. I appreciate the KL ablation but does Elastic Reset dominate PPO across all KL penalty values kept constant? Otherwise what is the reason for setting the KL penalty lower for elastic reset? - Again, to address this concern, I would love an investigation of how many resets affect performance, and any guidelines for choosing a set number of resets, in the same style as the nice ablations in Figure 5. Authors state that it is difficult to find a heuristic for how often to reset (L307), but graphs showing the effect of performance for different reset timescales and/or compute would be enormously informative, and help address the first main weakness I outlined above. ## Mathematical connection/intuition as to why elastic reset differs from distillation and/or KL - I get the feeling there are some mathematical connections to both KL penalty and alternating RL/distillation algorithms that are not fully explored in the paper. I haven't thought too deeply, but, for alternating RL/pretraining methods (S2P; Lowe et al., 2021) and iterated distillation, one can think of the distillation/SL phase as precisely doing a sort of model reset to the initial pretrained model via gradient descent to minimize KL between the online model and the initial model. For RLHF with KL penalties, we can also view this as a sort of bayesian inference, computing a model average of the prior pretrained model and the posterior (online) model trained via RL ([Korbak et al., 2022](https://arxiv.org/abs/2205.11275)). The number of steps taken is a hyperparameter that, if chosen carefully, results in a model after distillation that is likely some average of the online model and the initial pretrained model, as in elastic reset. If elastic reset is a simpler or more efficient way of doing the same thing as such a distillation phase, this is a great thing, but it's not quite clear to me now. Could authors clarify any differences between distillation and elastic reset? In particular, why might we expect elastic reset to be **more** performant than than alternating RL/distillation or a KL penalty, given that they seem to have the same motivation (which might even be mathematically formalized)? ## Minor - L99 "Elastic Reset takes inspiration from both of these"—it's not quite clear what "both" refers to. - Some qualitative examples of model outputs in the appendix could potentially be nice to have in the paper, perhaps even identifying what model outputs look like at different points on the pareto frontier curve for different training methods.

Questions

Main questions in Weaknesses section. Some more minor questions: - L100 "EMA on CPU"—are authors sure this is more efficient? If EMA is on CPU, don't you need a full model copy from GPU to CPU for each step? This seems expensive, and suggests that Elastic Reset would also benefit from keeping the EMA on the GPU. (You could theoretically keep the pretrained model used for KL penalties on the CPU and copy over, but this would significantly slow things down) - Authors could add some intuition (or experiments) that discuss how important is it that the reset resets to an **exponentially-weighted** moving average? Why not just a literal model average between the new model and the pretrained model?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

yes

Summary

Finetuning language models with reinforcement learning (RL) using human feedback, i.e. RLHF, has emerged as a promising paradigm for aligning large language models (LLM) to human preferences. Though RLHF has shown promising results when training models such as ChatGPT, RL has some inheritance drawbacks. Merely optimizing a reward model trained on human preferences can degrade performance, known as reward hacking, alignment tax, or language drift in the literature. This paper argues that the standard way for addressing this is insufficient and proposes a new idea. In particular, the authors propose, Elastic Reset, a technique to reset the model weights to address reward hacking.

Strengths

Addressing reward hacking is an age-old problem, and many solutions have been proposed. The strength of this paper is that the solution proposed has several key benefits: - Easy to implement: The authors implemented their idea on top of several existing RLHF frameworks and showcased the benefit of Elastic Reset to existing frameworks RL algorithms. - Written Presentation: The author's presentation of the proposed algorithm, explanations, and ablations studies were justified and thoroughly explained. - Experiments: The authors presented several experiments across three very different domains and showcased the benefit of their proposed method. - Clear Definition of Language Drift: Often language drift is not clearly defined, but found very clear examples to showcase language drift issues.

Weaknesses

Although the proposed algorithm is straightforward to implement and has shown good results across three challenging datasets, the approach has several weaknesses. - Resetting on-policy vs. off-policy: Resetting weights are typically done with off-policy algorithms because the policy has a problem exploring when increasing the replay ratio. Whereas for on-policy to deal with exploration, we typically add temperature parameter, entropy loss coefficient, or some other exploration bonus. Instead, in the on-policy case, the authors use reset to keep the policy close to the original policy so it does not experience language drift, which seems like the opposite use-case of the original intent. - Second reset: The first reset in which EMA is reset to the initial model means that you are still searching around an epsilon ball around the initial model. I am unsure why this would be much better than KL divergence, which explicitly does this.

Questions

- How did you decide on resetting of the first stage? Given stage one, how did you decide on resetting the second stage? - Why did you include the supervised learning (SL)+PPO results in your experiments? Training RL from scratch is known to be a difficult task. Most RLHF success stories warm start the RL model using SL. The results in the paper seem to use the RL with PPO results from the GRUE benchmark. But SL+RL results are always stronger. [Algorithm] - Sentiment Score - Perplexity [PPO +SL] - 0.626 - 35.045 [NLPO + SL.] - 0.611 - 33.82 Is the proposed approach is not compatible with RL+SL?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Post-Rebuttal Review

I thank the authors for their detailed answers to my questions and appreciate that they ran additional experiments as suggested, which seem to support the claims in the paper. Regarding the reward model used for evaluation, I agree with the authors that the experiments support some degree of generalization, particularly the ones on StackLLaMA. However, I was not convinced by this statement "Testing with a new reward model trained on held-out examples is possible but has not been done in previous literature and would deviate our results from existing baselines (by using a different training reward model without held-out examples)." First of all, more and more papers are using multiple models for evaluation ([1], [2], [3] to name a few) so I expect this to soon become standard practice, and in any case it provides a more robust evaluation. What I'd propose is to use two different reward models, the one that you already have which would allow easy comparison with existing literature, and another one that would strengthen the robustness of your evaluations. I also agree with reviewer tYGJ's comment that evaluating the approach on more challenging tasks beyond the IMDB one would further strengthen the paper and make the results more convincing and relevant to practical applications. In conclusion, I am willing to increase my score to 6 conditioned that the authors commit to including evaluations using another reward model to ensure the results are robust to this choice. References: [1]. AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback. Dubois et al. 2023. [2]. LIMA: Less Is More for Alignment. Zhou et al. 2023. [3]. Secrets of RLHF in Large Language Models Part I: PPO. Zheng et al. 2023

Thanks

Thanks to authors for the response. Upon looking at Figure 4a/b closely it seems like there are indeed still regimes where under 1 or 2 resets, doing elastic reset allows one to get higher reward with less drift. Figure 1(c) in the supplement is helpful for illustrating this point and would be useful to include in the supplementary material. The other responses addressing questions regarding elastic reset intuition, KL penalties, etc are also helpful. I'm raising my score to a 6.