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.

text: Findings of the Association for Computational Linguistics ACL 2024, pages 1048–1066 August 11-16, 2024 ©2024 Association for Computational Linguistics
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.

heading: 1 Introduction
text: Controllable text generation (CTG) focuses on generating text while adhering to specific constraints (Hu and Li, 2021; Zhang et al., 2022). 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 (Krause et al., 2021; Liu et al., 2021a; Zhang et al., 2022) 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 (Vaswani et al., 2017) 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 (Brown et al., 2020), 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 (Zhang and Song, 2022; Yang et al., 2022; Qian et al., 2022; Lu et al., 2022a; Zhou et al., 2023) 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 (Zou et al., 2021; Carlsson et al., 2022), 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 (Carlsson et al., 2022). 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 (Liu et al., 2021a; Yang and Klein, 2021; Lu et al., 2022b; Zhong et al., 2023), hindering the model’s capacity of content planning and thus limiting the fine-grained control. More recent decodetime methods (Li et al., 2022; Mireshghallah et al., 2022; Qin et al., 2022) 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 framework1.
1We compare the key features of different CTG methods in Appendix C.
1048
In this paper, we propose a novel CTG plugin named Residual Memory Transformer (RMT), which borrows the paradigm of residual learning (He et al., 2016; Zhang et al., 2020a) 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. (1) Flexibility: Theoretically, the proposed RMT has the capability to intervene in the generation process at any step with a non-intrusive manner.
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. (2) Control Granularity: We test our approach on a range of CTG tasks of different granularity levels (i.e., fine-grained control tasks including word inclusion and sentence length control; and attribute control based on sentiment). RMT achieves a control effectiveness comparable to the state-of-the-art approaches, while guaranteeing the text quality. (3) Efficiency: The results show that three-layer blocks of RMT is enough to achieve a competitive CTG effectiveness, approving the parameterefficient traits. The time-cost of text generation is almost the same to the original CLM, significantly outperforms the baselines.

heading: 2 Methodology

heading: 2.1 CTG Framework with RMT
text: As the dashed diagram in Figure 1 illustrates, RMT operates in a non-intrusive control mode as a CLM plug-in, where the control signal passes through RMT independently and affects the output distribution of a frozen CLM through residual learning (yt = y g t + y c t ) without interfering with the CLM’s free-generation passageway. Compared to intrusive controllable paradigms, e.g., prompt-based approaches or the attentionbased CTG (e.g., CoCon (Chan et al., 2021) inserts a control block into the middle layer of GPT in an intrusive manner), RMT allows for more flexibility in switching between freestyle generation and controllable generation modes. This trait is very important in the LLM era. On the one hand, it avoids harm to the generality of LLM, keeping the powerful capacity to multi-task. On the other hand, the gradient does not need backpropagation into CLM with large-scale parameters, bringing significant training efficiency. Hence, the non-intrusive fashion for CTG is challenging, but also more promising. More explanation can be seen in Appendix C.
Formally, given a partial generated text x<t and a sequence of control instruction C, the proposed framework aims to generate eligible text Xn = {x1, ..., xn} that meets the control conditions:
PΘ (Xn) =
n∏
t=1
PΘ (xt | x<t;C) , (1)
where Θ = {θ̃;ϕ} represents the CTG model’s parameters, which comprise both the frozen param-
eters inherited from the CLM (θ̃) and the tunable parameters derived from the RMT (ϕ). The entire generation process of our approach consists of the following four steps:
CLM’s Raw Generation. At the generation step t, CLM first maps the generated text x<t = {x1, ..., xt−1} to hidden states hg<t = {hg1, ...,hgt−1}. Afterward, a linear layer and softmax normalization are used to project the hidden state of the last token hgt−1 to the probability distribution ygt over the vocabulary:
ygt = Pθ̃ (xt|x<t) = softmax(Linear(h g t−1)), (2) where the CLM naturally generates the next token xt given the context x<t as per its training. RMT’s Control Encoding. Next, the RMT’s encoder is responsible for encoding the control instruction C = {c1, ..., cm} into the control memory C = {c1, ..., cm}, which is used to guide the controllable text generation in the RMT’s decoder:
C = RMT-Enc(C;ϕenc). (3)
RMT’s Control Decoding. After encoding the control signal, the RMT’s decoder maps the generated text x<t to new hidden states hc<t = {hc1, ...,hct−1} by considering both the control memory C and CLM’s vanilla hidden states hg<t, and synthetically predicting next token’s probability distribution yct over the vocabulary:
hc<t = {hc1...,hct−1} = RMT-Dec(x<t,hg<t,C;ϕdec), yct = Pϕ (xt|x<t;C) = softmax(Linear(hct−1)).
(4) Residual Learning to Generate. We employ residual learning to fuse the output distributions from the CLM (Eq. 2) and the RMT (Eq. 4), allowing the framework to obtain the joint predictions for the next token and achieve a noninvasive CTG:
yt = PΘ (xt|x<t;C) = ygt + yct . (5)

