Language Model Alignment with Elastic Reset

TL;DR

This paper introduces Elastic Reset, a novel algorithm for fine-tuning language models that improves reward achievement while reducing drift, without explicit objective modifications, demonstrated through various benchmarks.

Contribution

Elastic Reset is a new method that periodically resets the model and its EMA to enhance reward and reduce drift without changing the training objective.

Findings

Achieves higher reward with less drift compared to standard methods.

Outperforms baselines on pivot-translation and sentiment tasks.

Produces a more aligned and effective QA chatbot.

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.