Title: Language Model Alignment with Elastic Reset
Venue: NeurIPS 2023
Accepted: yes
Keywords: reinforcement learning from human feedback (rlhf)

language
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

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 at github.com/mnoukhov/elastic-reset.

1 Introduction

Dialogue agents that can effectively interpret and use language are a long-term challenge for NLP. The rise of large pretrained language models (LMs) [Brown et al., 2020] made language model finetuning one of the most promising research directions to achieving capable dialogue agents [Bender and Koller, 2020]. Recently, reinforcement learning (RL) has become a 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 to optimize the reward. But training on the RL objective moves the model away from its pretraining and can reduce performance on important benchmarks [Ouyang et al., 2022] and even drifting away from natural language syntax and semantics [Lazaridou et al., 2020].

'Language drift' [Lee et al., 2019, Lazaridou et al., 2020], 'alignment tax' [Askell et al., 2021], 'reward model overoptimization' [Gao et al., 2022], or LM-specific 'reward-hacking' [Clark and Amodei, 2016] is inherent to RLHF. In the extreme case, models learn to achieve high reward by generating nonsense text that is unintelligible to humans [Lewis et al., 2017]. Methods to mitigate this issue range from re-running pretraining [Lowe* et al., 2021], grounding in other modalities [Lee et al.,

2019], masking the LM generation [Ramamurthy et 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 to 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., 2019]. Still, all methods are insufficient over a large-enough training horizon so models are early-stopped before reaching a 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 seem improve this tradeoff and only serve to slow down training [Gao et al., 2022]. We posit that RLHF training requires a 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 a model has been shown to 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 a replay buffer with which to bootstrap learning after a 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 a large drop in performance. So we propose resetting to a model in-between our online and initial state, specifically to an exponential moving average (EMA) of our online policy, as EMA has been shown to 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 on top of regular RL methods such as REINFORCE [Williams, 1992] or PPO [Schulman et al., 2017].

