Language Model Alignment with Elastic Reset
Authors: Michael Noukhovitch
Samuel Lavoie
Florian Strub
Aaron Courville
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning from human feedback (rlhf)
language
TL;DR: In RLHF, Elastic Reset iteratively sets your model to its EMA and resets your EMA during training to achieve better reward for reduced alignment tax

Abstract: Finetuning language models with reinforcement learning (RL), e.g. from human feedback (HF), is a prominent method for alignment. But optimizing against a reward model can improve on reward while degrading performance in other areas, a phenomenon known as reward hacking, alignment tax, or language drift. First, we argue that commonly-used test metrics are insufficient and instead measure how different algorithms tradeoff between reward and drift. The standard method modified the reward with a Kullback-Lieber (KL) penalty between the online and initial model. We propose Elastic Reset, a new algorithm that achieves higher reward with less drift without explicitly modifying the training objective. We periodically reset the online model to an exponentially moving average (EMA) of itself, then reset the EMA model to the initial model. Through the use of an EMA, our model recovers quickly after resets and achieves higher reward with less drift in the same number of steps. We demonstrate that fine-tuning language models with Elastic Reset leads to state-of-the-art performance on a small scale pivot-translation benchmark, outperforms all baselines in a medium-scale RLHF-like IMDB mock sentiment task and leads to a more performant and more aligned technical QA chatbot with LLaMA-7B. Code available https://github.com/mnoukhov/elastic-reset

Language Model Alignment with Elastic Reset
Michael Noukhovitch Mila, Université de Montréal
Samuel Lavoie Mila, Université Montréal
Florian Strub Google Deepmind
Aaron Courville
Mila, Université de Montréal CIFAR AI Chair
Abstract
Finetuning language models with reinforcement learning (RL), e.g. from human feedback (HF), is a prominent method for alignment. But optimizing against a reward model can improve reward while degrading performance in other areas, a phenomenon known as reward hacking, alignment tax, or language drift. First, we argue that commonly-used test metrics are insufficient and instead measure how different algorithms tradeoff between reward and drift. The standard method modified the reward with Kullback-Lieber (KL) penalty between the online and initial model. We propose Elastic Reset, a new algorithm that achieves higher reward with less drift without explicitly modifying the training objective. We periodically reset the online model to an exponentially moving average (EMA) of itself, then reset the EMA model to the initial model. Through the use of an EMA, our model recovers quickly after resets and achieves higher reward with less drift the same number of steps. We demonstrate that fine-tuning language models with Elastic Reset leads to state-of-the-ee performance on a small scale pivot-trane translation benchmark, outperforms all baselines in medium-scale RLHF-like IMDB mock sentiment task and leads to a more performant and more aligned technical QA chatbot with LLaMA-7B. Code available at wwwpuuu...ommooooooollleaeee
Introduction
Dialogue agents that can effectively interpret and use language are long-term challenge for NLP. The rise of large pretrained language models (LMs) [Brown et al., 2020] made language model finetuning one the most promising research directions to achieving capable dialogue agents [Bender and Koller, 2020]. Recently, reinforcement learning (RL) has become key ingredient of finetuning large LMs for interaction with humans [Ziegler et al., 2019, Ouyang et al., 2022, Bai et al., 2022], notably shown in ChatGPT [OpenAI, 2022]. A reward model is learned on the alignment objective, such as learned human preferences [RLHF; Christiano et al., 2017, Stiennon et al., 2020], and the language model is finetuned optimize the reward. But training from its pretraining and can reduce performance on important benchmarks [Ouyang et al., 2022] and drifting away from natural language syntax and semantics [Lazaridou et al., 2020].
the RL objective moves the model away
"Language drift" [Lee et al., 2019, Lazaridou et al., 2020], "alignment tax" [Askell et al., 202 I], "reward model overoptimization" [Gao et al., 2022], or LM-specific "reward-hacking [Clark and Amodei, 2016] is inherent to RLHF, the extreme case, models learn to achieve high reward by generating nonsense text that is unintelligible to humans [Lewis al., 2017]. Methods to mitigate this issue range from re-rum re-running pretraining [Lowe** et al., 2021], I. grounding in other modalities [Lee et al.,
"Correspondance michael. wwo.hhoritcctummnzzeell.c
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Maintain an EMA
Elastic Reset every sleps
EMA
Online
Online
EMA
EMA
Pretrained
EMA Updata
Gopy
Copy
Figure 1: Elastic Reset. In actor-critic RL, we reset the policy but maintain the value function.
2019], masking the LM generation [Ramamurthy al., 2022] and iterated learning [Lu et al., 2020]. But the standard, and by far most popular approach, adds a Kullback-Lieber (KL) divergence penalty the reward in order to prevent the finetuned model from drifting too far from the pretrained model [Jaques et al., 2017, 2019, Ziegler et al., 2019J. Still, all methods are insufficient over a large-enough training horizon so models are early-stopped before reaching catastrophic level of drift.
Gao et al. [2022] find that achieving reward is proportional to drift from the initial model, but that not all drifts are equal. We wish to make small but effective changes that achieve high reward but maintain capabilities, yet auxiliary losses such as the KL. penalty don't 't seem improve this tradeoff and only serve to slow down training [Gao al., 2022]. We posit that RLHF training requires useful inductive bias that does not modify the training objective. Inspired by recent work in generalization for image classification [Zhou et al., 2022], sample-efficient RL [Nikishin et al., 2022, D'Oro et al., 2023], and inducing compositional language [Li and Bowling, 2019], we propose to use resets. Similar to iterated learning [Kirby, 2001], iteratively resetting model has been shown reduce overfitting in language and RL scenarios [Rita et al., 2022]. In this work, we show iteratively resetting a model also reduces drift while attaining equal or better reward than just a KL penalty.
Unlike previous work in sample-efficient RL [Nikishin et al., 2022], RLHF is typically on-policy so it does not maintain replay buffer with which to bootstrap learning after reset. In lieu of a replay buffer, we reset the policy but maintain the value function. Yet resetting the policy to its initial state can still cause large drop in performance. So we propose resetting to model in-between and initial state, specifically to an exponential moving average (EMA) of our online policy, as EMA has been shown be highly performant [Caron et al., 2021]. We still expect our EMA model to slowly drift, so we add a second step where we reset the EMA model to the initial model. We call this overall method Elastic Reset and illustrate it in Figure 1. Elastic Reset is implemented top of regular RL methods such as REINFORCE [Williams, 1992] or PPO [Schulman et al., 2017].
online
First, we test our method benchmark for drift, we outperform all previous baselines and demonstrate state-of-the-art perfor- mance. Next, re-evaluate how performance is measured in the field and argue for metrie of how each method trades off performance vs drift. We propose the Pareto Frontier Graph, a graphical measure that illuminates the trade-off between performance and drift and demonstrate that Elastic Reset dominates the baselines against this trade-off. Then, we scale up slightly to GPT2 and work on popular task closer RLHF, IMDB mock sentiment. Comparing to all baseline methods, we again show state-of-the-art performance on the benchmark. Through ablations, we show that Elastic Reset is robust to choices of byperparameters, even more so than baselines. Finally, we scale up even more to true RLHF finetuning of Llama-7B in order to create a helpful technical QA chatbot using a StackExchange dataset. We again outperform the baseline, demonstrating how Elastic Reset mitigates the alignment tax while better optimizing the human feedback reward.
small scale task: pivot translation with transformer. In this classic
2 Related Work
"It is often difficult or infeasible to capture exactly what we want an agent to do, and as a result we frequently end up using imperfect but easily measured proxies" [Clark and Amodei, 2016]. In RL, this proxy is how we construct our reward and the consequence can be Trward-hacckng and Amodei, 2016]; an agent optimizes the reward but does not accomplishing the meaningful task. RLHF aims to align an agent with human preferences while maintaining the capabilities of the pretrained model, but uses a learned reward model as proxy of human preferences [Christiano et al., 2017, Ziegler et al., 2019]. This can lead to LMs that optimize reward model but degrade in
[Clark

---

performance on general NLP benchmarks [Askell et al., 2021], overfit the reward model and do not generalize to true human preferences [Bai et al., 2022, Gao et al., 2022], or latch onto confounding factors hidden in the reward [Stiennon et al., 2020]. These effects are exacerbated if the reward model is updated during training such as in iterated RLHF [Bai et al., 2022] or related setups such as emergent communication [Lazaridou and Baroni, 2020], end-to-end dialogue [Lewis et al., 2017], learning RL policies through latent language [Andreas et al., 2018], and pivot translation [Utiyama and Isahara, 2007]. There, the phenomenon is known "language drift" [Lee et al., 2019] and can lead to incoherent and unnatural linguistic outputs Lewis et al. [2017].
This phenomenon is inherent to RLHF, Gao et al. [2022] show that improvement on alignment/ / reward is proportional to drift from the initial model, but also find that different methods and design choices achieve different proportions of performance to drift. Therefore, a major challenge of RLHF how to learn the reward in such a way as to minimize the drift, alignment tax, and sward-hacking. The standard approach used in most RLHF is to incorporate a KLL penalty between the training language model and some fixed model [Jaques et al., 2019, Ziegler et al., 2019, Stiennon et al., 2020, Ouyang et al., 2022, Steinert-Threlkele et al., 2022, Bai et al., 2022], usually the initial, pretrained model. Less common is to add the original pretraining task to the finetuning objective (termed S2P by Lowe* et al. [202]] 1) but this can be compute-intensing and requires maintaining the pretraining data which may be even more expensive for larger models [Brown et al., 2020]. On small-scale pivot-translation, Lu et al. [2020] propose iterated learning with student-teacher distillation but it too is relatively compute intensive. Recently, Ramamurthy et al. [2022] propose to maintain a delayed masking model and mask the LM to output only the top-p tokens. Elastic Reset takes inspiration from both of these, using an iterated process and maintaining an EMA model. Apart from better performance, our method is more space efficient and maintains the EMA CPU whereas both other methods require maintaining an extra model on GPU. It is also more compute efficient as resetting weights and EMA updates are very cheap operations, whereas Lu et al. [2020] requires long distillation phase and Ramamurthy et al. [2022] requires an extra forward pass with the masking model. There exist less- less-popular methods that have been applied to similar issues in RL: prompt-tuning [Singh et al., 2022], using fixed model to generate many options and re-ranking using a reward model [Lazaridou et al., 2020, Meta FAIR Diplomacy Team et al., 2022], and grounding the output in a separate modality [Lee et al., 2019] but none have been used for RLHF or at scale.
Our method is inspired by recent works that leverage resets for single agent RL [Nikishin al., 2022, D' O'Oro et al., 2023], image classification [Zhou et al., 2022], and emergent communication [Rita et al., 2022]. Those works generally train from scratch and reset to random initializations in order to improve generalization. Our scenario requires resetting to pretrained models and focuses on improving the tradeoff' between performance and drift from this pretrained model. Elastic Reset can be seen as an on-policy alternative to Nikishin et al.'s 's [2022] off-policy resets, whereas they maintain the old replay buffer, we maintain the value model and an EMA of our policy.
Finally, the pretrain-thes RL-finetteee setup with the goal of maintaining pretrained knowledge can be seen two-step, RL-specific instance of continual learning and therefore language drift has links to catastrophic forgetting [McCloskey and Cohen, 1989]. There is a clear similarity between mitigation methods: rehearsal [Robins, 1995] or experience replay [Rolnick et al., 2019] is equivalent to multitasking with the pretraining objective [Lowe" et al., 2021] and weight-update regularization (Kirkpatrick et al., 2017] has similarities to KL regularization [Jaques et al., 2019].
3 Elastic Reset
The standard method against drift is a KL. penalty, generally between the learning policy @ and the initial, pretrained model fg- is calculated empirically over the minibatch of training inputs x and outputs and used as an auxiliary reward with coefficient |3 on top of the regular reward model r
To To(y|2) 2)
R(x, v) = r(z, v) 33 log
no. (1) (r)
For Elastic Reset, we maintain an exponential moving average 0 of our learning model 0 and choose a decay hyperparameter parameter 7) We initialize 0 + Og and after every online model step, we update our EMA model 0 t- (1 m7)0 + n0. Every n steps, Elastic Reset sets the online model to the
3