heading: 2.2 Residual Memory Transformer (RMT)
text: The detailed structure of RMT is shown in Figure 1. Specifically, RMT adopts an encoder-decoder architecture and reuses the CLM’s suite of word and position embeddings.
RMT Encoder. The RMT encoder is responsible for encoding the control description C into the control memory C (Eq. 3). The encoder is composed of a stack of M identical blocks. Each block comprises a self-attention layer (Eq. 6) and a fully connected feed-forward layer. Additionally, it incorporates residual connections around above each layer, and is followed by a layer normalization.
RMT Decoder. The RMT decoder aims to predict the probability distribution of the next token yct (Eq. 4 ) in the control mode. The decoder is also composed of a stack of M identical blocks. Each block contains three carefully-designed attention layers and a fully connected feed-forward network. Similar to the encoder, residual connections and layer normalization are also applied. The three attention layers are directed at the already generated text x<t, the CLM’s output hidden states h g <t, and the control memory C, respectively: Causal Self Attention. The first attention layer utilizes a causal attention operation (Eq. 7 in Appendix A), whereby Q, K, and V are all firstly mapped from the original generated text x<t itself. This attention mechanism facilitates the identification and capturing contextual features of generated sequence from scratch.
Causal CLM Attention. The second attention layer also employs causal attention, but with a key difference: Q is sourced from the previous causal self-attention layer’s output, while K and V are obtained from the CLM’s last hidden states hg<t. This design establishes an inner residual connection with CLM, enabling RMT to consider the highlevel contextual features and maximally interact with the generative stream of CLM.
Cross Control Attention. The third attention layer is a cross-attention (Refer to the Appendix A) for the control memory from the RMT encoder. Specifically, Q is sourced from the previous causal CLM attention layer, while K and V are derived from the control memory C. This cross-attention layer bridges the RMT’s encoder and decoder, introducing the control signals to the generation. The cross-attention allows that every hidden vector of the current token can equally query the control instruction sequence encoded by RMT’s encoder, avoiding the negative effect of the control instruction’s token position.

heading: 2.3 Model Training
text: Pre-training. To enable the RMT’s semantic understanding capability and allow it align with the
CLM’s generation process, we utilize the denoising auto-encoder pre-training method, whereby the encoder processes the corrupted text X̂ regarded as control condition, and the decoder reconstructs the original text X .
The denoising pre-training schematic diagram could be seen in Figure2. Specifically, we corrupt the pretraining text (X → X̂) in four ways referring to (Lewis et al., 2020): (1) Token Masking: randomly masking tokens with special tokens. (2) Token Deletion: randomly deleting tokens in the text. (3) Span Replacing: randomly replacing the text spans of different lengths with special tokens. Different from token masking, each span is replaced with a single special token. (4) Text Rotation: randomly selecting a token and rotating the text around it, i.e., the text begins with that token and ends with its previous token. More details are supplied in Appendix B.1. Once the pre-training is complete, RMT can be fine-tuned to perform various controllable text generation tasks, without any need for re-pretraining.
Fine-tuning. This allows the CLM to generate sequences that meet specific control requirements in an auto-regressive manner. The specific objective can vary according to different tasks: we use the Maximum Likelihood Estimation (MLE) objective (please see details in Appendix A.3) for the word inclusion and length control tasks; and use the unlikelihood training strategy (Zhang and Song, 2022) for the attribute control task.
It is worth noting that RMT shares the trait of training efficiency. Firstly, RMT efficiently reuses the output of the base CLM, which makes RMT requiring fewer pre-training datasets and expediting the pre-training process. Additionally, RMT is lightweight and is built on the top layer of the
base CLM. Therefore, gradients do not need to propagate into the base CLM during the backpropagation process, which can save a significant amount of training time and GPU memory.

heading: 3 Experiments

