Language Model Alignment with Elastic Reset

1 Introduction

• RL finetuning can reduce performance on important benchmarks or cause language to drift from natural syntax and semantics, with extreme reward optimization producing unintelligible text.The introduction frames these effects as reward hacking, alignment tax, or language drift.Source: §1 · p. 1
• The standard approach adds a KL-divergence penalty between the finetuned and pretrained models, but methods remain insufficient over long horizons and require early stopping before catastrophic drift.The paper argues that auxiliary losses such as the KL penalty do not improve the reward–drift tradeoff and can slow training.Source: §1 · p. 2
• Elastic Reset periodically resets the online policy to an EMA of itself, then resets the EMA to the initial model while retaining the value function.The method is motivated by on-policy RLHF lacking a replay buffer to bootstrap learning after a reset.Source: §1 · p. 2
• Across pivot translation, IMDB mock sentiment, and LLaMA-7B technical QA, the paper reports state-of-the-art or baseline-beating performance while mitigating drift or alignment tax.The introduction also says Elastic Reset dominates baselines on the performance–drift tradeoff in the translation benchmark.Source: §1 · p. 2

2 Related Work

• Reward models proxy human preferences, so optimizing them can degrade general NLP performance, overfit preferences, or exploit confounding factors; in language settings, drift can yield incoherent outputs [Askell et al., 2021, Bai et al., 2022, Gao et al., 2022, Stiennon et al., 2020, Lee et al., 2019, Lewis et al., 2017].The passage relates reward hacking to optimizing a proxy without accomplishing the intended task.Source: §2 · pp. 2–3
• Most RLHF methods use a KL penalty against a fixed, usually pretrained, model, while alternatives such as pretraining-task augmentation and student–teacher distillation can be compute-intensive [Jaques et al., 2019, Ziegler et al., 2019, Stiennon et al., 2020, Ouyang et al., 2022, Lowe* et al., 2021, Lu et al., 2020].Delayed masking is another approach, and the paper positions Elastic Reset as drawing on iterative processes and an EMA [Ramamurthy et al., 2022].Source: §2 · p. 3
• Elastic Reset is described as more space- and compute-efficient than distillation and delayed masking because it keeps its EMA on CPU and avoids their extra GPU model or forward-pass costs [Lu et al., 2020, Ramamurthy et al., 2022].The passage says resetting weights and updating the EMA are cheap operations.Source: §2 · p. 3
• Unlike off-policy reset methods that retain a replay buffer, Elastic Reset is an on-policy alternative that retains the value model and policy EMA [Nikishin et al., 2022, D’Oro et al., 2023].The paper also connects preserving pretrained knowledge during RL finetuning to continual learning and catastrophic forgetting.Source: §2 · p. 3

3 Elastic Reset

• The standard KL method estimates divergence between the learning policy and the initial pretrained model over minibatch inputs and outputs, then adds it to reward with coefficient β.The penalty is applied as an auxiliary reward on top of the regular reward model.Source: §3 · p. 3
• Elastic Reset updates an EMA of the online policy after each step, then every n steps resets the online policy to that EMA and the EMA to the initial model.The EMA begins at the initial model and uses decay hyperparameter η.Source: §3 · pp. 3–4
• Elastic Reset can be combined with a KL penalty.Figure 2 identifies the Translation Game, IMDB mock sentiment, and StackLLaMA setups, including StackLLaMA reward modeling.Source: §3 · p. 4, Figure 2 · p. 4

4 Translation Game: Careful Comparison to SOTA

• The benchmark evaluates FR→EN→DE task performance and FR→EN drift using BLEU on held-out validation data.The setup trains FR→EN and EN→DE models on IWSLT, then finetunes FR→DE through English using paired French and German data without English inputs.Source: §4 · p. 4
• REINFORCE reaches the best possible task performance but drifts significantly, whereas Elastic Reset follows baselines until reset and then regains task performance faster than drift accumulates.FROZEN ENGLISH is the lower-bound baseline; REINFORCE trains FR→EN with policy gradients through discrete English tokens.Source: §4 · pp. 4–5
• Elastic Reset achieves nearly the best possible task performance while keeping drift at the initial model’s level, outperforming previous small-scale methods.After a reset, task performance recovers within 5k steps while drift does not return to its previous level; the first reset occurs at 23k steps.Source: Figure 3 · p. 5, §4 · p. 5
• FROZEN ENGLISH avoids drift but limits task performance, while MULTITASK reduces drift without a high task score and KL PENALTY scores highly on task performance but drifts substantially.The comparison includes REINFORCE, KL PENALTY, MULTITASK, SIL, and Elastic Reset, with runs spanning 50k updates and two resets for Elastic Reset.Source: §4 · pp. 4–5

