Critic-Guided Decoding for Controlled Text Generation (2023)

Abstract: Steering language generation towards objectives or away from undesired content has been a long-standing goal in utilizing language models (LM). Recent work has demonstrated reinforcement learning and weighted decoding as effective approaches to achieve a higher level of language control and quality with pros and cons. In this work, we propose a novel critic decoding method for controlled language generation (CriticControl) that combines the strengths of reinforcement learning and weighted decoding. Specifically, we adopt the actor-critic framework and train an LM-steering critic from reward models. Similar to weighted decoding, our method freezes the language model and manipulates the output token distribution using a critic to improve training efficiency and stability. Evaluation of our method on three controlled generation tasks, topic control, sentiment control, and detoxification, shows that our approach generates more coherent and well-controlled texts than previous methods. In addition, CriticControl demonstrates superior generalization ability in zero-shot settings. Human evaluation studies also corroborate our findings.

## introduction with recent advances in large language models (lms), generating natural-sounding text has become feasible @xcite @xcite . however, such text can still be undesirable; for instance, it may be off-topic or biased and otherwise offensive, reflecting the harms in the real-world data @xcite @xcite @xcite @xcite . to address this issue, several controlled text generation methods have been proposed @xcite @xcite . as shown in figure 1 , two major categories of approaches to controlled text generation currently exist: (a) reinforcement learning (rl) and (b) weighted decoding. in the rl approaches, generating a word at each time step is formulated as sequential decision-making leveraging an lm's probability distribution over words @xcite . for example, in the widely used actor-critic framework, the actor-a pretrained lm-predicts the next word, and the critica value network-evaluates the state, i.e., text generated thus far @xcite . unfortunately, lm fine-tuning through rl often suffers from noisy gradient estimation @xcite , which may lead to unstable training, and eventually mode collapse @xcite . even when fine-tuned correctly, optimizing an entire lm for each target attribute, e.g., detoxification, is computationally expensive and memory inefficient for real-world applications @xcite . in the weighted decoding approaches, the underlying lm is kept frozen, and only the final output probability distribution is adjusted @xcite @xcite . more specifically, to condition an lm's output p(x) on a target attribute, an approximation of bayesian decomposition p(x|a) ∝ p(x)p(a|x) is computed with p(a|x) from an external discriminator @xcite . note, these approaches involve a frozen lm with an independent discriminator for each target attribute; in comparison to the rl approaches, this plug-and-play structure allows more efficient training and memory use @xcite , but degrades the text quality in terms of fluency, diversity, etc. @xcite . in this paper, we propose a novel controlled text generation algorithm, critic-guided decoding for controlled text generation (criticcontrol), that re-weight the word distribution from an lm with predicted state-values from the critic network of rl. as a result, criticcontrol raises the likelihood of words that increases the value of the next state over the current state while lowering that of others. since the critic is trained with a frozen lm, the training is stable and efficient as in weighted decoding. moreover, different critics can be used in a plug-and-play manner depending on targeted attributes or reward models. in other words, critic-control combines the strengths of both the rl and the weighted decoding approaches. 1 we demonstrate the efficacy of criticcontrol through experiments on three controlled text generation tasks: topic, sentiment, and toxicity control. for all tasks, we find that criticcontrol consistently outperforms previous methods in terms of control success, fluency, and diversity. also, criticcontrol exhibits strong zero-shot controllability on unseen topics. finally, criticcontrol is compatible with widely used sampling methods like top-k and top-p sampling @xcite for improved text quality. ## related works reinforcement learning rl and the adversarial training formulation were first proposed in the context of language generation as an auxiliary algorithm to mitigate exposure bias in the teacher-forcing training of sequences @xcite @xcite . the main motivation is to incorporate readily-available sequencelevel reward signals into training, such as bleu or rouge @xcite . the success of utilizing rl has been observed in a wide range of tasks, including summarization @xcite @xcite , dialog modeling @xcite @xcite , neural machine translation @xcite , and style transfer @xcite . furthermore, rl has allowed models to capture high-level human feedback @xcite @xcite , which is out of the current work's scope, however. various rl approaches have been explored so far, such as reinforce @xcite @xcite @xcite , the actor-critic framework @xcite @xcite , and ppo @xcite @xcite . more recently, critic-guided methods have been studied to avoid the risk of rl's policy learning collapse. gpt-critic @xcite generates critic-guided texts for enriching task-oriented dialogue datasets. ilql @xcite 's critics guide supervised models to mimic the behavior of reinforced models. our work is the first to incorporate weighted decoding for sequence-level reinforcement learning. weighted decoding freezing the language model and controlling the output probability distribution to suit the purpose is being actively researched @xcite @xcite . this mainstream approach can be used in various scenarios, such as forcing the model to generate text that conforms to certain stylistic or content-based constraints or mitigating aggressive and toxic expressions @xcite @xcite . the pplm @xcite generates texts by plugging a steering layer into the top of language models p(x). then, the gradient from p(a|x) updates iteratively the last hidden representation to desired attributes. it only needs a few layers per each attribute but requires a lot of iterative computations; fudge @xcite economizes p(x|a) by bayesian decomposition p(x|a) ∝ p(x)p(a|x) reward model critic critic reward p(a | x1:end) complete x 1:end result the issue focused on the first two monts of 2015, ... is related to foods? topic prompt foods x4: on x2: issue x3: focused x1: the x4: on x5: ? x2: issue x3: focused x5: foods x1: the rl agent actor actor language models training inference the prefer foods want genius the prefer foods want genius the prefer foods want genius figure 2 : overall flow of criticcontrol. language model completes the sentence in response to the prompt 'the issue focused on', and the reward model evaluates the output by judging the relevance of the generated text to the specified topic. during the inference, critic modifies the output distribution of the language model to ensure that the generated text is appropriately related to the topic when choosing the next token. with p(a|x) computing classifier instead of gradient update methods. for better linguistic quality, gedi @xcite and dexperts @xcite take generative discriminator approaches with two lms conflicting polarly on the desired attribute(e.g., positive vs. negative). they interpret p(a|x) as the degree of disagreement between two conflicting lms. in contrast, criticcontrol is advantageous in taking a generative manner with only a steering layer, like pplm. ## criticcontrol criticcontrol is a controlled text generation framework that consists of a frozen pre-trained language model (actor network) and an extra network for the state-value prediction (critic network). given a part of the sentence, the actor tries to complete the attribute-aware continuations with the critic's support. as an rl formulation, at each step t, states are given as a set of tokens x 1:t-1 = {x 1 , ..., x t-1 } and the attribute a. the policy π θ of the actor samples x @xmath0 note that a t is the advantages function that measures how much the choice x t is better than the critic predicted, and we train the critic to minimize a t . in our algorithm, we freeze the actor model and only train the critic for flexible control, as described in the following section. ## experiments to evaluate the effectiveness of criticcontrol, we conduct experiments on a variety of controlled text generation tasks, including topic control, sentiment control, and detoxification. ## conclusion we propose criticcontrol, a controlled text generation method that takes advantage of both reinforcement learning and weighted decoding. through experiments on various controlled text generation tasks, we demonstrate that criticcontrol can effectively guide language models toward desired attributes while producing high-quality texts. additionally, we show that criticcontrol has a strong generalization ability in zero-shot attribute control tasks by using a general reward model. one of the limitations of this approach might be its high computational cost to explore with 'gpt3-scale' lan-guage models @xcite , and we expect that this can be addressed through offline reinforcement learning @xcite techniques in future research.