---

French
English
o.
German
StanExchonge
LLatA4 78
Worse Aroower
Transformer
Transformer
KocAy11 wit aigmment
This movie was
Witty and fun
Score
Sentment
below cass?
funciiion
OO LL 78
Voeginscadee
LLaMA
Roward
GPT-2
Frozon DEST
Figure 2: The Translation Game (top left), IMDB moek sentiment task (bottom left), and LtaM (right). We show all RL finetuning setups and StackLLaMA's reward modelling (top right).
EMA model 0 +- 0 and sets the EMA model to the initial model 0 t Bu. As with other methods, Elastic Reset can be easily combined with a KL penalty.
Translation Game: Careful Comparison to SOTA
Setup We first investigate the pivot-translation benchmark of Lee al. [2019], which was previ- ously popular for small-scale methods countering drift. Two translation models, French to English (FR-ENN and English to German (EN->DE) the models are finetuned on translating French to German through English (FR-EN-DEE) but given only paired French and German data from Multi30k [Elliott al., 2016, 2017] as shown in Figure 2. The models are not given English at finetune-time so the challenge is optimizing FR-DE while maintaining fluency in the intermediate English. Whereas larger benchmarks have only proxies for drift, we can exactly measure the performance degradation in our setup with the standard translation metric BLEU on held-out FR-EN validation set. Similarly, we measure success on the task with the FR-EN- BLEU score. Each model is an encoder-decoder Transformer [VVaswani et al., 2017] with 6 layers and all experimental details are available in Appendix A.
pretrained IWSLT [Cettolo et al., 2012]. Then,
Baselines The EN reward model is simply trained using cross-entropy between predicted and true DE. Our lower-bound baseline FROZEN ENGLISH, we freeze the FR-EN model and only update the EN-DE model. This models is guaranteed not to drift, but also cannot reach the best possible performance. For that, we need to update FR EN by backpropogating through the discrete EN tokens. We follow Lee et al. [2019] and train FR-EN using REINFORCE [Williams, 1992] to estimate the gradient. As our base model, we combine R REINFORCE with an exponentially moving baseline and, with previous work, add loss for entropy regularization.
When both FR ER and EN-DE are being updated, we tend to see reasonably large drift and we compare to the best previous methods that counter it on this benchmark. We follow Lee et al. [2019] to simulate the standard KL penalty method, KL PENALTY by training an LSTM LM on IWSLT English text and adding KL penalty with 3 = 0.05 to regularize the FR-EN model. MULTITASK learning, re- re-training, or S2P [Lowe* et al., 202 J. adds the supervised FR-EN objective on IWSLT pretraining data as an auxiliary task for the FR-EN model. Finally, implement Sceded Iterated Learning [SIL; Lu et al., 2020], which alternates between n finctuning steps and m steps of teacher-student distillation. FR-EN and EN-DE "teacher" models are finetuned on the translation game, then each distills knowledge into "student" model of itself, and finally the students are initialized as teachers for the next iteration. ELASTIC RESET is implemented
top of REINFORCE with very minimal KL penalty / = 0.001 and uses an EMA decay = 0.99. We run all models for 50k updates and reset every 23k steps to get 2 resets / 3 iterations within a run. Hyperparameters may differ between methods, e.g. 3, because we used minimal search to find the best hyperparameters for each method.
Experiments For each method, we run 5 seeds and plot the validation scores over training for the end-to-end task score, FR-EN-+DE BLEU, and drift score, FR-EN BLEU, in Figures 3b, 3a respectively. Following Lee et al. [2019], we also show the final validation score in Table 1. As a sanity check, FROZEN ENGLISH does not drift but also does not achieve very high task performance. line with previous results [Lu et al., 2020], all learning models initially improve FR-EN performance, likely because models are quickly, semi-supervised adapting from their pretraining (IWSLT) to the distribution of the finetune dataset (Multi30k). Afterwards, they start to overfit on their objective and FR- performance degrades. REINFORCE achieves the best

---

38
38
32-
FR--EN--
FR-EEN (drift)
Blou F8--
14
2.6
Frozon Englih
26
RENPORCE
NuRteok
20
Penalty
39
10000
50000
10000
20000
50000
Steps
Steps
FR-EN-DE (perf)
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 3: Comparing Elastic Reset to all baseline methods on the Translation Game We measure (a) Task Performance with FR-ENN-D BLEU and (b) Language Drift with FR->N BLEU, on the validation set during finetuning. We plot the mean and standard error over 5 seeds. To compare how methods trade off the two metrics, we plot (c) the best achieved drift vs task performance.
Table 1: Translation Game final validation scores
FR- DE
1 FR-+NN
FROZEN ENGLISH
30.8+0.2
36.3+-0.1
REINFORCE
33.2:10.3
29.611.33
+ SIL
28.2.1..4
27.314.4
MULTITASK (S2P)
32.2.1.33
35.211.0
KL PENALTY
33.2+0.2
30.8+0.4
ELASTIC RESET
32.9+0.1
36.3+0.1
possible task performance but drifts significantly. Despite extensive hyperparameter tuning and correspondence with the original authors, SIL does not manage to outperform the REINFORCE baseline so we exclude from the figures for visual clarity but show values in Table as well as full results in Appendix A. In line with previous work [Lee et al., 2019, Lu et al., 2020], we find that MULTITASK and KL PENALTY are both beneficial to reducing drift, but both represent tradeoff. Whereas MULTITASK strongly reduces drift, it does not achieve high task score. In contrast, KL PENALTY achieves a high task score but drifts quite drastically. Elastic Reset achieves nearly the best possible task score while maintaining the same drift score as the initial model. Visually, we see that our method track the baselines until the reset 23k steps. After the reset, we see slight performance drop but also big jump back in terms of FR-EN drift. While the task performance recovers within 5k steps, the drift performance does not degrade to previous levels. For the second reset, the EMA model is slightly more drifted and so the reset is less pronounced for both task and drift, leading to faster task recovery but slightly more drift. Overall, Elastic Reset shows state-of-the-an results on the benchmark and outperforms all previous small-scale methods.
Pareto Frontier Graph
Simply evaluating validation curves side-by-side or looking a table of final scores, it be unclear which method is better if one drifts less but the other achieves a higher reward e.g. MULTITASK vs KL PENALTY in Table 1. Previous work this [Lee et al., 2019] and other benchmarks (Ramamurthy et al., 2022] compare methods using simple point-estimates after training for specific number of epochs. But this number of epochs quite arbitrary models never fully converge reward, they are carly-stopped such that drift is not catastrophic. Since diffferent setups may admit different levels of drift, we believe that evaluation should reflect the continuous tradeoff between task and drift. We extend Ziegler et al. [2019], and create pareto frontier graph to plot each method's 's achieved task score drift metric on the validation set over training. We believe practioners will wish to choose the best model for some given task performance so, contrary to Ziegler et al. [2019], we plot the task score on x-axis and drift score on the y-axis. Improvement on drift metric can either mean lower scores (perplexity) to top such that, graphically, a better method will functionally dominate a worse method. We plot the best achieved reward vs drift over all validation steps for the Translation Game in Figure 3c. Not only
higher scores (BLEU) so we always plot task score as increasing from bottom
5

---

Fusik REaCE
Start
Reset
13.1
5061K Reset
Reszt
Start
32.4
Reaet
Tese122
Perslaaity (d00)
0.32
32.
0.36
22.0
0.34
Stut
Sast:
12.4
780
Most2
6.52
12.2
10
30
10
20
40
0.50
252
0.54
0.59
0.58
0..00
0.52
Epoch
Epoch
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 4: Plotting PPO vs Elastic Reset on IMDB but splitting the results visually between resets. We measure (a) Language Drift and (b) Task Performance via Semantic Score on the validation set over finetuning. All methods also include KI. penalty. We plot mean and standard error across 5 seeds.
Table 2: IMDB mock sentiment final test scores
SENTIMENT
PERPLEXITY
ZERO-SHOT
.489:0.01
32.45=0.13
PPO
.596=0.02
33.45#0.40
NLPO
558:10006
33.1220.74
ELASTIC RESET
.611+0.02
33.32+0.22
does Elastic Reset outperform the baselines at the final validation score, but functionally dominates such that it is the best method for all levels of task performance it achieves.
IMDB Mock Sentiment: Ablation Study for RLHE
Setup Next, we scale to a larger benchmark that more closely approximates the standard RLHF setup. We use the recently released GRUE benchmark for RL training of LMs [Ramamurthy et al., 2022] and use IMDB mock sentiment [Ziegler et al., 2019], the main task where language models susceptible to reward-hacking shown in Figure 2. The goal is to complete an IMDB movie review with as positive sentiment possible. The baseline LM GPT-2 [Radford al., 2019] with 117M parameters further pretrained on the IMDB domain [Maas et al., 2011]. We learn a DistilBERT [Sanh al., 2020] reward model on IMDB to output a sentiment score between (negative) and (positive). We then train our GPT-2 LM to complete different IMDB reviews while maximizing the sentiment reward. Following Ramamurthy al. [2022], we measure reward-hacking drift with our model's perplexity on the true IMDB data. If we consider knowledge of the IMDB data as useful capability, then our initial model was finetuned on IMDB to maximize log-probas prg--robability. i.c. minimize perplexity, and has the maximum capabilites. We measure divergence from the initial model, and decrease in capabilities, by the increase in our trained model's l's perplexity on ground truth IMDB data. In contrast to the previous task, lower perplexity
corresponds to less drift.
Baselines Our main baseline is PPO [Schulman et al., 2017] with Ziegler al. [2019] modifications for RLHF training, specifically adding a KI, penalty with the frozen initial model (equivalent to KI. WITH PRETRAINED) and dynamically decaying the coefficient over training. To further increase stability, Generalized Advantage Estimation [Schulman et al., 2015] is used for the advantage estimator. We also compare to NLPO [Ramamurthy al., 2022], a recent method that extends PPO with masking model to counteract drift. The masking model is initialized to the pretrained model and recieves delayed updates; is set the online model every n steps. During training. the online model's 's output probabilities are restricted to the mask model's top p tokens. We use the RL4LMs library Ramamurthy et al. [2022] and their default hyperparameters for both PPO and NLPO e.g. 0.1. We implement ELASTIC RESET on top of PPO with an EMA decay rate of 0.995 and greatly reduce the KI. coefficient 3 0.001 to allow the model drift more, then reset every 17 epochs such that get two resets / three iterations during training.

---

122
Bastic
122
PPO
12.
32.6
Perpicalty (drit)
Perglesity (drit)
12.0
33.0
0:30
33.6
PPD
33.8
0.50
0.52
0.54
0.55
0.58
0.69
0.2
0.64
0.50
0.52
0.54
0.56
0.58
0.60
Roward
Roward
(a) Reset Ablation
(b) EMA Decay Ablation
32.2
Easti Reset
12
32.4
(arri)
PPO
Gerplel (drit)
32.6
0:01
32.8
Elastic Reset
0.1
33.0
0.5
33.2
0.0001
0.001
PPO
0.50
0.55
0.60
0.65
0.70
0.75
0.50
0.52
0.34
0.56
0.58
0.60
0.82
Reward
Reward
(c) KL Ablation for PPO
KI. Ablation for Elastic Reset
Figure 5: Ablating Elastic Reset the IMDB mock sentiment task. We plot pareto graphs using mean and standard error across 5 seeds.
Experiments We run all experiments for 5 seeds and report mean and standard error on our validation set for reward, DistilBERT sentiment, and our drift score, perplexity. Following Ramamurthy et al. [2022], we run for 50 epochs (equivalent to 64k updates) and show our results in Figure 4. To make resets more visible, we plot validation scores every epoch for Elastic Reset. Since the benchmark provides a test set as well, we compare all final models in Table 2. The PPO baseline performs quite well because it already includes KL with the pretrained model. We find NLPO performs similarly to PPO, so we relegate NLPO results to Appendix B. Results from the original NLPO paper were stronger [Ramamurthy et al., 2022] but our reproduced numbers and curves were confirmed by the original authors [Ammanabrolu, 2023]. Elastic Reset achieves better semantic scores much faster by using smaller KL penalty coefficient (0.001 vs PPO 0.1) but also drifts more to achieve them. As with the previous task, this drift is then mitigated by the reset and we see semantic task score improve in relation drift over the iterations. Looking at the pareto graph in Figure 8c, we see that Elastic Reset far outpaces the baselines and provides better tradeoff of reward vs drift for every reward.
Ablations We empirically investigate our method through ablations. Throughout this section we experiments IMDB with the same hyperparameters, unless otherwise mentioned. For brevity, we plot only the pareto graphs but include all other graphs in Appendix D.3 along with these same ablation experiments for the Translation Game, with similar results.
To investigate the source of improvement in Elastic Reset, we ablate the two resets: online to EMA, and EMA to initial model. We discard the second reset to get Reset to EMA: model is reset to an EMA but the EMA is never reset. We also compare to the simplest reset idea, Reset to Init, and reset our policy to the initial model. We run all methods as previously and plot the pareto graph in Figure 5a, for task and drift graphs see Appendix D.3. We find that even simple resets are already performant but the two ablations have tradeoff: Reset maintains performance whereas Reset to Init does better at higher reward because it doesn't t drift as much. Elastic Reset combines the benefits of both and outperforms each method.
EMA is better at lower reward because
Next, we consider our method's robustness to hyperparameters. First, we search along different EMA decay rates 7) and plot our results in Figure 5b finding that our method is quite robust to choice of
7

---

4.0
Elsticc Reset
3.5
20
from Initiaal teiniht
KI from Initial
55
3.0
-0
1.5
600
Epochs
(a) Reward
(b) Drift
(c) Pareto Frontier Graph
Figure 6: Elastic Reset compared to PPO on StackLLaMA: A LLaMA-7B model RLHF finetuned on StackExchange as helpful, technical QA chatbot
decay. Next, we investigate robustness to the choice of KL penalty coefficient 3. We search across coefficients that range from 10x smaller 10x larger than our best KL penalty coefficient for PPO (3 = 0.1) and Elastic Reset (3 = 0.001). For visual clarity, we plot PPO in Figure 5e and Elastic Reset in Figure 5d and only plot four points for PPO 3 0, 0.01 to maintain visual scale. We find that PPO is not robust to choice of KL and larger values correspond to better pareto curves but slower training. Results with NLPO are similar and shown in Appendix D.3. contrast, Elastic Reset seems to be more robust to choice of KL with 0.001 producing the best curves while 10x larger and smaller values are similar. As opposed PPO, Elastic Reset even works reasonably well without KL penalty at all, (3 = 0), matching PPO's best performance with a KL. This demonstrates that the expensive KL penalty may be replaced with the cheap Elastic Reset, although the combination of the two is best. This is also in line with previous work that have argued that the KI. penalty may be unnecessary for RLHF [Bai et al., 2022, Gao al., 2022]. We also ablate the frequency of resets in Appendix D.4.1 and find that pareto curves are essentially unchanged.
Finally, we provide an empirical intuition for Elastic Reset: in Appendix E. show that resets iteratively improve the value function and in Appendix E.2 we show how EMA smoothes optimization but requires resetting order achieve high performance.
7 StackLLaMA: Practical RLHF
Setup Finally, apply Elastic Reset to larger-scale RLHF pipeline. We choose LLaMMA [TTuurron al., 2023] as it is a prominent open-source model that has demonstrated strong performance on benchmarks. We follow Beeching et al. [2023] to finetune LLaMA-7B with RLHF on the StackExchange dataset [Lambert et al., 2023] to output helpful answers to technical questions, as judged by humans. Users ask technical questions on StackExchange and upvote the best answers. We score answers from StackExchange based on the number of upvotes they received from users, score lo + upvotes) [Askell al., 2021] At most 10 answers text is cleaned, and HTML is converted to Markdown to make easier to parse. First, we finetune LLaMA-7B with language modelling the dataset get ALaMA-7B-SE We then further finetune it to get a reward model by learning to predict which two answers was more upvoted [Stiennon et al., 2020]. For given question and two answers V+,- (where y+ is preferred), the loss for our reward model rg is log or (ro(x, y+) ro y-)). Finally, we finetune LLLMA-7B-SE with RL against the reward model by sampling questions from the dataset and learning optimize the reward for model's s answer. All finetuning is done with a converted 8-bit model [Dettmers et al., 2022] and LoRA [Hu et al., 2021] for efficiency and to make the model training fit our GPUs. We rely on the HuggingFace Face trl [von Werra et al., 2023] and peft [Mangrulka et al., 2023] libraries. All technical details are described in Appendix C.
drawn per question,
Experiment We again compare to PPO with Ziegler al. [2019] modifications i.e. KL penalty with dynamically decaying coefficient. We run for 600 epochs (equivalent to 100k updates) and Elastic Reset every 260 epochs to get two resets three iterations. Each run takes 20 hours on 4 100s. Since only the LoRA parameters are being learned, we use Elastic Reset on those and therefore maintain only small percentage of parameters our EMA. We use a decay rate r = 0.995 and KI. penalty coefficient& 3 0.02 for both methods. Calculating perplexity for each epoch is computationally

