SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For PPO, we use the default hyperparameters 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/taskConfigs/imdb_text continua
C StackLlama
(b) KL from Initial
Figure 9: Rolling Window Mean (window = 10) version of the StackLLaMA results
We follow the original StackLlama and use theTRL 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 coefficient 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 config 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
(b) Language Drift
Figure 10: Plotting PPO vs Elastic Reset on IMDB but splitting the results visually between resets. Users with less compute can choose fewer resets. We didn't find improved results after 2 resets for IMDB
D.2 Reset Type
(b) Language Drift
Figure 11: Ablation of Elastic Reset, Reset to EMA, and Reset to Init on IMDB. We plot the mean and show error bars for standard error over 5 seeds.
D.3 KL Coefficient
As shown in Figure 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
(b) Language Drift
Figure 12: Ablation of KL and Multitask coefficients on Translation Game. We exclude KL 0.1 from the pareto graph since it fails. We plot the mean and show error bars for standard error over 5 seeds.
(a) Task Performance
Figure 13: Ablation of KL coefficients for Elastic Reset on Translation Game. We plot the mean and show error bars for standard error over 5 seeds.
(b) Language Drift
Figure 14: Ablation of PPO and NLPO coefficients on IMDB. We do more ablations of PPO since NLPO is always similar but worse than PPO. We plot the mean and show error bars for standard error over 5 seeds.
(a) Task Performance
Figure 15: Ablation of KL coefficients for Elastic Reset on IMDB. We plot the mean and show error bars for standard error over 5 seeds.
(a) Task Performance
Figure 16: Ablation with/without KL penalty for PPO and Elastic Reset on IMDB. We plot the mean and show error bars for standard error over 5 seeds. For clarity, we cut off PPO without KL to four data points on the pareto graph since it diverges too much
D.4 Decay Rate
Figure 17: Ablation of decay coefficients for Elastic Reset on Translation Game. Original decay is 0.99. We plot the mean and show error bars for standard error over 5 seeds.
(a) Task Performance
Figure 18: Ablation of decay coefficient for Elastic Reset on IMDB. Original decay is 0.995. We plot the mean and show error bars for standard error over 5 seeds.
D.4.1 Reset Frequency
(a) Task Performance
Figure 19: Ablation of reset frequency for Elastic Reset on IMDB. We plot the mean and show error bars for standard error over 5 seeds. Although 17 is optimal, results are relatively robust to choice of this hyperparameter
E Explaining Elastic Reset
Here we provide empirical intuition for why Elastic Reset is effective
E.1 Explaining Resets: Value Function
Drift 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.
(b) Language Drift
Figure 20: Ablation of learning vs frozen value function for PPO on IMDB.
E.2 Explaining Elastic Reset: The Benefit of EMA
The second is by smoothing optimization with an EMA. Work in other fields e.g. DINO in SSL, has leveraged EMA for stability and improved generalization. By resetting to an EMA and resetting our EMA, we further smooth the gradients to alignment. To demonstrate, we run PPO as usual but keep track of its EMA, without resetting to it. We then plot the accuracy of the online and EMA networks over training. To make the effect larger, we use a smaller KL  $\beta = 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.
(a) Task Performance
Figure 21: PPO trained on IMDB with a lower KL coefficient for 100 epochs, comparing its online network performance vs EMA model performance. We cut off PPO - online to just three data points in the pareto graph for clarity and visual scale.