heading: 3.1 Word Inclusion Experiment
text: Two different settings are experimented in this subsection. The first one is the without-context setting, i.e., steering the CLM to generate a single sentence containing the keywords without the context. The instruction is as follows:Continue to write a sentence, and include [word1, word2, ..., wordn]. The second is the with-context generation setting, which requires the CLM to continue to generate target text under an extended context. The instruction is:Continue to write a sentence, and include [word1, word2, ..., wordn]. Context: [context text]. The consideration of these two modes aims to test the CTG’s flexibility, capable of controlling the generation of CLM in a position-independent way.
Experimental Setting. Following the NRP (Carlsson et al., 2022), we use the GPT2large as the backbone CLM, and the training data for pre-training and fine-tuning also comes from Wikipedia. More details on training and inference are provided in Appendix B.2. We test our approach on CommonGen (Lin et al., 2020) (for without-context setting), and C2Gen (Carlsson et al., 2022)(for with-context setting).
Baselines. As for the without-context setting, we firstly choose those methods that are specifically trained for lexical constraints tasks. They are KG-BART (Liu et al., 2021b), which fine-tunes the whole parameters of BART while it is augmented by knowledge graph; and POINT (Zhang et al., 2020b), which is an insertion-based method and iteratively injects words around the target words. Finally, we also employ a pure decode-time approaches for lexical constraint, namely NeuroLogic (Lu et al., 2022b), as a baselines. The general CTG methods are also compared, including Non-Residual Prompt (Carlsson et al., 2022)(i.e., a recent state-of-the-art method using positionindependent prompting model to guide the generation of GPT2-large), and approaches directly instructing the GPT2-large and ChatGPT to generate target sentences. As for the with-context setting, KG-BART and POINT could not be applied to this scene and thus they are removed. The test results of POINT, KG-BART, and NRP are adopted from
the open source code3. Evaluation. Following the NRP (Carlsson et al., 2022), We employ coverage (Cov), measuring the average coverage rate of the generated text with respect to the target words, to assess the controllability. Perplexity (PPL) is used to evaluate the fluency. We calculate PPL using an off-the-shelf GPT2-large model. Additionally, we utilize SelfBleu-5 to evaluate the diversity of the generated text, with lower Self-Bleu values indicating greater syntactic diversity. Considering the non-intrusive setting (i.e., the CLM is fixed and the task is conducted on the open-ended text generation setting), our experiments omit that task-specific metrics like BLEU, METEOR, CIDEr, and SPICE for CommonGen dataset. Additionally, we conduct a human evaluation to assess the conformity of the generated text to common sense (CS) and text Fluency from human perspective. For the with-context setting, we introduce an additional manual evaluation metric called Relevance (Rel), where human evaluators scored the relevance of the generated text to the given context. More detailed information can be found in the Appendix D.
Result and Analysis. As is shown in Table 1, in the without-context setting, the coverage of RMT significantly outperforms the vanilla prompt method (i.e., Prompt+GPT2) and pure decode-time approach (i.e, NeuroLogit) with the same decode algorithm using by RMT. It shows slight advantages over ChatGPT and NRP. Compared to task-specific models (i.e., POINT, KG-BART), the coverage of RMT is close to them, yet with lower PPL. It is worth noting that, in the with-context setting, the strong baselines, including ChatGPT and NRP, respectively suffer from a 12.0% and 8.0% decline of control ability compared to the no-context setting. However, RMT maintains the same-level of control ability with the no-context setting, and outperforms NRP and ChatGPT, which shows the superiority of RMT in that the control ability will not degrade with the context length.
In term of human evaluation in CommonGen, POINT is an insertion-based generation framework that naturally suffered from worse text quality; thus the result of RMT is much better than POINT on terms of commonsense and fluency. Like NRP and Promp+GPT2, we all use frozen GPT-2-large as the backbone and do not introduce external knowledge.
3https://github.com/FreddeFrallan/ Non-Residual-Prompting
RMT achieves comparable results, which prove that the non-intrusive paradigm with residual learning would maintain the language model’s ability, thereby guaranteeing the text quality. Our approach falls behind KG-BART enhanced by knowledge graph and ChatGPT with hundred-billion level parameters. This is within exception. If inserting RMT into larger-scale CLM, e.g., GPT3, we believe the text quality of our approach will also be improved. As for C2Gen, RMT produces comparable results in contextual relevance with ChatGPT, but there is still a gap in common sense. More qualitative examples of different CTG approaches can be seen in Appendix E.

heading: 3.2 Sentence Length Control Experiment
text: We test the length-control ability of RMT, i.e., instructing the CLM to generate a sentence which satisfies the word inclusion objective under an exact number of words. The instruction is as follows:Continue to write a sentence, and include [word1, word2, ..., wordn], the sentence length should be within CL words. This setting increases the challenge of word inclusion task, effectively approving the fine-grained control ability.
We report control performance for CL = 15, 18, 20. As shown in Table 1, RMT can maintain same level word coverage under different control-
length. The longer the target length, the lower the generated text’s PPL. This result is quite intuitive since it is harder to produce fluent text within the shorter required length.
To quantitatively measure the controllability of sentence length, we conduct tests on the CommonGen validation dataset to analyze the control effect of our proposed method. Some examples are presented in Table 2, which demonstrate that RMT effectively steers the CLM to generate texts with specific keywords and required sentence lengths. Furthermore, we calculated the mean offset (Mean) and the standard deviation (SD) for different con-
trol lengths. The results are visualized in Figure 3. We observe that the offset and standard deviation both remain within one, indicating that RMT successfully achieves the sentence generation with desired sentence lengths as instructed.