---

Table Evaluations of the initial (zero-shot) and finetuned StackLLaMA models after 600 epochs. We measure alignment using an average over three reward model trained with three different seeds and drift with perplexity on the data. HumanEval is a programming benchmark that acts as a practical measure of drift alignment tax.
+ A REWARD + PERPLEXITY
1 HUMANEVAL (PASS @ 1, PASS @ 10)
ZERO-SHOT
4.43
11.0, 12.7
PPO
0.81 + 0.06
4.62
7.8, 10.7
ELASTIC RESET
0.96 + 0.09
4.57
11.0, 13.0
infeasible so we measure drift during training with the KL from the pretrained model over samples as done previously in other larger-scale RLHF [Bai et al., 2022, Gao et al., 2022].
Results We plot reward in Figure 6a 2 and KL from initial model over training in Figure 6b. As noted by Beeching al. [2023], the task much noisier at a larger scale and with real HF reward model. As a sanity check, we find that Elastic Reset tracks PPO until the first reset at 260 epochs where it drops only slightly, but also doesn't 't lose much performance. Around the second reset at 520 epochs, we see a much sharper drop but also maintaining the same approximate reward. At the end, Elastic Reset provides a non-trivial reduction in drift while aligning just as well as PPO. The pareto curve in Figure 6C shows Elastic Reset is equal or slightly worse at low reward but shows large improvements over PPO higher reward. Notably, Elastic Reset seems work nut-of-the-box with LoRA. To evaluate drift another way, we get the perplexity of the final models over the Stacknxe validation set as in Section 6. For more robust view of reward, we train two more reward models using different seeds and evaluate the increase in reward between initial and final models. We show mean and standard deviation across the three reward models in Table 3, we find that Elastic Reset achieves slightly better final reward than PPO while maintaining lower perplexity on the data. To examine a true alignment tax, we run our models on HumanEval Chen et al. [2021], programming benchmark that provides another view of drift. Answering human written coding questions is both a useful capability for our model and also falls within a similar domain to StackExchange. The benchmark tests for functional correctness such that pass @ corresponds to the percentage of problems solved by the model on the first try as shown in Table 3. Training with PPO degrades performance compared to the initial model, demonstrating large alignment tax. In contrast, Elastic Reset achieves a similar reward but maintains performance, even slightly improving on pass @ 10, creating alignment bonus [Askell et al., 2021] instead of tax.
human-written
8 Limitations
As a method, Elastic Reset is quite cheap computationally because both EMA updates and resets take negligable time compared to RLHF training and the EMA model can be stored on CPU. But our method is sensitive the choice of reset rate; we chose heuristically based on when it seemed the model was overfitting. It is also possible to reset the policy and EMA model at different time scales, which could be source of improvement. Our method also resets all of the trainable parameters, research in similar methods suggests that resetting larger models can benefit from resetting only part of the network [Zhou et al., 2022, Nikishin et al., 2022] or weighted-averrging instead of resets [D'Oro et al., 2023]. We leave both of these directions to future work.
Although we have thoroughly investigated our method three different tasks, we note that none of them are ideal RLHF benchmarks. As pointed out by Gao et al. [2022]. we measure our model's performance using the same reward model we optimize. This can lead to reward model overoptimiza- tion and our metric could mask overfitting and lack of generalization to the real world i.c. actual human preferences. An ideal benchmark could include a "gold" reward model a proxy for human preference [Gao et al., 2022], but no such benchmarks are open-sourced and available.
Note that our reward results differ from the original StackLLaMA likely due to a pending issue in their code affecting their reward model. After fixing the issue, the only difference seems to be an increase in base reward, and our reward curves visually quite similar theirs. See Appendix C for details

---

Finally, we note that we follow all previous RLHF work and investigate only on-policy methods [Ziegler et al., 2019, Stiennon et al., 2020, Askell al., 2021, Bai al., 2022]. Previous work in resetting for RL has focused on off- methods and demonstrated strong performance [Nikishin et al., 2022, D' Ororo et al., 2023]. As previously noted, our method can be seen as an adaptation of those to on-policy RL. In RLHF, PPO is by far the most popular method and on-policy is the dominant paradigm since it guarantees better local gradients. But it is possible that off off-policy methods could implicitly balance performance and drift by incorporating replay buffer with older data.
Conclusion
The problems drift [Lee et al., 2019], alignment tax [Askell ct al., 2021], reward model overopti- mization [Gao al., 2022], and reward hacking [Clark and Amodei, 2016] are inherent to RLHF and reduce its efficacy. We have introduced a simple but powerful new method, Elastic Reset, to tackle this problem and improve performance while maintaining linguistic capabilities. We have shown its ability on three different tasks and across three different scales: from layer Transformers to GPT2 to LLaMA-7B. The problem of drift is currently being addressed with standard KL penalty despite the computational cost, tradeoff with reward, and recent claims that is may be unnecessary [Bai et al., 2022, Gao et al., 2022]. Elastic Reset cheap and effective method to tackle the same problem, achieving a better tradeoff of reward and drift while reducing alignment tax. We hope our method leads to better RLHF and therefore models that closer aligned with human preferences [Ziegler et al., 2019]. As well, we hope this work invigorates more research into improving the reward drift tradcoff of RLHF with a focus on computationally efficient methods that scale.
Acknowledgments and Disclosure of Funding
MN is supported by Fonds de recherche du Québec - Nature et technologies and Sony. MN would like thank Issam Laradji, ServiceNow Research, Mila, and Compute Canada for providing resources used in the experiments.
References
P. Ammanabrolu. Re: Reproducing NLPO (RL4LMs), Jan. 2023.
J. Andreas, D. Klein, and S. Levine. Learning with Latent Language. In NAACL. arXiv, 2018. doi: 10448550/axxi..111..000 URL http://arxiiv http://ww.wggab//////1/// 00482. arXiv:1711.0045: [cs].
A. Askell, Y. Bai, A. Chen, D. Drain, D. Ganguli, T. Henighan, A. Jones, N. Joseph, B. Mann, N. DasSarma, N. Elhage, Z. HatfieldDDodss, Hernandez, J. Kernion, K. Ndousse, Olsson, D. Amodei, T. Brown, J. Clark, S. McCandlish, C. Olah, and J. Kaplan. A General Language Assistant as Laboratory for Alignment, Dec. 2021. URL attp://arxivv orgg 00861. 1xXix:2112.0088 [cs].
D. Bahdanau, K. Cho, and Y. Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. In ICLR. arXiv, 2015. URL http:://axtii wwrg/abs//en /409/ 0473. arXiv:1409 1409.0473 [cs, stat].
Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, N. Joseph. S. Kadavath, J. Kernion, T. Conerly, S. El-Showk, N. Elhage, Z. Hatfield- Dodds, D. Hernandez. T. Hume, S. Johnston, S. Kravec, Lovitt, N. Nanda, C. Olsson, D. Amodei. T. Brown, J. Clark, S. McCandlish, C. Olah, B. Mann, and J. Kaplan. Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback, Apr. 2022. URL attp://arxiv.. org/ abs/2204 05862. arXiv:220.00588
[cs].
E. Beeching. Y. Belkada, K. Rasul, I.. Tunstall, L. v. Werra, N. Rajani, and N. Lambert. StackLLaMA: An RI. Fine-tuned I.LMMA Model for Stack Exchange Question and Answering, 2023. URL https://hugging http:///wa....o/////////ccccccca
10

---

E. M. Bender and A. Koller. Climbing towards NLU: On Meaning, Form, and Understanding in the Age of Data. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 5185-5198, Online, July 2020. Association for Computational Linguisties. doi: 10. 10.18653/v12022..am acl m63n.463. URL htwp:::tp////ttt htattp:://acaantholoogyy/ www.oog......//oe org/2020/ acl-main 463.
T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Horsert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, McCandlish, A. Radford, Sutskever, and D. Amodei. Language Models are Few-Shot Learners. In Neural Information Processing Systems. arXiv, July 2020, doi: 10.48550//rX1v.2005. 14165, URL http:// http///axivv...////v///.//..aavt org/ htsp//bb/000550 14165. arXiv:2005.14165 [es].
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging Properties in Self-Supervised Vision Transformers, May 2021. URL http: www...rg/./..... 2104. 14294. utXiv:2104.14220 14294 [cs].
M. Cettolo, C. Girardi, and M. Federico. WIT3: Web Inventory of Transcribed and Translated Talks. In Proceedings of the 16th Annual conference of the European Association for Machine Translation, pages 261-268. Trento, Italy, May 2012. European Association for Machine Translation. URL https: /atp//aaaannooogyyyyyooooo
2012.222...00 eamt 1. 60.
M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. d. O. Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, A. Ray, R. Puri, G. Krueger, M. Petrov, H. Khlaaf, G. Sastry, P. Mishkin, B. Chan, S. Gray, N. Ryder, M. Pavlov, A. Power, L. Kaiser, M. Bavarian, C. Winter, Tillet, F. P. Such, D. Cummings, M. Plappert, F. Chantzis, E. Barnes, A. Herbert-Voss, W. H. Guss, A. Nichol, A. Paino, N. Tezak, J. Tang, Babuschkin, S. Balaji, S. Jain, W. Saunders, C. Hesse, A. N. Carr, J. Leikee, J. Achiam, V. Misra, E. Morikawa, A. Radford, M. Knight, M. Brundage, M. Murati, K. Mayer, P. Welinder, B. McGrew, D. Amodei, S. McCandlish, I. Sutskever, and W. Zaremba. Evaluating Large Language Models Trained on Code. July 2021. URL http://arxiv, org/ //bss/ 2107. 03374. vrXiv:2107) 200.0337733 [cs].
F. Christiano, J. Leike, T. Brown, M. Martic, Legg, and D. Amodei. Deep Reinforcement Learning from Human Preferences. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL petps://papeet peapers.nip..oo cctp:ccccpppp/00///// 2017/ htta://62200aadc00001999122200e0 Aabssetreeeee-bbaabrttete.e. html.
Clark and D. Amodei. Faulty Reward Functions in the Wild, Dec. 2016. htt.httpss//opeaai. honm///lo//aaally-eearrddd--n functions/.
T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer: LLM.int80: 8-bit Matrix Multiplication for Transformers at Scale. In NeurlPS. arXiv, Nov. 2022. doi: 10448500arxiv22200.77333 URL http://arxiv.orr org/.bs.g/...... abtpaabs.../..///////////// 07339. arXiv:2208.0777 00339 [es].
D'Oro, M. Schwarzer, E. Nikishin, P.-L. Bacon, M. G. Bellemare, and A. Courville. Sample- Efficient Reinforcement Learning by Breaking the Replay Ratio Barrier. In International Con- ference on Learning Representations, 2023. URL 4GBGWVIEYJ.
new.eeneeeeeeeeeeeee net/forum?id
D. Elliott, S. Frank, K. Sima' an, and L. Specia. Multi30K: Multilingual English- German Image Descriptions. In Proceedings of the 5th Workshop on Vision and Language, pages 70-74, Berlin, Germany, 2016. Association for Computational Linguistics. doi: 10. 18653/v1/W 16-3210 URL http://aclwas http:///wwwwb..e..... wwtaaaatoooooooo3/100
English-German
D. Elliott, S. Frank, L. Barrault, F. Bougares, and L. Specia. Findings of the Second Shared Task on Multimodal Machine Translation and Multilingual Image Description. In Proceedings of the Second Conference Machine Translation, pages 215-233, Copenhagen, Denmark, Sept. 2017. Association for Computational Linguistics. doi: 10. 10.18653/1//11774718.. URI. https: //aclanthology prg/W17-444
I. Gao, J. Schulman, and J. Hilton. Scaling Laws for Reward Model Overoptimization, Oct. 2022.
URL httb::///axiv.... org/ awtpaabs/...////////////aaaaaaaa 10760. arXiv:2210.1011 10760 [cs, stat].
11

---

S. Gugger, L. Debut, T. Wolf, P. Schmid, Z. Mueller, S. Mangrulkar, M. Sun, and B. Bossan. Accelerate: Training and inference scale made simple. efficient and adaptable, 2022. URL https://g" /tt:::///hhtii com/ conn//uugginggaccc//cceeee..ee
Hochreiter and Schmidhuber. Long Short-Temn Memory. Neural Computation, 98))11335- 1997.9.8.17335. URL https://doi.org///
Nov. 1997. ISSN 0899-7667. doi: 1162/neco. 1997.
1735.
E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. LoRA: Low-Rank Adaptation of Large Language Models, Oct. 2021. URL.htpp:/arxiv perg/abs/2100 09685. ar arXiv:2106.09635 [cs].
J. D. Hunter. Matplotli:: A 2D Graphics Enviromment. Computing in Science & Engineering, 9()): 90-95, May 2007. ISSN 1558-366X. doi: 10.1 109/CCF URL
beee.rrg/ocuunee 402655. Conference Name: Computing in Science & Engineering.
E. Jang, Gu, and B. Poole. Categorical Reparamcterization with Gumbel-Softmax, ICLR. arXiv, Aug. 2017. doi: 10.48550/arXiv 110S500/rXiv.1611..11 1611,01144, URL http:// owt/oog///////ooooooo
1rXiv:1611.0111 [cs, stat].
N. Jaques, S. Gu, D. Bahdanau, J. M. Hernández-Lobace R. E. Turner, and D. Eck. Sequence Tutor: Conservative Fine-TTuning Sequence Generation Models with KL-control, International Conference Machine Learning. arXiv, Oct. 2017. doi: 10448550/arXi..1611122777 [cs].
URL
http://arxiv/ wwwpsiiivorggaeeee wtgg/ba/////.l.oo/ 12t../2766.....6...
N. Jaques, A. Ghandeharioun, J. H. Shen, C. Ferguson, A. Lapedriza, N. Jones. S. Gu, and R. Picard. Way Off-Policy Batch Deep Reinforcement Learning of Implicit Human Preferences in Dialog, July 2019. URL http: // wwwp://..rraaab/////aaaae 10456 00456. 19XXvv197.00046 [cs, stat].
Kirby. Spontaneous evolution linguistic structure-an iterated learning model of the emergence of regularity and irregularity. IEEE Transactions on Evolutionary Computation, 5(2):102-110, Apr. 2001. ISSN 1941-002 -0026. doi: 10.1109/4235.918440 Conference Name: IEEE Transactions on Evolutionary Computation.
Kirkpatrick, R. Pascanu, N. Rabinowitz, J. Veness, G. Desjardins, A. A. Rusu, K. Milan, J. Quan, T. Ramalho, A. Grabska-Barwinskk. D. Hassabis, C. Clopath, D. Kumaran, and R. Hadsell. Overcoming catastrophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 114(13):3521-233 -3526, Mar. 2017. doi: 10. 0077/ppas.1611888 14. URL https: // /www./g www..paas... org//oi/11 1073 /ptp//aas.... 1611835114. Academy of Sciences.
Publisher: Proceedings of the National
Koehn, H. Hoang, A. Birch, C. Callison-Burch, M. Federico, N. Bertoldi, B. Cowan, W. Shen, C. Moran, R. Zens, C. Dyer, O. Bojar, A. Constantin, and E. Herbst. Moses: Open Source Toolkit for Statistical Machine Translation. In Proceedings of the 45th Annual Meeting of the Association for Computational Linguistics Companion Volume Proceedings of the Demo and Poster Sessions, pages 177-180. Prague. Czech Republic, June 2007. Association for Computational Linguistics. URL https: /ete:///aalaateee /aaaaaaaaannoooooceee wwaggg..rgggeee org/P07-2040
N. Lambert, L. Tunstall, N. Rajani, and T. Thrush. HuggingFace H4 Stack Exchange Preference Dataset, 2023. URL. https: https://hugginggfce. co/ http//atsst///uggiggaaaaaaa stack Srack-excia exxchang--preferee
A. Lazaridou and M. Baroni. Emergent Multi-Agent Communication in the Deep Learning Era, July 2020.
A. Lazaridou, A. Potapenko, and O. Tieleman. Multi-agent Communication meets Natural Language: Synergies between Functional and Structural Language Learning. In ACL. arXiv, May 2020. doi: 20048850/a2xxi.2000.77766 URL http://atre otg/ 07064. arXiv:2005.0700 00000 [cs].
Lee, Cho, and D. Kiela. Countering Language Drift via Visual Grounding, Sept. 2019. URL fttp://arxing wwwpxvivv.rr//.oo org/abe/// awa/aa////9/ 04499. arXiv:1909.04493 [cs].
12