• Single final scores or side-by-side validation curves can obscure which method is preferable when one drifts less but another earns higher reward.Source: §5 · p. 5
• Pareto frontier graphs plot validation task score against drift over training, showing the best achieved reward-versus-drift tradeoff across steps.Task score increases from bottom to top so better methods can be seen to functionally dominate worse ones.Source: §5 · p. 5
• Elastic Reset functionally dominates the baselines, outperforming them at the final validation score and at every task-performance level it achieves.Source: §5 · p. 6

Setup

• The task asks GPT-2, pretrained further on IMDB, to complete reviews for positive sentiment; increasing perplexity on true IMDB data measures drift from the initial model.Source: §6 · p. 6

Experiments

• Elastic Reset achieves a better reward–drift tradeoff than the baselines at every reward level, improving semantic scores faster but initially drifting more before resets mitigate drift.The reported PPO KL coefficient is 0.1, compared with 0.001 for Elastic Reset.Source: §6 · p. 7

Ablations

• Elastic Reset combines the strengths of Reset to EMA at lower rewards and Reset to Init at higher rewards, outperforming both ablations.Reset to EMA maintains performance, while Reset to Init drifts less at higher reward.Source: §6 · p. 7, Figure 5 · p. 7
• Elastic Reset is robust to EMA decay choices; it is also more robust than PPO to KL coefficients, works reasonably well without KL, and its reset-frequency ablation leaves Pareto curves essentially unchanged.The best curves use β = 0.001; combining KL and Elastic Reset performs best.Source: §6 · pp. 7–8
• The authors’ appendix intuition is that resets iteratively improve the value function, while EMA smooths optimization but requires resets to achieve high performance.Source: §6 · p. 8

Setup

• The study finetunes LLaMA-7B on StackExchange, trains a reward model from answer preferences, then uses RL to optimize answers for questions.Training uses an 8-bit model and LoRA; only LoRA parameters are included in Elastic Reset’s EMA.Source: §7 · p. 8

Experiment

• Because per-epoch perplexity is computationally infeasible, training-time drift is measured by KL from the pretrained model over samples.Source: §7 · p. 9

Results

• At high reward, Elastic Reset improves over PPO on the reward–drift tradeoff; at the end it reduces drift while aligning about as well, with slightly better final reward and lower StackExchange perplexity.At low reward, its Pareto curve is equal to or slightly worse than PPO’s.Source: §7 · p. 9, Table 3 · p. 9
• On HumanEval, PPO degrades performance from the initial model, while Elastic Reset maintains performance at similar reward and slightly improves pass@10.Source: §7 · p. 9, Table 3 · p. 9

8 Limitations

• The reset rate was chosen heuristically, and the method resets all trainable parameters rather than testing partial-network resets or weighted averaging.Partial resets and weighted averaging are left for future work.Source: §8 · p. 9
• None of the three evaluated tasks is an ideal RLHF benchmark, and using the optimized reward model may mask overfitting or poor generalization to human preferences.The authors note that no open-source benchmark with a gold reward model was available.Source: §8 · p. 9
• The reported reward results differ from the original StackLLaMA results, likely because of a pending issue affecting its reward model.After fixing the issue, the authors report only a higher base reward and visually similar reward curves.Source: §8 · p. 9
• The study evaluates only on-policy methods, leaving open whether off-policy methods with replay buffers could balance performance and drift differently.The authors describe this as a possibility, not a demonstrated result.Source: §8 · p. 10

9 Conclusion

• Elastic Reset is presented as a cheap method that improves the reward–drift tradeoff while maintaining linguistic capabilities across three tasks and scales from 6-layer Transformers to GPT2 and LLaMA-7B.The authors also call for more research on computationally efficient methods for this tradeoff.Source: §9 · p. 10

A.1 Experimental Details

• The translation experiments use five seeds, a single 16G V100 GPU per run, and report means and standard deviations.BLEU is scored with detokenized sacreBLEU.Source: §A.1 · p. 16

A.2 Translation Game with Transformers

• The appendix uses a Transformer rather than the earlier translation-game LSTM architecture, while reporting similar LSTM results separately.The authors use REINFORCE, noting that Gumbel-Softmax is also feasible and preliminary results were similar.Source: §A.2 · p. 16
• The appendix reports final translation validation scores in Table 4 and curves in Figure 7, but the supplied passages do not give their values.The cited table is labeled as containing Translation Game final validation scores.Source: §A.2 · p. 16, Table 4 · p. 16
• In the translation appendix, iteratively resetting FR→EN to pretrained weights performs well but not as well as Elastic Reset when EN→DE supplies a learning reward model.The authors omit this comparison from the main paper because they focus there on methods that also work with fixed reward models.Source: §A.2 · p. 16

A.3 Translation Game with LSTMs

