Plug-in Language Model: Controlling Text Generation with a Simple Regression Model (2024)

Abstract: Large-scale pre-trained language models have displayed unrivaled capacity in generating text that closely resembles human-written text. Nevertheless, generating texts adhering to specific conditions without fine-tuning or adding new parameters can be challenging. Contemporary approaches commonly rely on either prompts or auxiliary models to avoid modifying the language models. These auxiliary models are designed to assess whether a generated token contributes to meeting the desired requirements. These approaches adjust the distribution of the next token during the inference phase by leveraging the prediction score of the desired attribute to calculate gradients. However, these auxiliary models typically require the language model’s latent states. This prerequisite challenges integrating various existing black box attribute models or tools. We present the Plug-in Language Model (PiLM) as a solution to address the limitations. PiLM leverages reinforcement learning to utilize black box tools directly, adjusting the latent state to control text generation. However, performing backpropagation during the inference phase is time-consuming for PiLM. By replacing backpropagation with a simple regression model, PiLM can achieve an inference time comparable to that of the original LLM. Experiment results show that our approaches in this paper outperform existing state-of-the-art methods that rely on gradient-based, weighted decoding, or prompt-based methodologies.

## introduction large-scale language models can already generate text nearly indistinguishable from human-written content in terms of grammar and fluency. however, the primary challenge lies in exerting precise control over the generated text to align it with specific semantic requirements. without robust control * corresponding author mechanisms, there is a potential risk that the generated text may deviate from the intended meaning or even include offensive or derogatory language. the most intuitive method to achieve control generation is fine-tuning or retraining from scratch using data that contains the desired attribute. these approaches achieve notable breakthroughs in enhancing their performance. however, there is a trend of language models becoming increasingly larger, leading to a rise in the cost of training. therefore, there is a growing emphasis on methods for controlling at generation time. prior research involved training an auxiliary classification model to support the language model. gradient-based methods @xcite ) modified the hidden representation through gradient descent in the inference phase. weighted decoding methods @xcite integrate the original output distribution with an auxiliary attribute distribution. methods without an auxiliary model, such as prompt-based approaches @xcite concatenating prompt embedding and input text to influence the latent representation, thereby achieving control over the language model's output. energy-based methods @xcite view controlled generation as an optimization problem, iteratively seeking text with lower energy. in gradient-based methods, supplementary auxiliary classifiers typically use the language model's latent states as inputs to predict whether it aligns with the desired attribute. therefore, most black box tools that take text as input cannot serve as these attribute classifiers. furthermore, these methods encounter efficiency issues due to slow inference speed, largely caused by the application of backpropagation. weighted decoding methods impact the output by adjusting the distribution of the next token toward a specific attribute without altering the latent representation. avoiding gradient updating can significantly shorten the time. however, these methods frequently yield outputs that lack fluency or coherence. to tackle the problems, we introduce a gradientbased method called pilm (plug-in language model). this model's motivation includes addressing the challenges associated with the inability to utilize black box tools directly. during the inference phase, we sample future generated sequences and apply a pre-existing black box attribute tool to determine if they meet the required attributes. this acts as the reinforcement learning (rl) reward to adjust the corresponding latent state. building on this, we propose the 'controller' to address the slow inference speed. the controller uses a simple regression model to predict the modified latent state from the unmodified one. training pairs are easily gathered during reinforcement updating. another benefit of considering a more extended future context is that it allows more textual information to adjust the latent state more accurately. in this paper, we experiment with our proposed method on three distinct tasks: sentiment control, topic control, and language detoxification. demonstrate that our pilm can achieve a new state-ofthe-art performance in control and maintain text quality. our main contributions can be summarized as follows: (1) we propose a novel method that enables language models to utilize black box tools directly, eliminating the need to train attributespecific classifiers and making it more convenient when adding new attributes or switching to different language models. (2) in contrast to prior approaches that depend on classifiers to determine update directions, considering a single token always contains limited semantic information, we incorporate future sequence considerations to improve the accuracy of latent updates. (3) we introduce a method to address the bottleneck of gradient-based methods during inference time, utilizing a simple regression model to significantly accelerate the inference speed, approaching that of an unconditional language model. our code is available at https: //github.com/nike00811/naacl-2024-pilm . ## related work techniques that involve training a conditional language model from scratch or fine-tuning a pretrained language model-whether through rein-forcement learning @xcite , generative adversarial networks @xcite , or fit on attribute data @xcite @xcite @xcite by adjusting the model or additional parameters can produce outputs adhering to specific attributes. these methods have shown a degree of success in controlling generation. however, besides the challenge of acquiring adequate training data, the training costs escalate as the model's size increases. approaches that modify hidden representations during inference typically employ an auxiliary attribute discriminator @xcite . pplm @xcite uses a discriminator to measure whether the current latent representation can generate the text with the desired attribute. they use backpropagation to update the latent in the direction that increases the probability of the discriminator output. the weighted decoding methods @xcite @xcite @xcite only require access to the output logits or the distribution of the next token. fudge @xcite also employs a discriminator, and the difference is that the discriminator takes human-readable text as input instead of language model latent representations. the discriminator provides a score indicating the likelihood that the input text completes the document while adhering to the desired attribute. ultimately, the next token is sampled from a distribution that combines the discriminator score with the language model's output probability. the prompt-based methods @xcite @xcite concatenate the embedding before the input text, incurring almost negligible additional time cost. however, achieving more nuanced control is more challenging. discup @xcite utilizes unlikelihood training to receive soft prompts. preadd @xcite mixes the distribution of the next token generated with and without prompts, enabling more flexible control. energy-based methods @xcite have the advantage of having more relaxed conditions that only require access to the model's output text. m&m lm @xcite regards the output text as a state and utilizes bert to explore neighboring states. the decision to accept or reject a new state depends on the energy score, which can be obtained from black boxes. our proposed method amalgamates the advantages of the previously mentioned techniques while addressing their limitations. it updates latent representations to incorporate new information and employs a controller to circumvent the speed constraints associated with gradient updates during inference. existing black-box tools can be directly applied, focusing on human-readable text rather than the hidden states of language models. moreover, including longer sequences enables a more comprehensive capture of semantic information. ## method in this section, we will describe our proposed method. given an unconditional pre-trained generative model g, and prefix tokens x 1:@xmath0, x 2 , ..., x i }, and @xmath1:s denote as completed sequence, g is only learned to complete x with maximize p (x i+1 |x 1:i ), complete text with an additional attribute a can be modeled as p (x i+1 |x 1:i , a). according to bayes' theorem, p (x|a) ∝ p (x)p (a|x) we can decouple p (x|a) into unconditional language model p (x) and posterior probability of attribute p (a|x). for a controlled text generation task utilization of the language model, we can obtain the latent representation @xmath2, h 2 , ..., h i }, h i denote the key-value pair computed by the language model from a token x i . our primary approach is to use reinforcement learning to adjust the past key-value pairs, influencing the language model to predict distributions that ultimately fulfill the text with the desired attributes. we regard the pre-trained language model g and current latent representation h as the agents in the policy gradient algorithm. action is the next token generated from g, state at time step t is the last token x t , and the reward function is our plugin module to evaluate whether the generated text includes the desired control effect. to maximize expected rewards while mitigating the risk of adversely affecting pre-trained language models, we freeze the parameters in g and only update latent representation h. deviating from previous approaches @xcite @xcite , adjusting the distribution at all positions may result in excessive updates, resulting in a decline in text quality. we modify h for every n token and consider future n tokens in the update process τ = x 1:t+n , enabling a more comprehensive evaluation of h. eventually, we increase p(a|x) while remaining p(x). ## experiment we validate our proposed method on the gpt2medium model @xcite across three distinct controlled generation tasks: sentiment control, topic control, and language detoxification. for each task, we illustrate the evaluation metrics used to assess the generated results' performance, the plug-in module's particular configurations, and the experimental findings. to demonstrate the effectiveness of our approach on larger language models, in addition to the gpt2-medium model (355m), we also apply our method to the llama-7b model @xcite on sentiment control task. ## ethics statement we acknowledge the potential for controlled generation methods to be utilized to generate malicious content. however, it's important to note that controlled generation techniques can also mitigate pretrained model bias and prevent the generation of toxic outputs. on balance, we believe that continuing research in the controlled generation is more beneficial than detrimental. ## limitations pilm requires access to the complete model to update the hidden representation through gradients. this means that if a language model only allows for inference apis (e.g., gpt-3, gpt-4), it cannot be implemented with pilm. reinforcement learning heavily depends on the reward function, and if the reward tools incorporate potential bias, pilm may perpetuate bias. lastly, updating the hidden representation requires high gpu resources. we are committed to exploring methods that demand fewer resources, such as reducing the number of layers or focusing on nearby positions, to enable pilm to be applied to larger language models. method outputs gpt-2 once upon a time, i lived in the city. i loved it. i was so happy there. it felt so peaceful. i would go there every day to work, go to the movies, go to concerts. i'd go to the park. there was nothing pplm once upon a time, a while back i., a woman had a pregnancy and a man was in a lot of people, and it's not a lot of the same, the best, i think that's my brother has the best of it, there's a fudge once upon a time no (or else... or worse, i might even let bad-favoriti (i.a/i/a/i/a/something/anya /i.a / or / preadd once upon a time revelation starg exempt aspects inravings believearl capable prosecuted fromdownload 109thoroughlyeki nonethelessdallas towering animalslevant +-watkins fah protrisome solicement guess immune satisfactoryuggestabilityxml witatechacetpok 00e9monmidt cooper stocks plenty known 433studos realization bored dwind pilm-rl once upon a time, a person who was so depressed that he couldn't even function properly, and couldn't find a place in this world for himself, and couldn't even get any help from anyone, couldn't take any of it anymore, and couldn't even get pilm-controller once upon a time, i could not believe that so many men could be so evil. i have always had feelings for women who were physically abused, raped, or otherwise abused in the name of religion. but the world is so much worse than what we imagine.