---

M. Lewis, D. Yarats, Y. N. Dauphin, D. Parikh, and D. Batra. Deal or No Deal? End-to-Eme Learning for Negotiation Dialogues, June 2017. URL http: arXiv:100..05125 [es] version: 1.
ha/aab///10/ 05125.
Q. Lhoest, A. V. del Moral, Y. Jernite, A. Thakur, P. von Platen, S. Patil, Chaumond, M. Drame, Plu, L. Tunstall, J. Davison, M. Saško, G. Chhablani, B. Malik, S. Brandeis, T. L. Scao, V. Sanh, Xu, N. Patry, A. McMillan-Major P. Schmid, S. Gugger, C. Delangue, T. Matussière, L. Debut, S. Bekman, P. Cistac, T. Gochringer. V. Mustar, F. Lagunas, A. M. Rush, and T. Wolf. Datasets: Community Library for Natural Language Processing, Sept. 2021. URL http: //arxiv. org/ abs/abbb..../// 02846. arXiv:2109.0288. [cs].
F. Li and M. Bowling. Ease-of-Teaching and Language Structure from Emergent Communication. In NeurIPS, 2019. URL http:htttaaiiiiii owareiiii.orgaaae wwwggabs/19900 02403. arXiv: 1906.02403.
I. Loshchilov and F. Hutter. Decoupled Weight Decay Regularization. In International Confer- ence on Learning Representations, Feb. 2022. URL http::///pppreeeee.eeee/tooooae Bkg6RicqY7.
R. Lowe*, A. Gupta*, J. Foerster, D. Kiela, and J. Pineau. On interaction between supervision and self-play in emergent communication. In International Conference Learning Representations, Sept. 2021. URL https:/
Y. Lu, S. Singhal, F. Strub, O. Pietquin, and A. Courville. Countering Language Drift with Seeded Iterated Learning. In ICML. arXiv, Aug. 2020. doi: 10,485500arXxv.20003. 12694. URL http: //arxiv. org/ absaaba3...../////////////////// 12694. arXiv:2003. 12694 [cs].
A. L. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng. and C. Potts. Learning Word Vectors for Sentiment Analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142-150, Portland, Oregon, USA, June 2011, Association for Computational Linguisties. URL https: htteps//acllantholog" org/P11-1015 1015.
C. J. Maddison, A. Mnih, and Y. W. Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. In ICLR. arXiv, Mar. 2017. doi: 00448500/rXiv.1611..00 1.07000 URL [cs, stat].
http://arxivv http://axxiv..rgabbb//eee// arggabs/1611. 1tt.///..11111.......111111 00712. 1rXiv:1611.000
S. Mangrulka, S. Gugger, L. Debut, Y. Belkada, and S. Paul, PEFT: State-of-the-art Parameter- Efficient Tunein methods. May 2023. URL https: gitgggt..bb....gg cwt. wonm/huggingaace/peeeeee teett original-date: 2022-11-22 1-25T03:51:092.
M. McCloskey and N. J. Cohen. Catastrophic Interference in Connectionist Networks: The Sequential Bower, editor, Psychology of Learning and Motivation. volume 24, pages 109-165. Academic Press, Jan. 1989. doi: 10. 10.0006650007997722200 //www. www...ccieneedirec... http://ciiinneaaarrr////////////
Learning Problem. In
URL https:
Meta FAIR Diplomacy Team, A. Bakhtin, N. Brown. E. Dinan, G. Farina, C. Flaherty, D. Fried, A. Goff, J. Gray, H. Hu, A. P. Jacob, M. Komeili, K. Konath, M. Kwon, A. Lerer, M. Lewis, A. H. Miller, S. Mitts, A. Renduchintala, S. Roller, D. Rowe, W. Shi, J. Spisak, A. Wei, D. Wu, H. Zhang, and M. Zijlstra. Human-level play in the game Diplomacy by combining language models with strategic reasoning. Science, Nov. 2022. doi: 11116//ciennce.adeee. //www www..science.eomm org/doi/ 1126/science ade9097. Publisher: American Association for the Advancement of Science.
URL https:
E. Nikishin, M. Schwarzer, P. D'Oro, P.-.. Bacon, and A. Courville. The Primacy Bias in Deep Reinforcement Learning. In Proceedings of the 39th International Conference on Machine Learning. pages 16828- 16847. PMLR, June 2022. URL https:// mlr, press/ v162/nikishin22a html. ISSN: 2640-3498.
OpenAI. ChatGPT: Optimizing Language Models for Dialogue, Nov. 2022. URL //
ccmp://aaceoeeaaacce.......... :https:
13