First, we test our method on a small scale task: pivot translation with a transformer. In this classic benchmark for drift, we outperform all previous baselines and demonstrate state-of-the-art performance. Next, we re-evaluate how performance is measured in the field and argue for a metric 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 a popular task closer to 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 hyperparameters, 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.

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 'reward-hacking' [Clark 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 a proxy of human preferences [Christiano et al., 2017, Ziegler et al., 2019]. This can lead to LMs that optimize a reward model but degrade in

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 as '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 is how to learn the reward in such a way as to minimize the drift, alignment tax, and reward-hacking. The standard approach used in most RLHF is to incorporate a KL 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-Threlkeld 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. [2021]) but this can be compute-intensive 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 topp 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 on 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 a long distillation phase and Ramamurthy et al. [2022] requires an extra forward pass with the masking model. There exist less-popular methods that have been applied to similar issues in RL: prompt-tuning [Singh et al., 2022], using a 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 et al., 2022, D'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 [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-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 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 θ 0 . It is calculated empirically over the minibatch of training inputs x and outputs y and used as an auxiliary reward with coefficient β on top of the regular reward model r

For Elastic Reset, we maintain an exponential moving average ¯ θ of our learning model θ and choose a decay hyperparameter parameter η . We initialize ¯ θ ← θ 0 and after every online model step, we update our EMA model ¯ θ ← (1 -η ) θ + η ¯ θ . Every n steps, Elastic Reset sets the online model to the

EMA model θ ← ¯ θ and sets the EMA model to the initial model ¯ θ ← θ 0 . As with other methods, Elastic Reset can be easily combined with a KL penalty.

4 Translation Game: Careful Comparison to SOTA

Setup We first investigate the pivot-translation benchmark of Lee et al. [2019], which was previously popular for small-scale methods countering drift. Two translation models, French to English (FR → EN) and English to German (EN → DE), are pretrained on IWSLT [Cettolo et al., 2012]. Then, the models are finetuned on translating French to German through English (FR → EN → DE) but given only paired French and German data from Multi30k [Elliott et 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 a held-out FR → EN validation set. Similarly, we measure success on the task with the FR → EN → DE BLEU score. Each model is an encoder-decoder Transformer [Vaswani et al., 2017] with 6 layers and all experimental details are available in Appendix A.

Baselines The EN → DE reward model is simply trained using cross-entropy between predicted and true DE. Our lower-bound baseline is 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 REINFORCE with an exponentially moving baseline and, as with previous work, add a loss for entropy regularization.

When both FR → EN 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 a KL penalty with β = 0 . 05 to regularize the FR → EN model. MULTITASK learning, re-training, or S2P [Lowe* et al., 2021], adds the supervised FR → EN objective on IWSLT pretraining data as an auxiliary task for the FR → EN model. Finally, we implement Seeded Iterated Learning [SIL; Lu et al., 2020], which alternates between n finetuning 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 on top of REINFORCE with a 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. β , because we used a 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 a very high task performance. In 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 → EN performance degrades. REINFORCE achieves the best

0

10000

20000

30000

40000

50000

Steps

24

26

28

30

32

EN

DE

(a) Task Performance

0

10000

20000

30000

40000

50000

Steps

30

32

34

36

38

EN

Elastic Reset

REINFORCE

Multitask

KL Penalty

(b) Language Drift

28

30

32

FR

EN

DE (perf)

30

32

34

36

38

FR

EN (drift)

(c) Pareto Frontier Graph

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 it from the figures for visual clarity but show values in Table 1 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 a tradeoff. Whereas MULTITASK strongly reduces drift, it does not achieve a 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 at 23k steps. After the reset, we see a slight performance drop but also a 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-art results on the benchmark and outperforms all previous small-scale methods.

5 Pareto Frontier Graph

Simply evaluating validation curves side-by-side or looking at a table of final scores, it can 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 on this [Lee et al., 2019] and other benchmarks [Ramamurthy et al., 2022] compare methods using simple point-estimates after training for a specific number of epochs. But this number of epochs is quite arbitrary as models never fully converge to a reward, they are early-stopped such that drift is not catastrophic. Since different 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 a pareto frontier graph to plot each method's achieved task score vs 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 a drift metric can either mean lower scores (perplexity) or higher scores (BLEU) so we always plot task score as increasing from bottom 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

0

10

20

30

40

50

Epoch

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

Reward

Elastic Reset

Start

Reset 1

PPO

(a) Task Performance

0

10

20

30

40

50

Epoch

32.2

32.4

32.6

32.8

33.0

33.2

33.4

Perplexity

Elastic Reset

Start

Reset 1

PPO

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

Perplexity (drift)

Elastic Reset

Start

Reset 1

PPO

(c) Pareto Frontier Graph

does Elastic Reset outperform the baselines at the final validation score, but it functionally dominates such that it is the best method for all levels of task performance it achieves.

6 IMDB Mock Sentiment: Ablation Study for RLHF

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 are susceptible to reward-hacking, shown in Figure 2. The goal is to complete an IMDB movie review with as positive a sentiment as possible. The baseline LM is GPT-2 [Radford et al., 2019] with 117M parameters further pretrained on the IMDB domain [Maas et al., 2011]. We learn a DistilBERT [Sanh et al., 2020] reward model on IMDB to output a sentiment score between 0 (negative) and 1 (positive). We then train our GPT-2 LM to complete different IMDB reviews while maximizing the sentiment reward. Following Ramamurthy et 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 a useful capability, then our initial model was finetuned on IMDB to maximize log-probability, i.e. 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 perplexity on ground truth IMDB data. In contrast to the previous task, a lower perplexity score corresponds to less drift.

Baselines Our main baseline is PPO [Schulman et al., 2017] with Ziegler et al. [2019] modifications for RLHF training, specifically adding a KL penalty with the frozen initial model (equivalent to KL 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 et al., 2022], a recent method that extends PPO with a masking model to counteract drift. The masking model is initialized to the pretrained model and recieves delayed updates; it is set to the online model every n steps. During training, the online model'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 KL coefficient β = 0 . 001 to allow the model to drift more, then reset every 17 epochs such that we get two resets / three iterations during our training.

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

33.6

Perplexity (drift)

Elastic Reset

PPO

Reset to Init

Reset to EMA

(a) Reset Ablation

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

33.6

33.8

Perplexity (drift)

Elastic Reset

0.95

0.99

0.995

0.999

PPO

(b) EMA Decay

η

Ablation

0.50

0.55

0.60

0.65

0.70

0.75

Reward

32

33

34

35

36

37

38

Perplexity (drift)

Elastic Reset

PPO

0

0.01

0.05

0.1

0.5

1.0

(c) KL

β

Ablation for PPO

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

33.6

Perplexity (drift)

Elastic Reset

0

0.0001

0.001

0.01

PPO

(d) KL

β

Ablation for Elastic Reset

Experiments We run all experiments for 5 seeds and report mean and standard error on our validation set for our 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 our 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 a 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 a 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 to drift over the iterations. Looking at the pareto graph in Figure 8c, we see that Elastic Reset far outpaces the baselines and provides a better tradeoff of reward vs drift for every reward.

Ablations We empirically investigate our method through ablations. Throughout this section we run experiments on 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 : our 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 a tradeoff: Reset to EMA is better at lower reward because it maintains performance whereas Reset to Init does better at higher reward because it doesn't drift as much. Elastic Reset combines the benefits of both and outperforms each method.

Next, we consider our method's robustness to hyperparameters. First, we search along different EMA decay rates η and plot our results in Figure 5b finding that our method is quite robust to choice of

0

100

200

300

400

500

600

Epochs

1.5

2.0

2.5

3.0

3.5

4.0

Reward

(a) Reward

0

100

200

300

400

500

600

Epochs

0

5

10

15

20

KL from Initial

Elastic Reset

PPO

(b) Drift

2.5

3.0

3.5

4.0

Reward

0

5

10

15

20

KL from Initial (drift)

(c) Pareto Frontier Graph

decay. Next, we investigate robustness to the choice of KL penalty coefficient β . We search across coefficients that range from 10x smaller to 10x larger than our best KL penalty coefficient for PPO ( β = 0 . 1 ) and Elastic Reset ( β = 0 . 001 ). For visual clarity, we plot PPO in Figure 5c and Elastic Reset in Figure 5d and only plot four points for PPO β = 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. In 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 to PPO, Elastic Reset even works reasonably well without a KL penalty at all, ( β = 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 KL penalty may be unnecessary for RLHF [Bai et al., 2022, Gao et 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.1 we show that resets iteratively improve the value function and in Appendix E.2 we show how EMA smoothes optimization but requires resetting in order to achieve high performance.

7 StackLLaMA: Practical RLHF

Setup Finally, we apply Elastic Reset to a larger-scale RLHF pipeline. We choose LLaMA [Touvron et 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 = log 2 (1 + upvotes ) [Askell et al., 2021] . At most 10 answers are drawn per question, text is cleaned, and HTML is converted to Markdown to make it easier to parse. First, we finetune LLaMA-7B with language modelling on the dataset to get LLaMA-7B-SE. We then further finetune it to get a reward model by learning to predict which of two answers was more upvoted [Stiennon et al., 2020]. For a given question x and two answers y + , -(where y + is preferred), the loss for our reward model r θ is log σ ( r θ ( x, y + ) -r θ ( x, y -)) . Finally, we finetune LLaMA-7B-SE with RL against the reward model by sampling questions from the dataset and learning to optimize the reward for our model'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 on our GPUs. We rely on the HuggingFace trl [von Werra et al., 2023] and peft [Mangrulka et al., 2023] libraries. All technical details are described in Appendix C.

Experiment We again compare to PPO with Ziegler et al. [2019] modifications i.e. KL penalty with a 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 A100s. Since only the LoRA parameters are being learned, we use Elastic Reset on those and therefore maintain only a small percentage of parameters in our EMA. We use a decay rate η = 0 . 995 and a KL penalty coefficient β = 0 . 02 for both methods. Calculating perplexity for each epoch is computationally

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 et al. [2023], the task is much noisier at a larger scale and with a 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 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 at higher reward. Notably, Elastic Reset seems to work out-of-the-box with LoRA. To evaluate drift another way, we get the perplexity of the final models over the StackExchange validation set as in Section 6. For a 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 a 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], a 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@1 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 a large alignment tax. In contrast, Elastic Reset achieves a similar reward but maintains performance, even slightly improving on pass@10, creating an alignment bonus [Askell et al., 2021] instead of tax.

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 to 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 a 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-averaging instead of resets [D'Oro et al., 2023]. We leave both of these directions to future work.

Although we have thoroughly investigated our method on 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 overoptimization and our metric could mask overfitting and lack of generalization to the real world i.e. actual human preferences. An ideal benchmark could include a 'gold' reward model as a proxy for human preference [Gao et al., 2022], but no such benchmarks are open-sourced and available.

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 et al., 2021, Bai et al., 2022]. Previous work in resetting for RL has focused on off-policy methods and demonstrated strong performance [Nikishin et al., 2022, D'Oro 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-policy methods could implicitly balance performance and drift by incorporating a replay buffer with older data.

9 Conclusion

The problems of drift [Lee et al., 2019], alignment tax [Askell et al., 2021], reward model overoptimization [Gao et 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 6 layer Transformers to GPT2 to LLaMA-7B. The problem of drift is currently being addressed with a 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 is a 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 are closer aligned with human preferences [Ziegler et al., 2019]. As well, we hope this work invigorates more research into improving the reward / drift tradeoff of RLHF with a focus on computationally efficient methods that scale.

Acknowledgments and Disclosure of Funding

MNis supported by Fonds de recherche du Québec - Nature et technologies and Sony. MN would like to thank Issam Laradji, ServiceNow Research, Mila, and Compute Canada for providing resources used in the experiments.

References

Ammanabrolu. Re: Reproducing NLPO (RL4LMs), Jan. 2023.

Andreas, D. Klein, and S. Levine. Learning with Latent Language. In NAACL . arXiv, 2018. doi: 10.48550/arXiv.1711.00482. URL http://arxiv.org/abs/1711.00482 . arXiv:1711.00482 [cs].

Askell, Y. Bai, A. Chen, D. Drain, D. Ganguli, T. Henighan, A. Jones, N. Joseph, B. Mann, N. DasSarma, N. Elhage, Z. Hatfield-Dodds, D. Hernandez, J. Kernion, K. Ndousse, C. Olsson, D. Amodei, T. Brown, J. Clark, S. McCandlish, C. Olah, and J. Kaplan. A General Language Assistant as a Laboratory for Alignment, Dec. 2021. URL http://arxiv.org/abs/2112. 00861 . arXiv:2112.00861 [cs].

Bahdanau, K. Cho, and Y. Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. In ICLR . arXiv, 2015. URL http://arxiv.org/abs/1409.0473 . arXiv:1409.0473 [cs, stat].

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. HatfieldDodds, D. Hernandez, T. Hume, S. Johnston, S. Kravec, L. 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 http://arxiv.org/abs/2204.05862 . arXiv:2204.05862 [cs].

Beeching, Y. Belkada, K. Rasul, L. Tunstall, L. v. Werra, N. Rajani, and N. Lambert. StackLLaMA: An RL Fine-tuned LLaMA Model for Stack Exchange Question and Answering, 2023. URL https://huggingface.co/blog/stackllama .

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 Linguistics. doi: 10.18653/v1/2020.acl-main.463. URL https://aclanthology.org/2020.acl-main.463 .

B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-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, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language Models are Few-Shot Learners. In Neural Information Processing Systems . arXiv, July 2020. doi: 10.48550/arXiv.2005.14165. URL http://arxiv.org/abs/2005.14165 . arXiv:2005.14165 [cs].

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://arxiv.org/abs/ 2104.14294 . arXiv:2104.14294 [cs].

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://aclanthology.org/2012.eamt-1.60 .

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, P. 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, I. Babuschkin, S. Balaji, S. Jain, W. Saunders, C. Hesse, A. N. Carr, J. Leike, 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/abs/ 2107.03374 . arXiv:2107.03374 [cs].

F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei. Deep Reinforcement Learning from Human Preferences. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://papers.nips.cc/paper/2017/ hash/d5e2c0adad503c91f91df240d0cd4e49-Abstract.html .

Clark and D. Amodei. Faulty Reward Functions in the Wild, Dec. 2016. URL https://openai. com/blog/faulty-reward-functions/ .

Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer. LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. In NeurIPS . arXiv, Nov. 2022. doi: 10.48550/arXiv.2208.07339. URL http://arxiv.org/abs/2208.07339 . arXiv:2208.07339 [cs].

D'Oro, M. Schwarzer, E. Nikishin, P.-L. Bacon, M. G. Bellemare, and A. Courville. SampleEfficient Reinforcement Learning by Breaking the Replay Ratio Barrier. In International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id= 4GBGwVIEYJ .

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/W16-3210. URL http://aclweb.org/anthology/W16-3210 .

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 on Machine Translation , pages 215-233, Copenhagen, Denmark, Sept. 2017. Association for Computational Linguistics. doi: 10.18653/v1/W17-4718. URL https: //aclanthology.org/W17-4718 .

Gao, J. Schulman, and J. Hilton. Scaling Laws for Reward Model Overoptimization, Oct. 2022. URL http://arxiv.org/abs/2210.10760 . arXiv:2210.10760 [cs, stat].

Gugger, L. Debut, T. Wolf, P. Schmid, Z. Mueller, S. Mangrulkar, M. Sun, and B. Bossan. Accelerate: Training and inference at scale made simple, efficient and adaptable, 2022. URL https://github.com/huggingface/accelerate .

Hochreiter and J. Schmidhuber. Long Short-Term Memory. Neural Computation , 9(8):1735-1780, Nov. 1997. ISSN 0899-7667. doi: 10.1162/neco.1997.9.8.1735. URL https://doi.org/10. 1162/neco.1997.9.8.1735 .

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 http://arxiv.org/abs/2106.09685 . arXiv:2106.09685 [cs].

D. Hunter. Matplotlib: A 2D Graphics Environment. Computing in Science & Engineering , 9(3): 90-95, May 2007. ISSN 1558-366X. doi: 10.1109/MCSE.2007.55. URL https://ieeexplore. ieee.org/document/4160265 . Conference Name: Computing in Science & Engineering.

Jang, S. Gu, and B. Poole. Categorical Reparameterization with Gumbel-Softmax. In ICLR . arXiv, Aug. 2017. doi: 10.48550/arXiv.1611.01144. URL http://arxiv.org/abs/1611.01144 . arXiv:1611.01144 [cs, stat].

Jaques, S. Gu, D. Bahdanau, J. M. Hernández-Lobato, R. E. Turner, and D. Eck. Sequence Tutor: Conservative Fine-Tuning of Sequence Generation Models with KL-control. In International Conference on Machine Learning . arXiv, Oct. 2017. doi: 10.48550/arXiv.1611.02796. URL http://arxiv.org/abs/1611.02796 . arXiv:1611.02796 [cs].

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://arxiv.org/abs/1907.00456 . arXiv:1907.00456 [cs, stat].

Kirby. Spontaneous evolution of 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-0026. doi: 10.1109/4235.918430. 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-Barwinska, 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-3526, Mar. 2017. doi: 10.1073/pnas.1611835114. URL https: //www.pnas.org/doi/10.1073/pnas.1611835114 . Publisher: Proceedings of the National Academy of Sciences.

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://aclanthology.org/P07-2045 .

Lambert, L. Tunstall, N. Rajani, and T. Thrush. HuggingFace H4 Stack Exchange Preference Dataset, 2023. URL https://huggingface.co/datasets/HuggingFaceH4/ stack-exchange-preferences .

Lazaridou and M. Baroni. Emergent Multi-Agent Communication in the Deep Learning Era, July 2020.

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: 10.48550/arXiv.2005.07064. URL http://arxiv.org/abs/2005.07064 . arXiv:2005.07064 [cs].

Lee, K. Cho, and D. Kiela. Countering Language Drift via Visual Grounding, Sept. 2019. URL http://arxiv.org/abs/1909.04499 . arXiv:1909.04499 [cs].

Lewis, D. Yarats, Y. N. Dauphin, D. Parikh, and D. Batra. Deal or No Deal? End-to-End Learning for Negotiation Dialogues, June 2017. URL http://arxiv.org/abs/1706.05125 . arXiv:1706.05125 [cs] version: 1.

Lhoest, A. V. del Moral, Y. Jernite, A. Thakur, P. von Platen, S. Patil, J. Chaumond, M. Drame, J. Plu, L. Tunstall, J. Davison, M. Šaško, G. Chhablani, B. Malik, S. Brandeis, T. L. Scao, V. Sanh, C. Xu, N. Patry, A. McMillan-Major, P. Schmid, S. Gugger, C. Delangue, T. Matussière, L. Debut, S. Bekman, P. Cistac, T. Goehringer, V. Mustar, F. Lagunas, A. M. Rush, and T. Wolf. Datasets: A Community Library for Natural Language Processing, Sept. 2021. URL http: //arxiv.org/abs/2109.02846 . arXiv:2109.02846 [cs].

Li and M. Bowling. Ease-of-Teaching and Language Structure from Emergent Communication. In NeurIPS , 2019. URL http://arxiv.org/abs/1906.02403 . arXiv: 1906.02403.

Loshchilov and F. Hutter. Decoupled Weight Decay Regularization. In International Conference on Learning Representations , Feb. 2022. URL https://openreview.net/forum?id= Bkg6RiCqY7 .

Lowe*, A. Gupta*, J. Foerster, D. Kiela, and J. Pineau. On the interaction between supervision and self-play in emergent communication. In International Conference on Learning Representations , Sept. 2021. URL https://openreview.net/forum?id=rJxGLlBtwH .

Lu, S. Singhal, F. Strub, O. Pietquin, and A. Courville. Countering Language Drift with Seeded Iterated Learning. In ICML . arXiv, Aug. 2020. doi: 10.48550/arXiv.2003.12694. URL http: //arxiv.org/abs/2003.12694 . arXiv:2003.12694 [cs].

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 Linguistics. URL https://aclanthology. org/P11-1015 .

J. Maddison, A. Mnih, and Y. W. Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. In ICLR . arXiv, Mar. 2017. doi: 10.48550/arXiv.1611.00712. URL http://arxiv.org/abs/1611.00712 . arXiv:1611.00712 [cs, stat].

Mangrulka, S. Gugger, L. Debut, Y. Belkada, and S. Paul. PEFT: State-of-the-art ParameterEfficient Fine-Tuning methods, May 2023. URL https://github.com/huggingface/peft . original-date: 2022-11-25T03:51:09Z.

McCloskey and N. J. Cohen. Catastrophic Interference in Connectionist Networks: The Sequential Learning Problem. In G. H. Bower, editor, Psychology of Learning and Motivation , volume 24, pages 109-165. Academic Press, Jan. 1989. doi: 10.1016/S0079-7421(08)60536-8. URL https: //www.sciencedirect.com/science/article/pii/S0079742108605368 .

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 of Diplomacy by combining language models with strategic reasoning. Science , Nov. 2022. doi: 10.1126/science.ade9097. URL https: //www.science.org/doi/10.1126/science.ade9097 . Publisher: American Association for the Advancement of Science.

Nikishin, M. Schwarzer, P. D'Oro, P.-L. 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://proceedings.mlr.press/ v162/nikishin22a.html . ISSN: 2640-3498.

OpenAI. ChatGPT: Optimizing Language Models for Dialogue, Nov. 2022. URL https://webcache.googleusercontent.com/search?q=cache:qLONB_tyjdcJ:https: //openai.com/blog/chatgpt/&cd=1&hl=en&ct=clnk&gl=ca .

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 Linguistics. doi: 10.18653/v1/N19-4009. URL http://aclweb.org/anthology/N19-4009 .

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, P. Welinder, P. Christiano, J. Leike, and R. Lowe. Training language models to follow instructions with human feedback, Mar. 2022. URL http://arxiv.org/abs/2203.02155 . arXiv:2203.02155 [cs].

Paszke, S. Gross, F. 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, J. Bai, and S. Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library, Dec. 2019. URL http://arxiv.org/abs/1912.01703 . arXiv:1912.01703 [cs, stat].

Post. A Call for Clarity in Reporting BLEU Scores. In Proceedings of the Third Conference on Machine Translation: Research Papers , pages 186-191, Brussels, Belgium, Oct. 2018. Association for Computational Linguistics. doi: 10.18653/v1/W18-6319. URL https://aclanthology. org/W18-6319 .

Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language Models are Unsupervised Multitask Learners, 2019. URL https://www.semanticscholar.org/ paper/Language-Models-are-Unsupervised-Multitask-Learners-Radford-Wu/ 9405cc0d6169988371b2755e573cc28650d14dfe .

Ramamurthy, P. Ammanabrolu, K. Brantley, J. Hessel, R. Sifa, C. Bauckhage, H. Hajishirzi, and Y. Choi. Is Reinforcement Learning (Not) for Natural Language Processing: Benchmarks, Baselines, and Building Blocks for Natural Language Policy Optimization. In ICLR . arXiv, 2022. URL http://arxiv.org/abs/2210.01241 . arXiv:2210.01241 [cs].

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: //arxiv.org/abs/2209.15342 . arXiv:2209.15342 [cs, math].

Robins. Catastrophic Forgetting, Rehearsal and Pseudorehearsal. Connection Science , 7 (2):123-146, June 1995. ISSN 0954-0091. doi: 10.1080/09540099550039318. URL https://doi.org/10.1080/09540099550039318 . Publisher: Taylor & Francis _eprint: https://doi.org/10.1080/09540099550039318.

Rolnick, A. Ahuja, J. Schwarz, T. P. Lillicrap, and G. Wayne. Experience Replay for Continual Learning, Nov. 2019. URL http://arxiv.org/abs/1811.11682 . arXiv:1811.11682 [cs, stat].

Sanh, L. Debut, J. Chaumond, and T. Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter, Feb. 2020. URL http://arxiv.org/abs/1910.01108 . arXiv:1910.01108 [cs].

Schulman, S. Levine, P. Moritz, M. I. Jordan, and P. Abbeel. Trust Region Policy Optimization. In ICML . arXiv, 2015. doi: 10.48550/arXiv.1502.05477. URL http://arxiv.org/abs/1502. 05477 . arXiv:1502.05477 [cs].

Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal Policy Optimization Algorithms, Aug. 2017. URL http://arxiv.org/abs/1707.06347 . arXiv:1707.06347 [cs].

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://arxiv.org/abs/ 2206.08349 . arXiv:2206.08349 [cs].

Steinert-Threlkeld, X. Zhou, Z. Liu, and C. M. Downey. Emergent Communication Fine-tuning (EC-FT) for Pretrained Language Models. In Emergent Communication Workshop at ICLR 2022 , June 2022. URL https://openreview.net/forum?id=SUqrM7WR7W5 .

Stiennon, L. Ouyang, J. Wu, D. M. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei, and P. Christiano. Learning to summarize from human feedback. In NeurIPS . arXiv, 2020. URL http://arxiv.org/abs/2009.01325 . arXiv:2009.01325 [cs].

Sutskever, O. Vinyals, and Q. V. Le. Sequence to Sequence Learning with Neural Networks. In Neural Information Processing Systems . arXiv, Dec. 2014. URL http://arxiv.org/abs/1409. 3215 . arXiv:1409.3215 [cs].

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://github.com/ FSoft-AI4Code/CodeCapybara . original-date: 2023-04-21T10:28:53Z.

Touvron, T. 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 http://arxiv.org/abs/2302.13971 . arXiv:2302.13971 [cs] version: 1.

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 Linguistics; Proceedings of the Main Conference , pages 484-491, Rochester, New York, Apr. 2007. Association for Computational Linguistics. URL https://aclanthology.org/N07-1061 .

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.

von Werra, Y. Belkada, L. Tunstall, E. Beeching, T. Thrush, and N. Lambert. TRL: Transformer Reinforcement Learning, 2023. URL https://github.com/lvwerra/trl .

J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning , page 28, 1992.

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, Y. Jernite, J. Plu, C. Xu, T. Le Scao, S. Gugger, M. Drame, Q. Lhoest, and A. Rush. Transformers: State-of-the-Art Natural Language Processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, Oct. 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-demos.6. URL https: //aclanthology.org/2020.emnlp-demos.6 .

Zhou, A. Vani, H. Larochelle, and A. Courville. Fortuitous Forgetting in Connectionist Networks. In International Conference on Learning Representations , Mar. 2022. URL https: //openreview.net/forum?id=ei3SY1_zYsE .

M. Ziegler, N. Stiennon, J. Wu, T. B. Brown, A. Radford, D. Amodei, P. Christiano, and G. Irving. Fine-Tuning Language Models from Human Preferences, 2019. URL http://arxiv.org/abs/ 1909.08593 . arXiv:1909.08593 [cs, stat].

0

10000

20000

30000

40000

50000

Steps

24

26

28

30

32

34

EN

DE

EMA set

EMA reset

KL pretrained

REINFORCE

Mulitask

SIL

(a) Task Performance

0

10000

20000

30000

40000

50000

Steps

30

32

34

36

38

EN

EMA set

EMA reset

KL pretrained

REINFORCE

Mulitask

SIL

(b) Language Drift

29

30

31

32

33

FR

EN

DE (perf)

30

31

32

33

34

35

36

37

38

FR

EN (drift)

EMA set

EMA reset

KL pretrained

REINFORCE

Multitask

SIL

(c) Pareto Frontier Graph

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 et al., 2019]. All experiments are run with 5 seeds where each run uses a 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 et al. [2019] (as well as previous work on the Translation Game) use 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 we switch to a Transformer architecture. LSTMs results are similar and results are available in Appendix A.3. Also note that we use REINFORCE as our estimator but Gumbel-Softmax [Jang et al., 2017, Maddison et 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-en architecture from the fairseq library [Ott et al., 2019] with their default IWSLT DE-EN hyperparameters and training, specifically training with AdamW [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 are 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 we have a learning reward model (EN → DE), iteratively resetting our FR → EN 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 not include these results in the main paper because we focused on methods that could also work with fixed reward models.

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 et al. [2020] in Table 5

Next we finetune on Multi30k using the same hyperparameters as Lu et al. [2020]. We plot results in Table 6 and compare to published numbers from previous work. As usual, drift is the negative change in FR → EN BLEU from the pretrained model. Task performance is the positive change in FR → EN → DE BLEU from the pretrained models. Combined is 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.

Our results for the baseline are 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 baseline for REINFORCE whereas [Lee et al., 2019] use an Actor-Critic method, (2) Lu et al. [2020] uses 0.1 gradient clipping and we do not use gradient clipping (3) in preprocessing Multi30k, we first tokenize [Koehn et al., 2007] then lowercase whereas previous works did the opposite order.

A more reasonable explanation for the improvement in results is how we choose to evaluate. Previous work simply ran all methods for the same number of updates but this doesn't account for, even implicit, hyperparameter optimization. Previous methods show that the baseline's FR → EN → DE 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 is arbitrary, we believe that early stopping on a reasonable metric is 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 a student sender can outperform a teacher sender that it is distilling from. We could not reproduce this and believe it is a difference in the pretrained models.

B IMDB Mock Sentiment

0

10

20

30

40

50

Epoch

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

Reward

(a) Task Performance

0

10

20

30

40

50

Epoch

32.2

32.4

32.6

32.8

33.0

33.2

33.4

Perplexity

Elastic Reset

PPO

NLPO

(b) Language Drift

0.500

0.525

0.550

0.575

0.600

0.625

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

Perplexity (drift)

(c) Pareto Frontier Graph

B.1 Experimental Details

We run experiments and implement our method in the RL4LMs library [Ramamurthy et al., 2022]. All experiments are run with 5 seeds where run uses a single 40G A100 GPU. All plots show the mean and standard error over seeds.

For PPO, we use the default hyperparemters provided by Ramamurthy et al. [2022]. We also compared to NLPO, we found the defaults had a mistake and after communication with Ramamurthy et al. [2022] we changed the learning rate to 1e-6 and target update iterations to 50. We plot results in Figure 8. This still did not manage to reproduce the original NLPO test scores from Ramamurthy et al. [2022] but we found that our 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 NLPO in lieu of the original test scores.

Our best Elastic Reset hyperparameters are the default PPO parameters ( gpt2_ppo.yml ) from RL4LMs [Ramamurthy et al., 2022] with a 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 of 0.995 and reset every 17 epochs. Configs to reproduce the experiments can be found in the RL4LMs folder under scripts/training/task_configs/imdb_text_continuation

C StackLlama

0

100

200

300

400

500

600

Epochs

2.00

2.25

2.50

2.75

3.00

3.25

Reward

Elastic Reset

PPO

(a) Reward

0

100

200

300

400

500

600

Epochs

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

KL from Initial

Elastic Reset

PPO

(b) KL from Initial

C.1 Experimental Details

We follow the original StackLlama 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

[Lhoest et al., 2021]. We use the original authors' LoRA adapter weights to create our supervised model LLaMA-7B-SE 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 it and our final model achieves 64% accuracy compared to the original 67%. Although we follow the original authors method and use their codebase, 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 is run on 4x 80G A100 GPUs 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/stackllama/scripts/configs

For visual clarity, we provide a version of our StackLLaMA results with a rolling window mean in Figure 9

C.2 HumanEval

Our measure of alignment tax, HumanEval [Chen et al., 2021], is a programming benchmark where each question was hand-written by humans (OpenAI engineers) to be unseen in training data. We believe these questions were unseen in 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 i.e. we actually execute the code LLaMA wrote and see if it gets the right output. Since we decode by sampling, we 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 @1) 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

10

20

30

40

50

Epoch

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

Reward

Elastic Reset

Start

Reset 1

PPO

(a) Task Performance

0

10

20

30

40

50

Epoch

32.2

32.4

32.6

32.8

33.0

33.2

33.4

Perplexity

Elastic Reset

Start

Reset 1

PPO

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

Perplexity (drift)

Elastic Reset

Start

Reset 1

PPO

(c) Pareto Frontier Graph

D.2 Reset Type

0

10

20

30

40

50

Epoch

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

0.66

Reward

Elastic Reset

PPO

Reset to Init

Reset to EMA

(a) Task Performance

0

10

20

30

40

50

Epoch

32.2

32.4

32.6

32.8

33.0

33.2

33.4

33.6

Perplexity

Elastic Reset

PPO

Reset to Init

Reset to EMA

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

33.6

Perplexity (drift)

Elastic Reset

PPO

Reset to Init

Reset to EMA

(c) Pareto Frontier Graph

D.3 KL Coefficient

As shown in Figure 5, 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 a

0

10000

20000

30000

40000

50000

Steps

24

26

28

30

32

34

EN

DE

Multitask

0.05

0.1

0.5

1.0

KL Pretrained

0.005

0.01

0.05

0.1

(a) Task Performance

0

10000

20000

30000

40000

50000

Steps

30

31

32

33

34

35

36

37

38

EN

Multitask

0.05

0.1

0.5

1.0

KL Pretrained

0.005

0.01

0.05

0.1

(b) Language Drift

26

27

28

29

30

31

32

33

FR

EN

DE (perf)

30

31

32

33

34

35

36

37

38

FR

EN (drift)

Multitask

0.05

0.1

0.5

1.0

KL Pretrained

0.005

0.01

0.05

Elastic Reset

(c) Pareto Frontier Graph

0

10000

20000

30000

40000

50000

Steps

26

28

30

32

EN

DE

Elastic Reset

0.001

0.005

0.01

0.05

(a) Task Performance

0

10000

20000

30000

40000

50000

Steps

31

32

33

34

35

36

37

38

EN

Elastic Reset

0.001

0.005

0.01

0.05

(b) Language Drift

27

28

29

30

31

32

33

FR

EN

DE (perf)

30

31

32

33

34

35

36

37

38

FR

EN (drift)

Elastic Reset

0.001

0.005

0.01

0.05

Multitask

KL Pretrained

(c) Pareto Frontier Graph

0

10

20

30

40

50

Epoch

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

0.90

Reward

PPO

0.01

0.05

0.1

0.5

1.0

NLPO

0.001

0.01

0.1

(a) Task Performance

0

10

20

30

40

50

Epoch

32

34

36

38

40

42

44

Perplexity

PPO

0.01

0.05

0.1

0.5

1.0

NLPO

0.001

0.01

0.1

(b) Language Drift

0.50

0.55

0.60

0.65

0.70

0.75

Sentiment (perf)

32

33

34

35

36

37

38

Perplexity (drift)

PPO

0.01

0.05

0.1

0.5

1.0

NLPO

0.001

0.01

0.1

Elastic Reset

(c) Pareto Frontier Graph

0

10

20

30

40

50

Epoch

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

Reward

Elastic Reset

0.0001

0.001

0.01

0.1

(a) Task Performance

0

10

20

30

40

50

Epoch

32.2

32.4

32.6

32.8

33.0

33.2

33.4

Perplexity

Elastic Reset

0.0001

0.001

0.01

0.1

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

33.6

Perplexity (drift)

Elastic Reset

0

0.0001

0.001

0.01

PPO

(c) Pareto Frontier Graph

0

10

20

30

40

50

Epoch

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

Reward

Elastic Reset

KL 0.001

No KL

PPO

KL 0.1

No KL

(a) Task Performance

0

10

20

30

40

50

Epoch

32

33

34

35

36

Perplexity

Elastic Reset

KL 0.001

No KL

PPO

KL 0.1

No KL

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

0.66

Reward

32

33

34

35

36

Perplexity (drift)

Elastic Reset

KL 0.001

No KL

PPO

KL 0.1

No KL

(c) Pareto Frontier Graph

D.4 Decay Rate

0

10000

20000

30000

40000

50000

Steps

26

28

30

32

EN

DE

EMAset

0.95

0.99

0.995

0.999

(a) Task Performance

0

10000

20000

30000

40000

50000

Steps

35.0

35.5

36.0

36.5

37.0

37.5

EN

EMAset

0.95

0.99

0.995

0.999

(b) Language Drift

29.0

29.5

30.0

30.5

31.0

31.5

32.0

32.5

33.0

FR

EN

DE (perf)

32

33

34

35

36

37

38

FR

EN (drift)

EMAset

0.95

0.99

0.995

0.999

Multitask

KL Pretrained

(c) Pareto Frontier Graph

0

10

20

30

40

50

Epoch

0.50

0.55

0.60

0.65

0.70

Reward

Elastic Reset

0.95

0.99

0.995

0.999

(a) Task Performance

0

10

20

30

40

50

Epoch

32.5

33.0

33.5

34.0

34.5

Perplexity

Elastic Reset

0.95

0.99

0.995

0.999

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.2

32.4

32.6

32.8

33.0

33.2

33.4

33.6

33.8

Perplexity (drift)

Elastic Reset

0.95

0.99

0.995

0.999

PPO

(c) Pareto Frontier Graph

D.4.1 Reset Frequency

0

20

40

60

80

100

Epoch

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

Reward

Elastic Reset

10

17

26

35

(a) Task Performance

0

20

40

60

80

100

Epoch

32

33

34

35

36

37

Perplexity

Elastic Reset

10

17

26

35

(b) Language Drift

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

Reward

32

33

34

35

36

37

Perplexity (drift)

Elastic Reset

10

17

26

35

(c) Pareto Frontier Graph

E Explaining Elastic Reset

Here we provide empirical intuition for why Elastic Reset is effective

E.1 Explaining Resets: Value Function

Drift is a problem of noisy optimization, and Elastic Reset tackles this in two ways.

The first is by re-training with an improved value function. Though we begin with a sequence-level reward model, over training we learn a *token-level* value function. We believe that the value

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 if we don't train our value function at all but use a frozen reward model as our value function. Since our value function is GPT-2 but our reward model is DistilBERT, we first train a GPT-2 reward model similar to DistilBERT. We then compare regular PPO with a frozen value function i.e. just the sequence-level reward model vs a 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.

0

20

40

60

80

100

Epoch

0.50

0.52

0.54

0.56

0.58

0.60

0.62

0.64

Reward

Training Value

(a) Task Performance

0

20

40

60

80

100

Epoch

32.25

32.50

32.75

33.00

33.25

33.50

33.75

34.00

Perplexity

Training Value

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.25

32.50

32.75

33.00

33.25

33.50

33.75

34.00

Perplexity (drift)

Training Value

(c) Pareto Frontier Graph

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 over training. To make the effect larger, we use a smaller KL β = 0 . 01 and train our PPO for 100 epochs and plot results in Figure 21. We find that just keeping the EMA model is an effective way to improve the task / drift tradeoff. But we also see that it is noticeably slower than our method. PPO Online needs to reach nearly the maximum possible reward for its EMA to improve only slightly on performance.

0

20

40

60

80

100

Epoch

0.5

0.6

0.7

0.8

0.9

Reward

Elastic Reset

PPO - Online

PPO - EMA

(a) Task Performance

0

20

40

60

80

100

Epoch

35

40

45

50

55

60

Perplexity

Elastic Reset

PPO - Online

PPO - EMA

(b) Language Drift

0.50

0.52

0.54

0.56

0.58

0.60

0.62

Reward

32.5

33.0

33.5

34.0

34.5

Perplexity (drift)

Elastic Reset

PPO - Online

PPO - EMA

(c) Pareto Frontier Graph

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes, the authors clearly state the limitations of their study at he end of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors use Section 8 to describe the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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)\n\nWEAKNESSES:\n## 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.\n\nQUESTIONS:\nMain 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFinetuning 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.\n\nSTRENGTHS:\nAddressing 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.\n\nWEAKNESSES:\nAlthough 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}