• The appendix identifies Table 5 as IWSLT validation BLEU for pretrained LSTM models and Table 6 as Multi30k Translation Game BLEU scores with standard deviations.The supplied passages give the table captions but not their cell values.Source: Table 5 · p. 17, Table 6 · p. 17
• On Multi30k, drift is the negative change in FR→EN BLEU, task performance the positive change in FR→EN→DE BLEU, and Combined their sum.The authors caution that their detokenized sacreBLEU evaluation is not directly comparable to earlier work's non-detokenized, custom BLEU evaluation.Source: §A.3 · p. 17
• The authors report that their baseline outperforms Lee et al. and Lu et al., and identify differences in REINFORCE baselines, gradient clipping, and preprocessing order.They describe these as the differences they could identify, not established causes of the performance gap.Source: §A.3 · p. 17
• Because training-update counts are arbitrary and task scores can plateau while drift continues, the authors favor early stopping on a reasonable metric.They choose hyperparameters so methods plateau at the end of training.Source: §A.3 · p. 17
• In their fairseq setup, the authors could not reproduce SIL's reported result that a student sender outperforms its teacher and suspect differences in pretrained models.They report no performance improvement from SIL in their setup.Source: §A.3 · p. 17

B IMDB Mock Sentiment

• Figure 8 compares Elastic Reset, PPO, and NLPO on validation-set language drift and semantic task performance over finetuning, with KL penalties for all methods.The figure reports means and standard errors across five seeds; the supplied caption does not state which method performs best.Source: Figure 8 · p. 18

B.1 Experimental Details

• The RL4LMs experiments use five seeds, one 40G A100 GPU per run, and report means and standard errors.Source: §B.1 · p. 18
• For NLPO, the authors change the learning rate to 1e-6 and target update iterations to 50 after finding an error in the defaults.They use reproducible test scores because they could not reproduce the original NLPO test scores, although validation curves matched those supplied in the appendix.Source: §B.1 · p. 18
• The best Elastic Reset configuration uses PPO defaults with target KL 1.0, KL coefficient 0.001, EMA decay 0.995, and resets every 17 epochs.Source: §B.1 · p. 18

C StackLlama

• Figure 9 presents a rolling-window mean with window 10 for the StackLLaMA results.The caption does not specify the plotted metrics or report an outcome.Source: Figure 9 · p. 18

C.1 Experimental Details

• The StackLLaMA setup uses original LoRA adapters but a newly trained reward model, and evaluates perplexity on 4,000 supervised-training examples.The reward model achieves 64% accuracy versus the original 67%; the authors note their results may differ amid codebase updates and a possible original reward-model creation issue. Training used four 80G A100 GPUs for 20 hours, and Figure 9 presents a rolling-window mean.Source: §C.1 · pp. 18–19

C.2 HumanEval

• HumanEval measures alignment tax using hand-written programming questions intended to be unseen in training, scoring executable functional correctness with pass@1 and pass@10.Evaluation uses the CodeCapybara harness.Source: §C.2 · p. 19

D.1 Reset Number

• For IMDB, results did not improve after two resets, so users with less compute can choose fewer resets.Figure 10 compares PPO and Elastic Reset with results visually split between resets.Source: Figure 10 · p. 19

D.2 Reset Type

• The ablations compare Elastic Reset with EMA versus initialization resets, and test decay coefficients around originals of 0.99 for translation and 0.995 for IMDB.The figures report means and standard errors over five seeds.Source: Figure 11 · p. 20, Figure 17 · p. 22, Figure 18 · p. 22

D.3 KL Coefficient

• Elastic Reset is relatively unaffected by KL, unlike PPO and REINFORCE; multitask and KL-with-pretrained methods have slightly different translation Pareto frontiers.The KL 0.1 setting fails; PPO without KL diverges too much, and NLPO is similar to but worse than PPO.Source: §D.3 · p. 20, Figure 12 · p. 20, Figure 16 · p. 21

D.4.1 Reset Frequency

• On IMDB, reset frequency 17 is optimal, but results are relatively robust to this hyperparameter.Figure 19 reports the reset-frequency ablation.Source: Figure 19 · p. 22

E.1 Explaining Resets: Value Function

• At the same reward, a learning value function yields less drift than using a frozen reward model as the value function.Both reach the same reward; the comparison is shown in Figure 20.Source: §E.1 · p. 23, Figure 20 · p. 23

E.2 Explaining Elastic Reset: The Benefit of EMA

• Keeping PPO’s EMA improves the task–drift tradeoff, but more slowly than Elastic Reset; PPO’s online model nears maximum reward before its EMA improves performance slightly.The comparison uses PPO with KL β = 0.01 for 100 epochs and is shown in Figure 21.Source: §E.2 · p. 23, Figure 21 · p. 23