Reward-Augmented Decoding: Efficient Controlled Text Generation With a Unidirectional Reward Model (2023)

Abstract: While large language models have proven effective in a huge range of downstream applications, they often generate text that is problematic or lacks a desired attribute. In this paper, we introduce Reward-Augmented Decoding (RAD), a text generation procedure that uses a small unidirectional reward model to encourage a language model to generate text that has certain properties. Specifically, RAD uses the reward model to score generations as they are produced and rescales sampling probabilities to favor high-reward tokens. By using a unidirectional reward model, RAD can cache activations from prior generation steps to decrease computational overhead. Through experiments on generating non-toxic and sentiment-controlled text, we demonstrate that RAD performs best among methods that change only the generation procedure and matches the performance of state-of-the-art methods that involve re-training the language model. We further validate that RAD is effective on very large language models while incurring a minimal computational overhead.

## introduction large language models (llms, @xcite @xcite are seeing widespread adoption thanks to the fact that they can perform many language tasks and generate coherent long-form text. as llms are deployed in situations where they interact with humans, it can be beneficial to control the language model so that it generates text with certain properties @xcite -for example, we might desire generations that are unbiased, non-toxic, and helpful. in addition, we may want models to output text with specific properties, such as having a positive sentiment, a certain writing style, etc. typically, llms pre-trained on uncurated large-scale text corpora can generate text that does not have these desired attributes @xcite ## reward-augmented decoding at a high level, reward-augmented decoding, as shown in fig. 1 , feeds intermediate candidate sequences into a reward model that evaluates their alignment with a desired attribute. then, at each decoding step, rad uses the predicted reward of each candidate sequence to modify the token probabilities output by the language model. in this section, we describe these steps in detail. refer to ## experiments we now evaluate rad's performance in two standard settings: preventing language models from generating toxic text @xcite and controlling the sentiment of generated text @xcite . baselines in both settings, we consider the same set of baselines as @xcite , namely: the performance of the base language model itself without any interventions; pplm @xcite , which uses a bag-of-word classifier to update lm hidden states during decoding; gedi @xcite and dexperts @xcite , which use signals from auxiliary language models to modify lm probabilities in one pass; rectification @xcite , which adjusts lm probabilities proportional to the risk of resulting in a toxic generation; dapt @xcite , which further trains the model on data that has the desired property; ppo @xcite , which updates the lm with gradients from the reward model; quark @xcite , which performs parameter-efficient fine-tuning on attributeannotated data @xcite ; and ctrl @xcite , a language model trained to condition on control codes. unless otherwise mentioned, we report results directly from @xcite and @xcite , which can be consulted for further baseline details.