SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\eta$  and plot our results in Figure 5b finding that our method is quite robust to choice of
(a) Reward
(b) Drift
Figure 6: Elastic Reset compared to PPO on StackLLaMA: A LLaMA-7B model RLHF finetuned on StackExchange as a helpful, technical QA chatbot
decay. Next, we investigate robustness to the choice of KL penalty coefficient  $\beta$ . We search across coefficients that range from 10x smaller to 10x larger than our best KL penalty coefficient for PPO  $(\beta = 0.1)$  and Elastic Reset  $(\beta = 0.001)$ . For visual clarity, we plot PPO in Figure 5c and Elastic Reset in Figure 5d and only plot four points for PPO  $\beta = 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,  $(\beta = 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₂(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_{\theta}$  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 HuggingFaceTRL [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  $100\mathrm{k}$  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  $\eta = 0.995$  and a KL penalty coefficient  $\beta = 0.02$  for both methods. Calculating perplexity for each epoch is computationally
Table 3: Evaluations of the initial (zero-shot) and finetuned StackLLaMA models after 600 epochs. We measure alignment using an average over three reward model trained with three different seeds and drift with perplexity on the data. HumanEval is a programming benchmark that acts as a practical measure of drift / alignment tax.
↑Δ REWARD	↓PERPLEXITY	↑HUMANEVAL (PASS@1,PASS@10)
ZERO-SHOT	0	4.43	11.0, 12.7
PPO	0.81 ± 0.06	4.62	7.8, 10.7
ELASTIC RESET	0.96 ± 0.09	4.57	11.0, 13.0
infeasible so we measure drift during training with the KL from the pretrained model over samples as done previously in other larger-scale RLHF [Bai et al., 2022, Gao et al., 2022].
Results We plot reward in Figure  $6a^2$  and KL from initial model over training in Figure 6b. As noted by Beeching 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.
As a method, Elastic Reset is quite cheap computationally because both EMA updates and resets take negligible 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.
Note that our reward results differ from the original StackLLaMA likely due to a pending issue in their code affecting their reward model. After fixing the issue, the only difference seems to be an increase in base reward, and our reward curves are visually quite similar to theirs. See Appendix C for details
Finally, we note that we follow all previous RLHF work and investigate only on-policy methods [Ziegler et al., 2019, Stiennon et al., 2020, Askell 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.
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.
MN is 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.
P. Ammanabrolu. Re: Reproducing NLPO (RL4LMs), Jan. 2023.
J. 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].
A. 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].
D. 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].
Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, N. Joseph, S. Kadavath, J. Kernion, T. Conerly, S. El-Showk, N. Elhage, Z. Hatfield-Dodds, D. Hernandez, T. Hume, S. Johnston, S. Kravec, 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].
E. 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.
N. 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].
I. 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].
H. Q. To, N. D. Q. Bui, and M. Nguyen. CodeCapybara: Open Source LLaMA Model that Follow Instruction-Tuning for Code Generation., May 2023. URL https://github.com/ FSoft-AI4Code/CodeCapybara. original-date: 2023-04-21T10:28:53Z.
H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Roziere, 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.
M. Utiyama and H. Isahara. A Comparison of Pivot Methods for Phrase-Based Statistical Machine Translation. In Human Language Technologies 2007: The Conference of the North American Chapter of the Association for Computational Linguistics; Proceedings of the Main Conference, pages 484-491, Rochester, New York, Apr. 2007. Association for Computational Linguistics. URL https://aclanthology.org/N07-1061.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is All you Need. In Neural Information Processing Systems, page 11, 2017.
L. von Werra, Y. Belkada, L. Tunstall, E. Beeching, T. Thrush, and N. Lambert. TRL: Transformer Reinforcement Learning, 2023. URL https://github.com/lvwerra/trl.
R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, page 28, 1992.
T. Wolf, L. Debut, V. Sanh, J. Chaumont, 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.
H. 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.
D. M. Ziegler, N. Stiannon, 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].
(b) Language Drift
Figure 7: All transformer model methods on the Translation Game. We measure (a) Task Performance with FR→EN→DE BLEU and (b) Language Drift with FR→EN BLEU, on the validation set during finetuning. We plot the mean and show error bars for standard deviation over 5 seeds. To compare how methods do on both metrics, we plot (c) the best achieved drift vs task performance across finetuning.
A Translation Game
We implement the translation game in the faireq 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 tokenized 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 faireq 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  $\rightarrow$  DE), iteratively resetting our FR  $\rightarrow$  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.
Table 4: Translation Game final validation scores
↑ FR→EN→DE	↑ FR→EN
PRETRAINED	23.8	36.2
FROZEN SENDER	30.8±0.2	36.3±0.1
REINFORCE	33.2±0.3	29.6±0.3
+SIL	28.2±0.4	27.3±4.4
+MULTITASK (S2P)	32.2±0.3	35.2±1.0
+KL PRETRAINED	33.2±0.2	30.8±0.4
+RESET	32.5±0.1	33.3±0.1
+RESET TO EMA	32.9±0.1	36.3±0.1
+ELASTIC RESET	33.0±0.1	36.3±0.1
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
Table 5: BLEU score of IWSLT-pretrained LSTM models on IWSLT 2013 validation set
FR→EN	EN→DE
LEE ET AL. [2019]	34.1	22.0
LU ET AL. [2020]	32.2	20.2
OURS	38.5	23.2
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  $\rightarrow$  EN BLEU from the pretrained model. Task performance is the positive change in FR  $\rightarrow$  EN  $\rightarrow$  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.
Table 6: BLEU scores and ± standard deviation on the Multi30k Translation Game using IWSLT-pretrained LSTM models.
METHOD	FR→EN	FR→EN→DE	DRIFT	PERF	COMBINED
LEE ET AL. [2019]	PRETRAINED	27.2	16.3
REINFORCE	12.4 ± 0.7	24.5 ± 1.5	-14.8	+8.2	-6.6
+ LM	23.6 ± 1.1	27.7 ± 0.4	-3.6	+11.4	+7.8
+ LM + G	24.8 ± 0.4	28.1 ± 0.7	-2.4	+11.8	+9.4
LU ET AL. [2020]	PRETRAINED	29.4	15.7
GUMBEL-SOFTMAX	14.5 ± 0.8	27.1 ± 0.1	-14.9	+11.4	-3.5
SIL	29.4 ± 0.3	28.3 ± 0.2	0	+12.6	+12.6
OURS	PRETRAINED	32.6	18.0
FROZEN-SENDER	32.6 ± 0	25.1 ± 0.1	0	+7.1	+7.1
REINFORCE	30.0 ± 0.2	30.3 ± 0.2	-2.6	+12.3	+9.7
LM λ = 0.01	31.5 ± 0.1	30.2 ± 0.2	-1.1	+12.2	+11.1
MULTITASK λ = 0.1	32.9 ± 0.2	28.5 ± 0.3	+0.3	+10.5	+10.8
SIL	27.7 ± 0.7	24.3 ± 0.1	-4.9	6.3	+1.4
ELASTIC RESET	32.6 ± 0.1	30.0 ± 0.1	0	+12.0	+12.0
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  $\mathrm{FR} \rightarrow \mathrm{EN} \rightarrow \mathrm{DE}$  BLEU scores plateau and there is significant language drift in  $\mathrm{FR} \rightarrow \mathrm{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
(b) Language Drift
Figure 8: IMDB mock sentiment validation set scores for Elastic Reset, PPO, and NLPO. We measure (a) Language Drift and (b) Task Performance via Semantic Score on the validation set over finetuning. All methods also include a KL penalty. We plot mean and standard error across 5 seeds.
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.