Controllable Text Generation with Residual Memory Transformer (2024)

Abstract: Large-scale Causal Language Models (CLMs), e.g., GPT3 and ChatGPT, have brought great success in text generation. However, it is still an open challenge to effectively control the generation process of a CLM while balancing the flexibility, control granularity, and generation efficiency. In this paper, we provide a new alternative for controllable text generation (CTG), by designing a non-intrusive, lightweight control plugin, namely Residual Memory Transformer (RMT), to accompany the generation of CLM at arbitrary time steps. With an encoder-decoder setup, RMT can accept any types of control conditions and cooperate with the base CLM through a residual learning paradigm, to achieve a more flexible, general, and efficient CTG. Extensive experiments are carried out on various control tasks, in the form of both automatic and human evaluations. The results demonstrate the superiority of RMT over a wide range of state-of-the-art CTG approaches. The code implementation of our work is available at: https://github.com/Residual_Memory_Transformer.

## introduction controllable text generation (ctg) focuses on generating text while adhering to specific constraints @xcite . these constraints can range from high-level semantic elements, such as emotions, topics, and toxicity avoidance, to finer-grained content, e.g., inclusion of specific concepts or key elements in the generated text. recently, ctg has been regarded as critical for establishing safer, more reliable and practical @xcite @xcite ai applications in the real world. the current state-of-the-art ctg methods are based on large language models (llms) that build upon the transformer structure @xcite and have gained significant attention due to their remarkable ability to understand and generate text. recently, large-scale causal language models (clms), i.e., decoder-only language models, show particular advantages in zero/fewshot scenarios @xcite , resulting in a series of successors such as gpt3 and gpt4. this has been seen as a milestone towards the realization of artificial general intelligence. considering the significant scale of clms and the substantial cost of training such models, current mainstream clm-based ctg methods fall into two categories, i.e., prompt-based and post-processing approaches. the prompt-based methods @xcite @xcite @xcite concatenate the control-prompts or demonstrations with the input head of the generative model to instruct more controllable text generation. as previous studies have revealed @xcite , the control effectiveness tends to deteriorate with the increasing distance from the prompt. additionally, inserting a control-prompt into a well-trained model may harm the model's original generative stream, thus losing the flexibility of control @xcite . on the other hand, most postprocessing methods leverage an auxiliary module to adjust the probability of naturally producing a token by the generative model at the decoding phase @xcite @xcite , hindering the model's capacity of content planning and thus limiting the fine-grained control. more recent decodetime methods @xcite @xcite improve the control granularity through iterative sampling or editing, but at the expense of generation efficiency. therefore, the flexibility, control granularity and generation efficiency need to be better balanced, which demands a more versatile clm-oriented ctg framework foot_0 . in this paper, we propose a novel ctg plugin named residual memory transformer (rmt), which borrows the paradigm of residual learning @xcite and only makes late fusion with a frozen clm to noninvasively steer the generation process. unlike the prompt-based approaches, this paradigm does not disturb the original generative stream of the base clm model, allowing for a better flexibility of ctg (i.e., control with a plug-and-play manner). in addition, the rmt architecture consists of an encoder-decoder structure, where the encoder handles different types of control information and influences the generation process, so as to achieve fine-grained control. meanwhile, rmt utilizes cross-attention to uniformly apply control conditions to each generated token, avoiding the negative effect varying with context length. in particular, different from the vanilla decoder of transformer, an additional clm causal attention is introduced to extract the prior knowledge from the generative stream of clm, allowing rmt not to deviate too far away from the original generative model while its implied high-level semantics is leveraged. the reuse of the base-clm's output allows rmt to achieve effective control with a tiny network, resulting in an improved generation efficiency. the training of rmt includes two stages: pretraining and fine-tuning. the pre-training of rmt aims to reconstruct noisy text into a complete sentence, facilitating rmt's understanding of the semantics of various control conditions, while aligning with the generative process of clm. during the fine-tuning stage (i.e., residual learning), the logit of rmt is directly added to that of the fixed clm, and the goal is to learn the parameters of rmt such that the joint model distribution gets close to the desired text distribution. since the gradient does not need to be backpropagated to base-clm, the training of rmt is efficient. for instance, in our experiments based on gpt2-large, the entire pretraining stage of rmt with 4m samples was completed in approximately 30 hours, using a single nvidia a6000 gpu. for comparison, the computational cost of rmt is approximately 1/10 of nrp citecarlsson-etal-2022-fine, a baseline ctg model used in our experiments. we conduct extensive experiments to explore the superiority of our approach in three aspects. experimentally, it maintains the same level of control effectiveness in both with-context and withoutcontext settings, showing that rmt enables longdistance control throughout the generation process. ## related work as summarized by @xcite , ctg approaches could be divided into three categories, i.e., retraining/refactoring, fine-tuning, and postprocess. the first two categories refer to the training approaches, which aim to fine-tune (e.g., reinforcement learning is used to fine-tune pretrained language model(plm) @xcite , optimizing continuous prompts @xcite , and prompting model @xcite to steer text generation) or retrain/refactor @xcite @xcite a plm to generate texts that meet the desired control condition. these methods have shown significant performance improvements in the field. however, with the increasing size of plm, they have become resource-intensive to fine-tune or retrain. as a result, post-process approaches have become more popular in the research community. the post-process methods are dedicated to guiding language model toward desired texts in the decode-time stage using an auxiliary module. achieving the goal of generating attribute-specific texts, pplm @xcite leverages a simple attribute classifier to update the head hidden layer of lm by gradient feedback. then, diffusion-lm @xcite combines diffusion-lm and classifier-guided gradient updates to the continuous sequence of latent variables, achieving plug-andplay controllable generation. cold @xcite proposes a novel decoding framework, by combining energy-based model and gradient-based sampling. fudge @xcite by a frozen gpt2 model. neurologits @xcite incorporates the lexical constraints into the decode-time algorithm. in order to accelerate the generation process, gedi @xcite and dexpert @xcite ) train another smaller language model as generative discriminators to guide the generation from a base clm. recently, more subtle and efficient decoding approaches @xcite are also explored to achieve the better controllability and text quality. moreover, controlling multiple control elements is also expolred @xcite . most decoder-time ctg either solely intervene with the lm during the token selection phase of generation, lacking planning capabilities, or require multiple iterations, resulting in excessive generation times. rmt shares a plug-and-play trait with decoder-time approaches, while the key distinction is that we use a lightweight residual model to integrates control information and multi-level contextual streams from the clm, enabling fine-grained content planning and efficient text generation. ## conclusions we have proposed a new ctg approach, which leverages a residual model to steer the clm to generate desired text noninvasively. additionally, we propose a residual memory transformer, a novel encoder-decoder architecture, to fuse the raw contextual text, the generative stream of clm, and the control information in one shot, thus better collaborating and controlling the text generation. experimental results show that rmt exhibits a better performance in flexibility, control granularity and efficiency over a range of state-of-the-art baselines, making it a compelling solution for ctg.