Language Model Alignment with Elastic Reset
Michael Noukhovitch Samuel Lavoie Florian Strub Aaron Courville
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.