---

M. Ott, S. Edunov. A. Baevski, A. Fan, S. Gross, N. Ng. D. Grangier, and M. Auli. fairseq: A Fast, Extensible Toolkit for Sequence Modeling. In Proceedings of the 2019 Conference of the North, pages 48-53, Minneapolis, Minnesota, 2019. Association for Computational Linguisties. doi: 10.18655/11/N11-40099 URL http://aclweb. wwwpa//wwwww..ee.eeeee.ee.eoooee wwtgg/anthollogg//M--000000
L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, J. Schulman, J. Hilton, F. Kelton, L. Miller, M. Simens, A. Askell, Welinder, P. Christiano, Leike, and R. Lowe. Training language models follow instructions with human feedback, Mar. 2022. URL http: /atp:v/aavvv..//a org/ abs/2203 02155. arXiv:2203.02155 [cs].
A. Paszke, S. Gross, Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang. Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, Bai, and S. Chintala. PyTorch: An Imperative Style, Perforreffomance Deep Learning Library, Dec. 2019. URL atttp://arxinn wwwaaii...g/abb..... 1912. 01703. arXiv:1912.01703 [cs, stat].
M. Post. A Call for Clarity in Reporting BLEU Scores. In Proceedings of the Third Conference Machine Translation: Research Papers, pages 191, Brussels, Belgium, Oct. 2018. Association for Computational Linguisties. doi: 3xg/W18-6319
18-6319. URL https: / httppts//ccllaathhlllge
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language Models Unsupervised Multitask Learners, 2019. URL https://wwit http:://www.emmattccccho.. semantccscholar org/ htppr//anguage-MMdll---reeeeee Unsuprrru wew.eutttadddll -ellitttsk-Leareess Learners- Radford-wu/ #40C:00881096677177776766666666
R. Ramamurthy, Ammanabrolu, K. Brantley, J. Hessel, R. Sifa, Bauckhage, H. Hajishirzi, and Y. Choi. Is Reinforcement Learning (Not) for Natural Language Processing: Benchmarks, Baselines, and Building Blocks for Natural Language Policy Optimization. ICLR. arXiv, 2022. abs/ 2200022111.2 01241. arXiv:2210.012 [cs].
URL http: / //w.xv..vv/// orggabs/////oooaooooogg.og...srr
M. Rita, C. Tallec, P. Michel, J.-B. Grill, O. Pietquin, E. Dupoux, and F. Strub. Emergent Communication: Generalization and Overfitting in Lewis Games, Sept. 2022. URL http: //arxivvorrg
http://oor/abb///////// abs/2209....../// 15342. arXiv:2209. 15342 [cs, math].
A. Robins. Catastrophic Forgetting, Rehearsal and Pseudorehearsal. Connection Science, (2):123-146, June 1995. ISSN 0954-0091. https://doi. org/10. 100000005000995000333330 http:://oi..or//0.110000 wtt..///000.10000000000555001111
doi: 10. 10.180009540995503331
URL Publisher: Taylor Francis _eprint:
D. Rolnick, A. Ahuja, J. Schwarz, T, P. Lillicrap, and G. Wayne. Experience Replay for Continual Learning, Nov. 2019. URL http:///axxiaa org/ absaaas//i11 18111 11682. 1rXiv:1811.1163 [cs, stat].
V. Sanh, L. Debut, J. Chaumond, and T. Wolf. DistilBERT,, distilled version of BERT: smaller, faster, cheaper and lighter, Feb. 2020. URL http:://arziv.org/abbas. 0110. [cs].
01108. 111.v:1910001108
Schulman, Levine, Moritz, M. Jordan, and Abbeel. Trust Region Policy Optimization. In ICML. arXiv, 2015. doi: 00448550aaxx..15022555/77 URL htttp://arxiaa org/ www.xii.org/abs/55000 abs/1 05477. arXiv: 502.05577 [cs].
J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal Policy Optimization Algorithms, Aug. 2017. URL httt:://arxivv org/ htgg//bs/1707/// 06347. MrXiv:1707.003477 [cs].
A. K. Singh, D. Ding, A. Saxe, F. Hill, and A. K. Lampinen. Know your audience: specializing grounded language models with the game of Dixit, June 2022. URL http:://axxvvvxr///aae/ 2206. 08349. arXiv:2206.0832 08349 [cs].
S. Steinert-Threlkell X. Zhou, Z. Liu, and C. M. Downey. Emergent Communication Fine-tuning (EC-FT) for Pretrained Language Models. Emergent Communication Workshop at ICLR 2022, June 2022. URL https: BtRthtt::/oopereeeeeeee opwrrevieeeeeeeeeeeeeee wtw.eewwweeeeeeeeeee mett:/oorumiid--SSrrrrttt
14

---

