SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michael Noukhovitch*
Florian Strub
CIFAR AI Chair
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.
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.,
*Correspondance to michael.noukhovitch@umontreal.ca
Figure 1: Elastic Reset. In actor-critic RL, we reset the policy but maintain the value function.
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.
"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 top- $p$  tokens. Elastic Reset takes inspiration from both of these, using an iterated process and maintaining an EMA model. Apart from better performance, our method is more space efficient and maintains the EMA 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  $\theta$  and the initial, pretrained model  $\theta_0$ . It is calculated empirically over the minibatch of training inputs  $x$  and outputs  $y$  and used as an auxiliary reward with coefficient  $\beta$  on top of the regular reward model  $r$
$$
R (x, y) = r (x, y) - \beta \log \frac {\pi_ {\theta} (y | x)}{\pi_ {\theta_ {0}} (y | x)} \tag {1}
$$
For Elastic Reset, we maintain an exponential moving average  $\bar{\theta}$  of our learning model  $\theta$  and choose a decay hyperparameter parameter  $\eta$ . We initialize  $\bar{\theta} \gets \theta_0$  and after every online model step, we update our EMA model  $\bar{\theta} \gets (1 - \eta)\theta + \eta\bar{\theta}$ . Every  $n$  steps, Elastic Reset sets the online model to the
Figure 2: The Translation Game (top left), IMDB mock sentiment task (bottom left), and StackLLaMA (right). We show all RL finetuning setups and StackLLaMA's reward modelling (top right).
EMA model  $\theta \gets \bar{\theta}$  and sets the EMA model to the initial model  $\bar{\theta} \gets \theta_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  $\rightarrow$  EN and EN  $\rightarrow$  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  $\beta = 0.05$  to regularize the FR  $\rightarrow$  EN model. MULTITASK learning, re-training, or S2P [Lowe* et al., 2021], adds the supervised FR  $\rightarrow$  EN objective on IWSLT pretraining data as an auxiliary task for the FR  $\rightarrow$  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  $\rightarrow$  EN and EN  $\rightarrow$  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  $\beta = 0.001$  and uses an EMA decay  $\eta = 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.  $\beta$ , 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
(a) Task Performance
(b) Language Drift
Figure 3: Comparing Elastic Reset to all baseline 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 standard error over 5 seeds. To compare how methods trade off the two metrics, we plot (c) the best achieved drift vs task performance.
(c) Pareto Frontier Graph
Table 1: Translation Game final validation scores
↑FR→EN→DE	↑FR→EN
FROZEN ENGLISH	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 PENALTY	33.2±0.2	30.8±0.4
+ELASTIC RESET	32.9±0.1	36.3±0.1
possible task performance but drifts significantly. Despite extensive hyperparameter tuning and correspondence with the original authors, SIL does not manage to outperform the REINFORCE baseline so we exclude 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  $23\mathrm{k}$  steps. After the reset, we see a slight performance drop but also a big jump back in terms of FR  $\rightarrow$  EN drift. While the task performance recovers within  $5\mathrm{k}$  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 practitioners 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
(b) Language Drift
Figure 4: Plotting PPO vs Elastic Reset on IMDB but splitting the results visually between resets. We measure (a) Language Drift and (b) Task Performance via Semantic Score on the validation set over finetuning. All methods also include a KL penalty. We plot mean and standard error across 5 seeds.
Table 2:IMDB mock sentiment final test scores
↑ SENTIMENT	↓ PERPLEXITY
ZERO-SHOT	.489±0.01	32.45±0.13
PPO	.596±0.02	33.45±0.40
NLPO	.558±0.06	33.12±0.74
ELASTIC RESET	.611±0.02	33.32±0.23
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 capabilities. 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  $\beta$  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 recovers 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.  $\beta = 0.1$ . We implement ELASTIC RESET on top of PPO with an EMA decay rate of 0.995 and greatly reduce the KL coefficient  $\beta = 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.
(a) Reset Ablation
(c) KL  $\beta$  Ablation for PPO
(b) EMA Decay  $\eta$  Ablation
(d) KL  $\beta$  Ablation for Elastic Reset
Figure 5: Ablating Elastic Reset on the IMDB mock sentiment task. We plot pareto graphs using mean and standard error across 5 seeds.