heading: 3.3 Attribute Control Experiment
text: Setting and Baselines. We follow DisCup (Zhang and Song, 2022), and use a discriminator-involved training strategy to fine-tune RMT, and the checkpoint of attribute-classifier comes from the open source code4. More details could be seen in Appendix B.3. Different from DisCup, which optimizes unique prompt for every attribute, we use the RMT module to directly encode different attribute instructions uniformly. The instructions for sentiment control are “The sentiment: positive” and “The sentiment: negative”, respectively. GPT2-large is used as the backbone CLM, and the training data is the widely used Stanford Sentiment Tree (SST-5) (Socher et al., 2013) collected from movie reviews. We follow the commonly used setting in previous work (Krause et al., 2021; Zhang and Song, 2022; Lu et al., 2022a). Specifically we use 5K neutral prompts, 2.5K positive prompts, and 2.5K negative prompts as test dataset (provided by DEXPERT (Liu et al., 2021a)), and every CTG approach generates another 20 continuations based on given prompts, to achieve sentiment (i.e., positive and negative) controllable text generation. We collect the primary baselines reported by DisCup (Zhang and Song, 2022) and DEXPERT (Liu et al., 2021a), including the decodetime approaches with GPT2-large as base-CLM
4https://github.com/littlehacker26/ Discriminator-Cooperative-Unlikelihood-Prompt-Tuning
(i.e., PPLM (Dathathri et al., 2020),GEDI (Krause et al., 2021),and DEXPERT (Liu et al., 2021a)), and the training approaches (i.e., CTRL (Keskar et al., 2019) and DisCup (Zhang and Song, 2022)).
Evaluation. Following previous work (Krause et al., 2021; Zhang and Song, 2022; Lu et al., 2022a), we use an external sentiment classifier provided by Huggingface.co to classify the generated texts, and get sentiment control accuracy (i.e., Correctness). PPL and Dist-1/2/3 are reported to test their fluency and diversity, respectively. For human evaluation, we introduce Relevance, i.e., how the text conforms to required sentiment; Topicality, i.e., how the generated continuations are contextconsistent with the given prompts, and Fluency, i.e., the text’s fluency evaluated from the human perspective. More details can be seen in Appendix D.
Result and Analysis. As shown in Table 3, RMT in top-k setting demonstrates a superior control performance compared to all baselines, while maintaining comparable text quality and diversity. And RMT in the top-p setting shows better text quality yet weaker control performance. The closest performers to our approach are DEXPERT and DisCup. However, DEXPERT requires fine-tuning an external pair of CLM, resulting in tuning two time more parameters of the base CLM. In contrast, RMT is parameter-efficient, requiring tuning relevant parameters which are only around 16% of the base CLM’s parameters. Regarding DisCup, both approaches employ the same loss function. However, DisCup optimizes different continuous prompts for each attribute, while we utilize the RMT to steer the CLM using residual learning. DisCup requires fine-tuning fewer parameters, while RMT is completely decoupled from the CLM and controls different attributes using a unified model,

heading: Method Relevance (↑) Fluency (↑) Topicality (↑)
text: providing a greater flexibility and a more powerful control capability.
The human evaluation presented in Table 4 indicates that DEXPERT performs slightly better in terms of fluency and contextual relevance, but exhibits lower control ability, and RMT excels in attribute relevance. We speculate that the use of a classifier-based objective in RMT and DisCup might result in overly strong control, leading to the generation of abrupt turning points that may compromise fluency and topicality to some extent. DisCup performs relatively poorly, as too few control parameters limit its expression ability. More detailed examples are presented in Appendix E.

heading: 3.4 Further Analysis
text: Block Layers and different base-CLMs. we investigate the influence of the number of block layers on the RMT’s control performance over different base-CLMs, i.e., GPT2, GPT-j-6b (Wang and Komatsuzaki, 2021) and Llama2-7b (Touvron et al., 2023). The results shown in Figure 4 indicate that a value of M = 3 is deemed appropriate for achieving effective word inclusion and sentence length control. Moreover, RMT achieves a consistently high and stable control performance for different CLMs, showing its model-agnostic nature
and adaptability across different CLMs. Ablation Study. To assess the effectiveness of the Residual Memory Transformer (RMT), we conduct an ablation study. The study involves three settings: removing the causal self-attention, excluding the causal CLM attention, and evaluating RMT’s performance without pre-training. The results are presented in Table 6. Interestingly, under all three settings, we observe a significant decline in performance. This finding serves as compelling evidence that every component in RMT, as well as the pre-training stage, is crucial and indispensable for achieving optimal results.
Generation Efficiency. RMT is a lightweight plugin, allowing for efficient inference speeds comparable to the original CLM. Table 5 demonstrates that RMT outperforms typical CTG approaches in term of efficient generation speed, approaching the speed of the pure CLM (GPT2-large).