N. Stiennon, L. Ouyang, J. Wu, D. M. Ziegler. R. Lowe, C. Voss, A. Radford, D. Amodei, and P. Christiano. Learning summarize from human feedback. In NeurIPS. arXiv. 2020. URL http:://axxiv.rr/bbbss.sss 000.0.0................. 01325. arXiv:2009.011 [es].
1. Sutskever, O. Vinyals, and Q V. Le. Sequence to Sequence Learning with Neural Networks. In Neural Information Processing Systems. arXiv, Dec. 2014. URL bttp:://arxiii www.iiivoo///oooo arg/abs/1409 3215. #tXiv:1409.32155 [es].
H. Q. To, N. D. Q. Bui, and M. Nguyen. CodeCapybara: Open Source LLaMA Model that Follow Instruction-Tuning for Code Generation., May 2023. URL https://githut. com/ #Soft-AI4Code SI14Code/Codeeee /Codeaapyaaaaa.a...aaaaaaa original-date: 2023-04-2IT10:22:532
H. Touvron, Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. LLaMA: Open and Efficient Foundation Language Models, Feb. 2023. URL wttthhtt///rrtt org hwtg/ab//330...11 13971. arXiv:2302.1399 13977
[cs] version: 1.
M. Utiyama and H. Isahara. A Comparison of Pivot Methods for Phrase-Based Statistical Machine Translation. In Human Language Technologies 2007: The Conference of the North American Chapter of the Association for Computational Linguisties; Proceedings of the Main Conference, pages 491, Rochester, New York, Apr. 2007. Association for Computational Linguistics. URL https: /ttpps://aclls /awpatholaanotoyyo.yooooo org/ N061
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is All you Need. In Neural Information Processing Systems, page 11, 2017.
L. von Werra, Y. Belkada, L. Tunstall, E. Beeching, T. Thrush, and N. Lambert. TRL: Transformer Reinforcement Learning, 2023. URL https: pepttts:://ittueee git.ub. com/Lvwerra/trl
R. Williams. Simple statistical gradient-following algorithms connectionist reinforcement learning. Machine Learning, page 28, 1992.
T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Jernite, J. Plu, C. Xu, T. Le Scao, oftte-of-the-Ar Natural Lan-
S. Gugger, M. Drame, Q. Lhoest, and A. Rush. Transformers: guage Processing. In Proceedings of the 2020 Conference on Empirical Methods in Natu- ral Language Processing: System Demonstrations, pages 38-45, Online, Oct. 2020. Associ- ation for Computational Linguistics. doi: 10. (0055//1122220.emmppeeemee // //aclanthology org/ wwwgg//220..2/ emnlp-demos 6.
URL https:
H. Zhou, A. Vani, H. Larochelle, and A. Courville. Fortuitous Forgetting in Connectionist Net- works. In International Conference on Learning Representations, Mar. 2022. URL https: //openreview ftsts/forumidddei3311
D. M. Ziegler, N. Stiennon, J. Wu, T. Brown, A. Radford, D. Amodei, Christiano, and G. Irving. Fine- Fine-Tuning Tuning Language Models from Human Preferences, 2019. URL http://arxiv.... org. /bs// 1909. 08593. arXiv: 1909. 1909.0899 [es, stat].
15

---

as
Promssendo
Facce
Poter sude
1046
20000
2005
4009
S000
20000
30000
C000
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 7: All transformer model methods on the Translation Game We measure (a) Task Performance with FR ER +DE BLEU and (b) Language Drift with FR BLEU, the validation set during finetuning, We plot the mean and show error bars for standard deviation over seeds. To compare how methods do on both metrics, plot (c) the best achieved drift task performance across finetuning.
A Translation Game
A.1 Experimental Details
We implement the translation game in the fairseq library [Ott et al., 2019] on top of Pytorch [Paszke al., 2019]. All experiments are run with 5 seeds where each run uses single 16G V100 GPU. All plots show the mean and standard deviation over seeds. We score BLEU using detokenized sacreBLEU [Post, 2018]. Plots are made with Matplotlib [Hunter, 2007]
A.2 Translation Game with Transformers
Note that Lee al. [2019] (as well as previous work the Translation Game) seq2seq [Sutskever et al., 2014] LSTMs [Hochreiter and Schmidhuber, 1997] with attention [Bahdanau et al., 2015]. For consistency with other experiments and relevance to mainstream research switch Transformer architecture. LSTMs results are similar and results are available in Appendix A.3. Also note that we REINFORCE
estimator but Gumbel-SSftmax [Jang et al., 2017, Maddison al., 2017] is also feasible [Lu et al., 2020] and preliminary results have been similar.
We follow Lee et al. [2019] for preprocessing both IWLST and Multi30k. We use the iwslt-de- architecture from the fairseq library [Ott et al., 2019] with their default IWSLT DE-EN hyperparam- eters and training, specifically training with Adam [Loshchilov and Hutter, 2022] using default hyperparameters. We pretrain our models on IWSLT following Lee et al. [2019] and early-stop on the tst2013 validation set. Our final validation scores 43.85 BLEU for FR +EN and 29.4 BLEU for EN-DE.
We report all final validation scores in Table 4 and all curves in Figure 7.
Given that have learning reward model (EN->DE) iteratively resetting FR EEN model to its pretrained weights can be an effective strategy [Rita et al., 2022]. We report this as RESET and find it performs quite well but not as well as Elastic Reset. We did include these results in the main paper because we focused on methods that could also work with fixed reward models.
Table 4: Translation Game final validation scores
FR ERN +DE
FR++NN
PRETRAINED
23.8
36.2
FROZEN SENDER
30.8+0.2
36.3=0.1
REINFORCE
33.2+0.3
29.6+0.3
+ SIL
28.2+0.4
27.3=4.4
+ MULTITASK (S2P)
32.2+0.33
35.2+1.0
+ PRETRAINED
33.2-1..2
30.8+0.4
+ RESET
32.5+10.1
33.3-0.1
+ RESET TO EMA
32.9+0.1
36.3-0.1
+ ELASTIC RESET
33.0+0.1
36.3-0.1
16

---

A.3 Translation Game with LSTMs
We follow Lee et al. [2019] for preprocessing and pretraining on IWSLT. We compare our pretraining results to both Lee et al. [2019] and Lu al. [2020] in Table
Table 5: BLEU score of IWSLT-pretrained LSTM models on IWSLT 2013 validation set
FR-+EN
EN-+DE
LEE ET AL. [2019]
34.1
22.0
LU ET AT.. [2020]
32.2
20.2
OURS
38.5
23.2
Next we finetune on Multi30k using the same hyperparameters Lu et al. [2020]. We plot results in Table 6 and compare published numbers from previous work. As usual, drift the negative change in FR EEE BLEU from the pretrained model. Task performance is the positive change in FR-EEN+DE BLEU from the pretrained models. Combined the sum of drift and task performance. We show the pretrained, baseline, and best-performing model from previous work. Note that our results are not directly comparable to previous work because we evaluate using detokenized sacreBLEU [Post, 2018] whereas previous work wrote their own BLEU evaluation code and did not detokenize.
Table 6: BLEU scores and + standard deviation on the Multi30k Translation Game using IWSLT- pretrained LSTM models.
METHOD
FR+EN
FR-EN-DE
DRIFT
PERF
COMBINED
LEE ET AL. [2019]
PRETRAINED
27.2
16.3
REINFORCE
1.74
1.5
-14.8
+8.2
-6.6
LM
23.6 1.1
27.7
-3.6
+11.4
+7.8
LM +G
0.4
28.1
-2.4
+11.8
+9.4
LU ET AT.. [2020]
PRETRAINED
29.4
15.7
GUMBEL-SOFIMAX
14.5 0.8
27.1 0.1
-14.9
+11.4
-3.5
SIL
0.3
0.2
+12.6
+12.6
OURS
PRETRAINED
32.6
18.0
BOZEN-SENDER
3
2.11+
+7.1
+7.1
REINFORCE
0.2
3.2
-2.6
+12.3
+9.7
LM 0.01
3.11
3.2
-1.1
+12.2
+11.1
MULTITASK À 0.1
32.9
0.2
2.3
+0.3
+10.5
+10.8
SIL
27.7
0.7
2.1
-4,9
6.3
+1,4
BLASTIC RESET
0.1
30.0
0
+12.0
+12.0
Our results for the baseline notably better than Lee et al. [2019], Lu et al. [2020]. The only difference between our code and theirs as far as we can tell is (1) we use an exponential moving average bascline for REINFORCE whereas [Lee al., 2019] use Actor-Crit
method, (2) Lu et al. [2020] uses 0.1 gradient clipping and we do not use gradient clipping (3) preprocessing Multi30k, first tokenize [Koehn et al., 2007] then lowercase whereas previous works did the opposite order.
more reasonable explanation for the improvement in results is how we choose evaluate. Previous work simply ran all methods for the same number of updates but this doesn't 't account for, even implicit, hyperparameter optimization. Previous methods show that baseline's FR TRR-ENDDD BLEU scores plateau and there is significant language drift in FR-EN without real improvements to the task score. We hypothesize that these extra training episodes only serve to increase the drift without measuring what we actually care about: performance gain for drift. Since the number of updates arbitrary, we believe that early stopping on a reasonable metric a better evaluation protocol and choose hyperparameters such that methods plateau at the end.
We also note that our results with the SIL method of Lu et al. [2020] are negative. We do not manage to gain any improvement in performance. We collaborated with the authors of Lu et al. [2020] for many months but, in our setup using the fairseq library [Ott et al., 2019] could not reproduce their results. One of their fundamental results is that student sender can outperform a teacher sender that it is distilling from. We could not reproduce this and believe it is difference the pretrained models.
17

---

single 40G A100 GPU. All plots show the B IMDB Mock Sentiment
Bostic Reset
0.62
PPO
32.4
0.60
13.2
NLPO
Perplesity (oritt
12.6
0.58
23.0
32.8
12.8
33.0
32.6
33.2
0:50
32.2
30
20
30
40
10
20
30
40
50
3,500
0,525
0.550
0.600 Reward
0.575
D/625
Epoch
Epoch
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure IMDB mock sentiment validation scores for Elastic Reset, PPO, and NLPO. We measure (a) Language Drift and (b) Task Performance via Semantic Score on the validation set over finetuning. All methods also include a KL penalty. We plot mean and standard error across 5 seeds.
B.1 Experimental Details
We run experiments and implement our method the RLALMs library [Ramamurthy et al., 2022]. All experiments
with 5 seeds where mean and standard error over seeds.
For PPO, use the default hyperparemters provided by Ramamurthy et al. [2022]. We also compared to NLPO, we found the defaults had mistake and after communication with Ramamurthy al. [2022] we changed the learning rate le-6 and target update iterations 50. We plot results in Figure 8. This still did not manage to reproduce the original NLPO test scores from Ramamurthy al. [2022] but found that validation curves matched their provided curves in the appendix. After communications [Ammanabrolu, 2023], we both agreed that we should use our reproducible test scores for NIPO in lieu of the original test scores.
Our best Elastic Reset hyperparameters are the default PPO parameters (gpt2_pppo. yml from RL4LMs [Ramamurthy et al., 2022] with few modifications: target kl is set from 0.5 to 1.0 and KL coefficient is set much lower to 0.001. We use an EMA decay 0.995 and reset ev- ery 17 epochs. Configs to reproduce the experiments can be found in the RL4LMs folder under tttp:/tsssssottttt configes iimeiiissmmmiiieeee httlimmdbbbbeee ttaibb_txt___ett conett_coninnaaii...
C
StackLlama
Elastic Reset
17.5
Elastic Reset
PPO
15.0
PPO
3.00
12.5
Reward
2.75
10.0
2.50
7.5
2.25
5.0
1/m
2.5
0.0
100
200
300 Epochs
400
500
100
200
300
500
600
Epochs
Reward
(b) KI. from Initial
Figure 9; Rolling Window Mean (window = 10) version of the StackLLaMA results
C.1
Experimental Details
We follow the original StackI.lama and use the trl library from Huggingface on top of the transformers library [Wolf et al., 2020] to train the model on the StackExchange dataset, loaded with datasets
18

---

[Lhoest et al., 2021]. We use the original authors' LoRA adapter weights to create our supervised model LLaMA-7B-S but train our own reward model as there were issues loading the pretrained reward model. As noted by [Beeching et al., 2023], the reward modelling task is difficult enough that humans struggle with and our final model achieves 64% accuracy compared the original 67%. Although we follow the original authors method and use their codehase, we note that our results may be different but valid. There have been many updates and fixes to the trl codebase since the authors' original blog post and specifically a possible issue in the code for creating reward models could have affected the original authors' run.
Furthermore, to speed up training, we used a 2x smaller KL coeffient and ran with half the number of GPUs. Specifically, RL training run on 4x 80G A100 GPÜs for 20 hours using Accelerate [Gugger et al., 2022]. We evaluate perplexity of the model on the validation set used in supervised finetuning. 4000 examples from the supervised training dataset. The configs for reproducing our experiments are in the trl library folder under examples wetees/steeee stackllama/scrr
conf igs
For visual clarity, we provide a version of StackLLaMA results with rolling window in Figure 9
C.2
HumanEval
Our measure of alignment tax, HumanEval [Chen et al., 3021], is a programming benchmark where each question was hand-written by humans (OpenAl engineers) to be unseen training data. We believe these questions were unseen LLaMA training data as well. We prompt the model with the question and it writes the corresponding code. The measure of success is functionally correct code ie. we actually execute the code LLaMA wrote and see it gets the right output. Since we decode by sampling, generate N 100 continuations of the prompt and then follow Chen et al. [2021] to estimate how often our model would get the right answer on the first generated continuation (pass @ and within the first 10 generated continuations (pass @ 10). We evaluate HumanEval using CodeCapybara's evaluation harness [To et al., 2023].
D Ablations Full Graphs
D.1 Reset Number
0.42
0.52
Rport 2
Resst
13.
13.6
Start
Resett
Revrtt
925
0.58
PF
purseeds
33.0
PPO
12.0
Diatc
330
0.34
12.6
0.52
0.55
12.2
10
20
20
45
0.50
0.50
0.50
0.00
Spoch
Spoch
Remaadd
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 10: Plotting PPO Elastic Reset IMDB but splitting the results visually between resets. Users with less compute can choose fewer resets. We didn't find improved results after 2 resets for IMDB
19

---

D.2 Reset Type
1.44
Clatiic Fare:
750
33
Glatic SALEE
Euit Raii
Saurrto ENA
Ezech
Ecech
Adwovrt
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 11: Ablation of Elastic Reset, Reset to EMA, and Reset to Init on IMDB. We plot the mean and show error bars for standard error over 5 seeds.
D.3 KL Coefficient
As shown in Figure Elastic Reset performance in relatively unaffected by the presence or absence of KL. In contrast, PPO and REINFORCE performance is known to be sensitive [Lu et al., 2020]. This raises the question of whether different coefficients of the KL loss would lead to different pareto frontiers. We find in Figure 12 that both Multitask and KL with Pretrained methods do have slightly different pareto frontiers. Choosing
Macoaa 0,0
0bc. 0b. (1f--th--.
6.0
04
10099
19000
38300
.0089
5900
10099
20003
35300
40093
50000
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 12: Ablation of KL and Multitask coefficients on Translation Game We exclude KL 0.1 from the pareto graph since it fails. We plot the mean and show error bars for standard error over 5 seeds.
33
Latic Raid
Blc.
6.03
0.02
505
23000
1000
Staps
FR.-0h-00 (pa7)
(a) Task Performance
(b) Language Drift
(c) Parcto Frontier Graph
Figure 13: Ablation of KL coefficients for Elastic Reset on Translation Game, We plot the mean and show error bars for standard error over 5 seeds.
20

---

P00
34-
bcech
tp3is
Sentiment
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 14: Ablation PPO and NLPO coefficients on IMDB. We do more ablations of PPO since NLPO is always similar but worse than PPO. We plot the mean and show error bars for standard error over seeds.
Chite Tapst
Dustc Teyt
0.03
6
866-
332
paplodty (ants
32.
33.
Daatc
22.6
532
0.001
0.52
400
6.0
252
0.54
0.5
3.70
0.50
0.52
beech
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 15: Ablation of KI. coefficients for Elastic Reset on IMDB. We plot the mean and show error bars for standard error over seeds.
Purpoeatty Parppexity (6ri))
Purplonta
12
CL3.1 Ne
0.72
No
0.55
X
Rcaci KL00 No
0.32
0.50
Epoch
Epoch
Damard
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 16: Ablation with/without KI. penalty PPO and Elastic Reset on IMDB. We plot the mean and show error bars for standard error over 5 seeds. For clarity, we cut off PPO without KL to four data points on the pareto graph since diverges too much
21

---

D.4 Decay Rate
213
146e
900
spat)
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 17: Ablation of decay coefficients for Elastic Reset Translation Game. Original decay is 0.99. We plot the mean and show error bars for standard error over 5 seeds.
222
6d0
6.90
0.56
Perttctt
32.5
0.5
3.54
8.56
0.50
0.00
0.52
Saach
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 18: Ablation of decay coefficient for Elastic Reset on IMDB, Original decay is 0.995. We plot the mean and show error bars for standard error over seeds.
D.4.1 Reset Frequency
0.32
Parpiaiity 6aritt
Parplexily
0.70
0.53
0.60
0.30
12
0.5
6.75
Epoch
Epoch
Reward
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 19: Ablation of reset frequency for Elastic Reset on IMDB. We plot the mean and show error bars for standard error over 5 seeds. Although 17 is optimal, results are relatively robust to choice of this hyperparameter
E Explaining Elastic Reset
Here we provide empirical intuition for why Elastic Reset is effective
E.1
Explaining Resets: Value Function
Drift problem of noisy optimization, and Elastic Reset tackles this in ways.
The first is by re-training with an improved value function. Though we begin with sequence-level reward model, over training we learn a #token-leve1* value function. We believe that the value
22

---

function greatly improves over time so early training may have exacerbated drift because it uses the early, sub-optimal value function. With each reset of the policy while maintaining the value function. we re-train with less noise and more direct gradients to alignment. Another way, both the reward model and optimal value function point towards high reward, but not in the same way. The reward model points towards high reward but not necessarily in the most direct way. Using the frozen reward model as our value function leads to equivalent reward but higher drift compared to using an optimal value function.
To demonstrate this phenomenon, we compared what happens we don't 't train our value function at all but use frozen reward model as our value function. Since our value function is GPT-2 but our reward model DistilBERT, we first train GPT-2 reward model similar DistilBERT. We then compare regular PPO with frozen value function i.c. just the sequence-level reward model vs training value function. In Figure 20 we find that the learning value function is clearly better than the frozen one. They both reach the same reward but a more optimal value function leads to less drift for the same reward.
112
Turis
3375
Ta VaLE
12.12
Turephaas
666-
33.99
24.02
32.53
4170
035
Exech
1.00
1pech
530
032
2:100
0.50
032
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 20: Ablation of learning vs frozen value function for PPO on IMDB.
E.2 Explaining Elastic Reset: The Benefit of EMA
The second is by smoothing optimization with an EMA. Work in other fields e.g. DINO in SSL, has leveraged EMA for stability and improved generalization. By resetting to an EMA and resetting our EMA, we further smooth the gradients to alignment. To demonstrate, we run PPO as usual but keep track of its EMA, without resetting to it. We then plot the accuracy of the online and EMA networks training. To make the effect larger, epochs and plot results in Figure 21. We find that just keeping the EMA model an effective way to improve the task / drift tradeoff. But we also that it is noticeably slower than method. PPO Online needs to reach nearly the maximum possible reward for its EMA to improve only slightly on performance.
smaller KL B 0.01 and train our PPO for 100
madi Reian
Plcze Rood
Foii FAS PM
PRI Pis
pap satorrtyoo cettra
4.4
42
52
335
06-
24.
11.
20
130
130
3.22
0.34
8:56
0.50
0.00
0.2
Epars
Epacr
(a) Task Performance
(b) Language Drift
(c) Pareto Frontier Graph
Figure 21: PPO trained on IMDB with a lower KL coefficient for 100 epochs, comparing its online network performance vs EMA model performance. We cut off PPO online to just three data points in the pareto graph for clarity and visual scale.
23


Reviews:
[{"review_id": "es9vz2IezD", "reviewer": "Reviewer_1Gya", "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.", "questions": "- How did you decide on resetting of the first stage? Given stage one, how did you decide on resetting the second stage?\n\n- 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.\n\n[Algorithm] - Sentiment Score - Perplexity \n\n[PPO +SL] - 0.626   - 35.045 \n\n[NLPO + SL.] - 0.611    - 33.82 \n\nIs the proposed approach is not compatible with RL+SL?", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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:\n- 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.\n- Written Presentation: The author's presentation of the proposed algorithm, explanations, and ablations studies were justified and thoroughly explained.\n- Experiments: The authors presented several experiments across three very different domains and showcased the benefit of their proposed method.\n- 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.\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! We’re glad you found our problem compelling and clearly defined, our method easy to implement, and our experiments thorough in demonstrating the benefit of our method.\n\n**On-policy Elastic Reset vs prior work off-policy Reset**\n\nA major difference between our work and prior work in resets is that we reset to an EMA whereas prior works (like Nikishin, Schwarzer et al, 2022) reset to a new random initialization. Resetting to random can indeed help with exploration, but we don’t think that is the benefit of our method. Instead we believe that\n1. After a reset, the value function is maintained and provides a better direction, see Appendix D.4 for an in-depth summary. Supporting this, we find that each run after a reset is better than the last as shown in Figure 1 in the additional pdf at the top.\n2. The EMA model drifts less than the online model because it smoothes out the gradient steps, see Appendix D.5. Just using an EMA is too slow, though, and it is difficult to achieve high reward with an EMA alone. Iteratively resetting to an EMA strikes a good balance. \n\n**Elastic Reset vs KL Divergence**\n\nKL simply pulls all gradients towards the pretrained model. Not all change from the pretrained model is bad, as argued by [Gao et al (2022)](https://arxiv.org/abs/2210.10760) there are good gradient directions to move in, even outside the epsilon ball. Elastic Reset stays close to the model but the EMA also seems to provide an inductive bias against drifting in non-linguistic directions (see Appendix D.4)\n\n**Hyperparameter Choices**\n\nWe chose all our hyperparameters on the validation set of the IMDB and then applied them to StackLLaMA. We tested resetting 1,2,3,5,and 10 times but found no improvement over 2. We tested resetting the EMA model 2,4x less and more frequently but found no improvement over the same reset schedule.\n\n**SL + PPO**\n\nWe do include SL + PPO for the appropriate tasks (Translation Game and StackLLaMA). We don’t do SL+PPO for IMDB following the GRUE benchmark that found stronger results without SL [(Ramamurthy et al, 2022)](https://arxiv.org/abs/2210.01241). \n"}, {"review_id": "9QvJ7YbdgN", "reviewer": "Reviewer_Xquz", "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. \n\nThe 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.\n\nOverall 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.", "questions": "Main questions in Weaknesses section. Some more minor questions:\n- 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)\n- 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?", "limitations": "yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "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).\n- Good comparison to other sensible baselines, with some sensible ablations on the IMDB mock sentiment task, but some areas for improvement here (see Weaknesses)\n- 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\n\n- 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*.\n\t- 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.\n- 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).\n- 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).\n- 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) \n\n## Could use more carefully controlled baselines\n\n- 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.\n- 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.\n\t- 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?\n- 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.\n\n## Mathematical connection/intuition as to why elastic reset differs from distillation and/or KL\n\n- 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)?\n\n## Minor\n- L99 \"Elastic Reset takes inspiration from both of these\"—it's not quite clear what \"both\" refers to.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review! We are glad you find our method simple but effective and tested against sensible baselines across a breadth of experiments. We also believe pareto curves are the right way to think about language drift and hope that our work increases their adoption as a standard in RLHF evaluation. \n\n**Compute Constraints**\n\nOur method does require 2 resets to outperform baselines but notably uses the exact same compute as the baseline to achieve these results. This is in line with all other works using resetting that frequently require many resets to outperform baselines.\n\nWe also agree that compute is an interesting axis of comparison so we’ve created a new graph where we plot each run between resets as its own pareto curve so practitioners can choose at which reset to stop (see Figure 1 in additional pdf attached in the global response). Thank you for the suggestion. A user with less compute can also make resets more frequent to get lower reward but less drift (see Figure 3 in additional pdf). \n\n**Baselines**\n\nWe always compare to the exact model without Elastic Reset (i.e. 0 resets)\n- KL Penalty (Reinforce + KL) in Translation Game\n- PPO (includes KL Penalty) in IMDB\n- PPO (includes KL Penalty) in StackLLaMA\n\n**Fair baselines use different KL coefficients**\n\nBecause PPO drifts so much, it requires a very high KL coefficient and does not work well otherwise. Elastic Reset accounts for drift by resetting so if you use a high KL coefficient, resets will not be useful within IMDB’s 50 epochs as there will not be much drift to counter. But Elastic Reset can still work with that large KL coefficient: simply train for more than 50 epochs to get more drift and reset less frequently (e.g. every 50 epoch instead of 17). \n\nFor a fair comparison, we limit both runs to the same compute. So following all previous work, we choose the best KL coefficient for each baseline and our own method within a fixed compute budget i.e. 50 epochs for IMDB following the [GRUE benchmark](https://rl4lms.apps.allenai.org/grue). In Figures 5c,d we also show that Elastic Reset works across a wide range of smaller KL coefficients that PPO doesn’t and all of them outperform PPO’s best setting. \n\n**Choosing Number of Resets**\n\nOn IMDB, we tested resetting 1,2,3,5,and 10 times but found no improvement over 2. We hope the new reset-pareto graph described above helps choose the number of resets and allows future work to demonstrate the tradeoff.\n\n**Elastic Reset Intuition**\n\nDistillation is one possible way to view the EMA and so Elastic Reset can be seen as a compute-efficient alternative to Iterated Learning (which alternates learning and distillation). We can see that an EMA model drifts less than its online model (see Appendix D.5) which could imply it is an effective distillation for RLHF. \n\nAnother reason why Elastic Reset works is it doesn’t reset the value function. So after each reset, the model starts with a better initialization (EMA) but also better value function that can more directly point towards high reward (see Appendix D.4). This also helps explain the improvements after each reset in Figure 1 in the additional pdf.\n\n**EMA on GPU vs CPU**\n\nIt is indeed easier to keep EMA on GPU and we do this for our experiments. If memory is tight, it is possible to do EMA updates every $n$ steps instead of every step and therefore the GPU to CPU copy is less frequent. On IMDB, we found updating the EMA every 100 steps achieves similar results to updates every step.\n\n**Weighted average instead of EMA**\n\nWe use an EMA since it is a simple and efficient update with each gradient update and our own experiments suggest it mitigates drift (Appendix D.5). Simple weighted averages are an interesting idea and will likely work but we leave it for future work as we don’t expect any improvements or advantages over EMA.\n"}, {"review_id": "pDoTVo2l1W", "reviewer": "Reviewer_tYGJ", "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.  ", "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? ", "limitations": "The authors use Section 8 to describe the limitations. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\nThe 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. \n\n- 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. \n\n- 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.  \n\n- 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. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad the reviewer finds our method simple and effective. We hope the following clarifies the method and importance of our work.\n\n**Novelty**\n\nWe are not aware of any method that resets to an EMA to counter overfitting, and we can’t find any method that also resets the EMA model as we do with Elastic Reset. We would be happy to cite others and explain how our method compares if the reviewer can give references.\n\n**Definition and Importance of “Drift”**\n\nWe define drift as “performance degradations of a language model as a result of RL finetuning and improvement on a reward objective”, see [Lazaridou et al (2020)](https://arxiv.org/abs/2005.07064) for a full taxonomy of drifts. The phenomenon also known as “alignment tax” is generally agreed to be a major downside of RLHF finetuning [(Askell et al, 2021)](https://arxiv.org/abs/2112.00861) and many previous works that we cite have demonstrated the issue in toy tasks [(Lee et al, 2019)](https://arxiv.org/abs/1909.04499) to real-world RLHF chatbots [(Bai et al, 2022)](https://arxiv.org/abs/2204.05862). \n\n**Connection to Continual Learning**\n\nWe agree there are many links to CL. We focus on RLHF-specific terms, known methods, and benchmarks to make our work most useful to the many RLHF projects currently in progress. We are happy to add a summary of CL connections to our related work as well as other references the reviewer may have:\n\nThe pretrain-then-RL-finetune setup with the goal of maintaining pretrained knowledge can be seen as a two-step, RL-specific instance of continual learning and therefore language drift has links to catastrophic forgetting (McCloskey & Cohen, 1989). There is a clear similarity between mitigation methods: rehearsal (Robins, 1995) or experience replay (Rolnick et al, 2019) is equivalent to multitasking with the pretraining objective (Lowe* et al., 2021) and weight-update regularization (Kirkpatrick et al., 2017) has similarities to KL regularization (Jaques et al, 2019).\n\n**Are Tasks Appropriate?**\n\nThe IMDB task is the most popular RLHF benchmark, used as a test-bed for many papers including the original work [(Ziegler et al, 2019)](https://arxiv.org/abs/1909.08593). We use the standard version from the GRUE benchmark [(Ramamurthy et al, 2022)](https://arxiv.org/abs/2210.01241) because it includes a strong, tuned PPO baseline and allows us to demonstrate trends and do quantitative evaluations that are reasonable given our compute. \n\n**Choosing Hyperparameters**\n\nFollowing the GRUE benchmark, we chose all our hyperparameters on the validation set of the IMDB and test on the test set (Table 2). StackLLaMA was too expensive to tune so we simply applied what we knew from IMDB for our hyperparameters. \n\n**Elastic Reset uses smaller KL coefficient**\n\nBecause PPO drifts so much, it requires a very high KL coefficient and does not work well otherwise. Elastic Reset, accounts for drift by resetting so if you use a high KL coefficient, resets will not be useful within IMDB’s 50 epochs as there will not be much drift to counter. We choose an optimal, smaller KL so that there is more drift for Elastic Reset to counter.\n\nWe compare baselines against our own method and choose the best KL coefficient for each within a fixed compute budget i.e. 50 IMDB epochs in the GRUE benchmark. In Figures 5c,d we also show that Elastic Reset works across a wide range of smaller KL coefficients that PPO doesn’t and all of them outperform PPO’s best setting. We find that Elastic Reset even works without KL and is competitive with PPO+KL's best setting (see Figure 2 in the additional pdf included in the global response at the top) \n"}, {"review_id": "BwiDeCejND", "reviewer": "Reviewer_8rTL", "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.", "questions": "1. Why don't you show results with the same KL coefficients for PPO and Elastic Reset in Figure 5?\n Do smaller values of KL not work at all for PPO? What about larger values for Elastic Reset? \n\n2. 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.\n\n3. 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?\n\n4. In Figures 5 and 6, the light blue line can barely be seen on printed paper, can you please use a different / stronger color?", "limitations": "Yes, the authors clearly state the limitations of their study at he end of the paper. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed approach is simple and effective. \n\n2. The evaluation is quite thorough considering multiple settings and thus demonstrating the generality of the approach. \n\n3. The paper is well written and the problem setting is clearly stated. \n\n4. 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. \n", "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.\n\n2. 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?\n\n3. 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. \n\nReferences:\n[1]. Stiennon et al. 2020, Learning to summarize from human feedback. \n\n[2]. Bai et al. 2022, Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review, we are glad you find the approach simple and effective, the writing clear, and the problem of high importance to the community. We agree that drift is a major limitation for RLHF methods and wish to address it robustly.\n\n**Testing with another Reward Model**\n\nWe agree that evaluation is key and find the proposed evals interesting but hope to demonstrate why our protocol is already robust to reward-hacking. Testing with models trained on other reward datasets (e.g. HHH) doesn’t make sense for any of our tasks since they are optimizing very different rewards (StackExchange helpfulness is good coding != HHH Helpfulness which is friendliness / general QA. Summarization is just human-preferred summaries). 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). \n\nWe think our protocol is robust to reward hacking as\n1. The pareto graph measure is robust\n\n A method that does better on the pareto graph must achieve higher reward while staying closer to the original model (lower KL) which, by definition, is reducing drift. By staying closer to the original model, it is likely to better transfer, maintain more of the original model’s ability, and less likely to be over-optimized (reward-hacking) \n\n2. External tests for IMDB and StackLLaMA show robustness\n\nWe evaluate on a set of separate test prompts for the IMDB task (Table 2) and demonstrate that our method outperforms the baselines there, transferring well to unseen prompts. On StackLLaMA, we use a separate coding benchmark HumanEval that should be correlated with good StackExchange answers, but is completely separate from our reward model. We find better coding performance with the Elastic Reset-trained model, clearly showing that it has learned better coding, not reward hacking.\n\n**Varying Reset Interval**\n\nWe provide new results on IMDB varying the reset interval from 10, 17, 26, and 35 epochs, keeping everything else constant. We plot new results in Figure 3 in the additional pdf attached in the global response at the top. As shown in the pareto curve, Elastic Reset is robust to choice of reset interval, we simply chose 17 to fit two resets into 50 IMDB epochs. \n\nSadly, we could not run this for StackLLaMA as it takes too long for this rebuttal period.\n\n**Elastic Reset vs PPO with same KL**\n\nAs shown in Figure 5c, PPO diverges quickly with KL 0.01, and does even worse with smaller KL coefficients\n\n**Experiments without KL**\n\nWe ran both PPO and Elastic Reset without KL penalty in IMDB and plot results in Figure 2 in the additional pdf. PPO diverges quickly but Elastic Reset without KL actually performs comparable to PPO’s best result with KL. \n\n**Cleaner Figure 6**\n\nWe have added smoothing with a rolling mean of window size 10 but could not include it in the additional pdf due to lack of space. We will add this to the appendix for the final paper.\n\n**Partial Resetting**\n\nWe agree this idea is interesting and include it as possible future work in Section 8.\n\n**Light Blue is Too Light**\n\nWe have made the lightest blue stronger/darker for the final paper, please see the graphs in the additional pdf. \n"}]

Author rebuttal:
We've run some extra experiments in response to reviewer comments, please find three figures attached:
1. Plotting each reset of Elastic Reset separately
2. PPO and Elastic Reset with best KL vs without KL (coefficient 0)
3. Elastic Reset over a range of reset frequencies 

Meta-review: The paper proposes a method called elastic reset which adds periodic resets (to an exponentially moving average) to RL-based language generation training with the aim of improving the trade-off between reward optimization and language drift. Experiments on pivot-based translation and IMDB review generation show an improved trade-off between language drift and reward (task performance). There is also a small improvement on RLHF training of LLaMA on StackExchange. 

The paper is well-written and of interest since it aims to address an important current problem. While method is simple it leads to improvements across multiple tasks. The reviews raised a number of questions related to the robustness of the experimental evaluation: The rebutals mostly addressed those and provided some additional results, while the use of open-source LLMs enables reproducibility. More experiments with large-scale LMs and/or more RLHF tasks would further strengthen the results. The connection between the proposed method and related concepts in continual learning and distillation can also be strengthened. However on balance the paper's contribution is strong enough to be accepted.
Decision: Accept (poster)