heading: 4 Related Work
text: As summarized by (Zhang et al., 2022), 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) (Ouyang et al., 2022; Lu et al., 2022a), optimizing continuous prompts (Yang et al., 2022; Zhang and Song, 2022), and prompting model (Carlsson et al., 2022) to steer text generation) or retrain/refactor (Chan et al., 2021; Yu et al., 2021; Keskar et al., 2019; Gururangan et al., 2020) 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 (Dathathri et al., 2020) leverages a simple attribute classifier to update the head hidden layer of LM by gradient feedback. Then, DiffusionLM (Li et al., 2022) combines Diffusion-LM and classifier-guided gradient updates to the continuous sequence of latent variables, achieving plug-andplay controllable generation. COLD (Qin et al., 2022) proposes a novel decoding framework, by combining energy-based model and gradient-based sampling. Fudge (Yang and Klein, 2021) uses the discriminator that contains future attribute information to re-rank the token distribution produced

heading: Models Cov (%)
text: by a frozen GPT2 model. NeuroLogits (Lu et al., 2022b, 2021) incorporates the lexical constraints into the decode-time algorithm. In order to accelerate the generation process, GeDi (Krause et al., 2021) and DEXPERT (Liu et al., 2021a) train another smaller language model as generative discriminators to guide the generation from a base CLM. Recently, more subtle and efficient decoding approaches (Zhong et al., 2023; Deng and Raffel, 2023) are also explored to achieve the better controllability and text quality. Moreover, controlling multiple control elements is also expolred (Kumar et al., 2021; Yang et al., 2022).
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.

heading: 5 Conclusions
text: 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.

heading: Acknowledgments
text: This work is funded in part by the Natural Science Foundation of China (grant no: 62376027) and Beijing Municipal Natural Science Foundation (grant no: 4222036 and IS23061).

heading: Limitations
text: RMT still suffers from the following limitations. (1) Challenges in applying RMT to close-ended CLMs. Presently, the application of RMT needs to obtain the last hidden states or the logits of CLM, thus directly applying RMT to certain commercial CLMs, e,g., GPT-4, is a problem that needs further investigation. This is indeed a common problem of all plugin-style CTG approaches. (2) RMT does not focus on commonsense and may result in the generation of texts that do not conform to commonsense. This issue could be relieved in the future by incorporating some external knowledge graphs.
Up to now, it is still challenging for large-scale causal language models to avoid the factual errors in the generated text, which RMT does not address either. A promising future work is to combine RMT and information retrieval techniques to enhance the factual accuracy of those generative models. Moreover, RMT could also be used to encode personal profiles and build personalized chatbots, or fusion with image information so as to be applied to multimodal scenes.

heading: Appendices
text: In order to enhance the reproducibility of the experimental results presented in our paper and provide additional supporting details for the conclusions outlined on the main page, we have included supplementary materials in this section.

heading: A Preliminary

heading: A.1 Attention in Transformer
text: Self-Attention mechanism is a critical component in Transformer (Vaswani et al., 2017), which is used to effectively capture long-range dependencies between words in the input sequence. Specifically, it is defined as a mapping between a query and a collection of key-value pairs. The values are weighted according to a compatibility function between the query and each corresponding key, and then summed up, eventually obtaining the output vector. The self-attention mechanism can be formatted as follows:
Att(Q,K, V ) = softmax ( QKT√
dk
) V, (6)
where Q,K, V represent the transformed representations of the sequence using the corresponding learnable matrix, and √ dk is the scaling factor. In the Transformer model, the key, query and value vectors are always derived from a shared sequence.
Causal Attention is a particular branch of selfattention, also called masked self-attention, which is usually applied in the decoder module of Transformer. Different from normal self-attention, the queries in causal attention are only confined to their preceding key-value pairs’ positions and their current position to maintain the auto-regressive property. Usually, it can be implemented by a mechanism that masks the invalid positions and sets them to negative infinite:
Att(Q,K, V ) = softmax ( QAKT√
dk
) V, (7)
Aij = { 1 if i ≥ j, −∞ else. (8)
Cross Attention is another branch of selfattention in the decoder part of Transformer, which aims to capture the interaction between the encoder and decoder. Specifically, The key and value parts are obtained from the previous step outputs of the
encoder, and the query is from the decoder, enabling the encoder to attend to each position in the decoder module at the same time.

heading: A.2 Causal Language Model
text: Causal Language Model (CLM) eliminates the need for a separate encoder component and typically leverages a decoder-only Transformer (Vaswani et al., 2017), in which only the causal attention is used to perform step-wise density estimation, i.e., predicting the next token. Suppose a CLM is parameterized with θ. Given a partial sequence x<t, it assigns a probability Pθ(xt|x<t) over a vocabulary V for next-token xt generation. When generating a sequence of text Xn = {x1, x2, . . . , xn}, it can be formulated by the chains rule as below:
Pθ (Xn) =
n∏
t=1
Pθ (xt | x<t) . (9)
The entire generation process is carried out iteratively. First, a token is sampled from the distribution Pθ (xt | x<t). Then the selected token is concatenated with the input for the next step of generation.

heading: A.3 Model Training with MLE
text: An encoder-decoder generation model is usually trained with Maximum Likelihood Estimation (MLE). Given a finite set of training data D consisting of a target and source data per sample, and a language model with parameters θ, the optimized object is defined as follows:
L (θ,D) = − |D|∑
i=1
|x(i)|∑
t=1
log pθ ( x (i) t | x (i) <t , S (i)),
(10) where |D|, |x(i)| represents the number of samples in dataset and the length of a target sequence respectively. S(i) represents the corresponding source data encoded by encoder of RMT. x(i)t is the nexttoken of x(i)<t, and x (i) 1:t is a partial sequence truncated from target training sample. The parameters θ is updated by maximizing the log likelihood (i.e., probability of decoder’s next-token prediction).

heading: B Experimental Details
text: In our paper, all the experiments are conducted on a single NVIDIA A6000 GPU with 48GB of memory. We utilize the PyTorch deep learning
framework as our implementation tool and employ the HuggingFace Transformers package to load a pre-trained GPT model. To implement the Residual Memory Transformer (RMT), we utilize the Transformer package in PyTorch-1.14.0, and the number of multi-head attention is set to 8. The post-normalization is applied to achieve a better performance of fine-tuning.
In the implementation of RMT for Llama2 and GPT-6-J, we replaced the Causal-Self-Attention and Causal-CLM-Attention sublayer with LlamaFlashAttention to add the position information, and Cross-Control-Attention are same to that of GPT2. As for the encoder of RMT, we also add the Rotary Position Embedding to the bidirectional self-attention.

heading: B.1 Pre-training
text: Data Details. The pre-training stage is consistently applied across all experimental settings in our study. For this stage, we collect a dataset of 4 million sentences sampled from the Wikipedia corpus to serve as the training data. To optimize computational resources, we filtered out sentences that exceeded a length of 200 words. These samples are then noised using the following methods:
(1) Token Masking: We randomly mask tokens with a special token. The proportion of masked tokens in a sample is 20%.
(2) Token Deletion: We randomly delete tokens from the text. The proportion of deleted tokens in a sample is 20%.
(3) Span Replacing: Text spans of various lengths are randomly replaced with special tokens. Each span is replaced with a single special token. The proportion of replaced span tokens in a sample is 20%.
(4) Text Rotation: A token is randomly selected, and the text is rotated around it. This means the text began with the selected token and ended with its previous token.
In the total dataset, the ratios of Token Masking, Token Deletion, Span Replacing, and Text Rotation are 50%, 16.7%, 16.7%, and 16.7%, respectively. Since RMT shares GPT’s tokenizer, which does not reserve a special token, we select the word “xxx” with no special meaning as the special token. we did not specifically tune the hyperparameters related to the ratios of different denoised tasks during
pretraining. Instead, we followed BART’s [12] approach and selected the most important tasks, such as Text Infilling (included by Token Masking in RMT), to constitute the largest proportion of the task. The remaining tasks were divided equally among the other three tasks. In our experiments, we observed that the specific choice of pretraining tasks did not significantly impact the results in CTG. For instance, when we used Token Masking as the only pretraining task, the control performance on CommonGen reached 92.2%. By adding the other three tasks, the performance improved slightly to 93.9%.
Training Details. During the pre-training process, we set the following parameters for model training: the batch size of 64, learning rate of 5e5, random seed of 42, and pre-training epoch of 1. We employ the AdamW as the optimizer; the pre-training of the RMT is conducted on a single NVIDIA A6000 GPU. Since the backpropagation process does not need to propagate into the Casual Language Model (CLM), the training of the model was efficient. As a result, the entire pre-training stage with GPT2-large is completed in approximately 30 hours.

heading: B.2 Word Inclusion
text: Baselines. We closely follow the methodology of previous work, specifically the Non-Residual Prompt (NRP) (Carlsson et al., 2022). To establish a reliable reference point for the CommonGen dataset, we rely on the generation results of the baselines provided in the NRP paper. And for the NeuroLogic baseline, in order to keep the comparability, we use the same decoding algorithm setting with RMT, the beam search size is 4 and max generation length is set to 32. For the C2Gen dataset, we utilize the checkpoint provided by NRP, to generate the responding texts. However, the result shows NRP could not work well. As a result, we give up using NRP as a baseline in the C2Gen dataset.
Regarding the ChatGPT baseline, we employ prompts specifically designed for generating results on the CommonGen dataset. The specific prompt is as follows: Continue to write a sentence, and include [word1, word2, ..., wordn], the sentence length should be within 15 words. When it comes to the C2Gen dataset, the prompt instruction of ChatGPT is: Continue to write a sentence, and include [word1, word2, ..., wordn], the sentence length should be
within 15 words. Context: [context text]. Details of RMT. In our approach, we randomly select training data from the Wikipedia corpus and fuse the training data of CommonGen. Specifically, for the setting involving without-context word inclusion and sentence length control, we gather a total of 158,343 samples. To ensure manageable sentence lengths, we impose a limit of 32 tokens per sentence. To choose the appropriate checkpoint during training, we utilize the CommonGen validation dataset and selected the checkpoint based on coverage metrics. For the with-context setting, we collect 132,170 samples, each with a sentence length constraint of 128 tokens, and constructed a validation dataset to choose the checkpoint.
Every sample is composed of two parts: context and target. We use spacy library5 to extract the keywords in the target text, of which are tagged with VERB, NOUN. The number of keywords in each sentence is restricted to the range from 3 to 5. The context length, target length, and keywords are used as control instructions to steer the generation process of the CLM. All above training data will be released and open source, once we have cleaned up our project.
During the fine-tuning stage, the learning rate of the AdamW optimizer is set to 1e-4, and the max epoch and batch size are 6 and 64, respectively. During the inference stage, For the control task without context (i.e., CommonGen dataset), we use the specific word “The” as context; the core parameters of the decoding algorithm are shown in Table 7.

heading: B.3 Sentiment Control
text: Baselines. In this part, we build upon the baselines established by DisCup (Zhang and Song, 2022) and DEXPERT (Liu et al., 2021a), adopting their settings, checkpoints, and generated results. The results of baselines, including PPLM (Dathathri et al.,
5https://spacy.io/
2020), CTRL (Keskar et al., 2019), GEDI (Krause et al., 2021), and DEXPERT, are calculated from the generated texts provided DEXPERT (Liu et al., 2021a). As for DisCup, we use the checkpoint from the authors, the depth of re-ranked candidate token is set to top-k=90 to generate continuations using the prompts in test datasets.
Details of RMT. In our own approach, during the training stage, we incorporate the discriminatorinvolved loss introduced by DisCup and utilize the attribute classifier provided by their open-source code. During the training, we optimize RMT using two different strategies, one is to use top-k=90 algorithm to select re-rank candidate tokens and the other setting uses top-p=0.95 to select candidate tokens. The top-k setting increases the depth of token sampling, thus the diversity and control ability of the model will increase, yet with the burden of decreasing text fluency. We choose a low temperature value of 0.005 and a learning rate of 2e-4. The top-p keeps the re-ranked tokens within the distribution of base-CLM, which could potentially achieve higher fluency but lower text quality and control ability. The instructions for sentiment control are “The sentiment: positive” and “The sentiment: negative”, respectively. During the inference stage, we employ the same generation parameters as those used in the word inclusion experiments described in Table 7.

heading: C CTG Approaches Comparison

heading: C.1 CTG Comparison
text: We select some notable works in controllable text generation (CTG) and compare their characteristics. The results are shown in Table 8. We compare them according to three key features: (1) Fine-grained Control suggests whether the method can potentially be applied to control various elements, such as emotion, topic, toxicity avoidance, word constraints, and structural control (e.g., personal profile, syntax tree). It indicates the method’s ability to handle different control aspects effectively. (2) Efficient Generation evaluates whether the time-cost of text generation using the method is comparable to that of the original language model (LM). (3) Flexibility mainly considers two aspects. The first is the non-intrusive fashion, i.e, the method can operate without disturbing the language model’s original textual stream, mitigating the harms of original LMs. The second requires that controleffect is distance-independent and enables long-
distance control throughout the generation process. Flexibility is a very important trait in the LLM era. Specifically, the non-intrusive fashion of RMT decouples the control module from the base-CLM and the gradient does not need backpropagation into CLM with large-scale parameters, bringing significant training efficiency. Secondly, we no longer need to fine-tune the base CLM for specific control tasks, avoiding any adverse impact on its generalization ability (i.e., keeping the CLM from becoming task-specific model). Moreover, this design enables us to dynamically adjust the fine-grained control requirements during the text generation process at any time. For instance, we can enhance a well-trained CLM, such as ChatGPT, with external grounded documents, making it possible to determine whether external control information should be intervened through an automatic selection model (i.e., switching b/w controlled and free-style generation).
The above analysis suggests that controlling the generation process of LMs while balancing control flexibility, control granularity, and generating efficiency remains an open challenge. RMT shows potential in addressing these issues, making it a compelling solution for controllable text generation.

heading: C.2 RMT v.s NRP
text: The most similar to RMT is Non-Residual Prompt (NRP) approach(Carlsson et al., 2022), which also requires pretraining and fine-tuning. Therefore, we make a thorough discussion of differences of RMT w.r.t NRP.
Principle reasonability. The principle of RMT appears intuitively more reasonable. RMT utilizes cross-attention to uniformly apply control conditions to each generated token, alleviating the “distance bias” problem that arises with increasing distance from the control prompts. Intuitively, NRP still uses distance-aware causal attention to encode control conditions, and seems inherently difficult to fully overcome the “distance bias” issue. In contrast, cross-attention in RMT enables the final representation of the current generation step to look up the control condition position-independently, which is more intuitively appealing.
Training Paradigm. The pretraining stage of RMT is general and task-independent, relying solely on self-supervised training (e.g., utilizing Wikipedia data). Various CTG tasks can be finetuned using task-specific corpora without the need
for additional data reconstruction. However, pretraining and finetuning stage of RMT is coupled and interrelated, limiting the applicability to various CTG tasks.
Training efficiency. RMT holds a clear advantage over NRP in computational source cost. RMT efficiently reuses the output of the base CLM, requiring fewer pretraining training data and expediting the pretraining process. Additionally, RMT is lightweight and is built on the top layer of the base-CLM. Hence, the backpropagation process does not need to propagate into the CLM, which will save a huge training time and GPU memory. Conversely, NRP, despite being initialized from the base CLM, necessitates training a separate baseCLM and involves multi-stage pretraining, leading to higher computational costs. As discussed in Appendix B.1, RMT’s pre-training was conducted on a single NVIDIA A6000 GPU, completed in approximately 30 hours. In contrast, according to the report in the paper on NRP, they incurred approximately 400 GPU hours (using DGX-100 machines) for their training process.
Parameter efficiency. RMT is more parameterefficient compared to NRP. NRP requires training additional base-CLM, whereas RMT only necessitates approximately 1/6 of the base-CLM’s parameters.

heading: D Human Evaluation
text: We conduct the human evaluation for the word inclusion and sentiment control tasks. To ensure the reliability of the evaluation, we invite three highly educated experts as evaluators, and each evaluator
has expertise in English. For the evaluation process, we establish a unified criterion using a scale of 10 for each aspect. The scale ranges from 1 to 3 for poor performance, 4 to 7 for satisfactory performance, and 8 to 10 for excellent performance. This criterion allows evaluators to provide consistent and comparable scores for each aspect.
Regarding the CommonGen dataset and C2Gen, we randomly select 25 samples from test dataset for each CTG approach, resulting in a total of 25*7 =175 samples for CommonGen and 50 samples for C2Gen. For the sentiment control experiments, we create 10 positive-steering prompts composed of neutral and positive prompts, as well as 10 negativesteering prompts composed of neutral and negative prompts. Each prompt is evaluated using three CTG approaches, resulting in a total of 60 samples for the final human evaluation. To obtain the final results, we calculate the average score for each metric for every CTG method under comparison, based on the evaluations provided by the experts.

heading: E Qualitative Examples
text: In this section, we list some examples generated by RMT and other different CTG approaches. Specifically, Table 9 shows the cases on the CommonGen dataset, and Table 10 presents the cases in the C2Gen dataset. As for attribute control, we show the examples on the adversarial steering setting (i.e., the polarity of the prompt’s sentiment is opposite from the target sentiment) as shown in table 11.

heading: Future Work
text: This issue could be relieved in the future by incorporating some external knowledge graphs. Up to now, it is still challenging for large-scale causal language models to avoid the factual errors in the generated text, which RMT does not address either. A promising future work is to combine RMT and information retrieval techniques to enhance the factual accuracy of those generative models. Moreover, RMT could also be used to encode personal profiles and build personalized chatbots, or fusion with image information so as to be applied to multimodal scenes