Reinforcement Learning with Token-level Feedback for Controllable Text Generation (2024)

Abstract: To meet the requirements of real-world applications, it is essential to control generations of large language models (LLMs). Prior research has tried to introduce reinforcement learning (RL) into controllable text generation while most existing methods suffer from overfitting issues (finetuning-based methods) or semantic collapse (post-processing methods). However, current RL methods are generally guided by coarse-grained (sentence/paragraph-level) feedback, which may lead to suboptimal performance owing to semantic twists or progressions within sentences. To tackle that, we propose a novel reinforcement learning algorithm named TOLE which formulates TOken-LEvel rewards for controllable text generation, and employs a “first-quantize-then-noise” paradigm to enhance the robustness of the RL algorithm. Furthermore, TOLE can be flexibly extended to multiple constraints with little computational expense. Experimental results show that our algorithm can achieve superior performance on both single-attribute and multi-attribute control tasks. We have released our codes at https://github.com/WindyLee0822/CTG.

## introduction large autoregressive language models (llms) trained on extensive corpus can generate highquality texts. however, to satisfy real-world applications, making the generation more controllable is urgent. it is desired to enhance specific attributes of generated texts for practical needs (e.g. positive sentiment for psychological escort, formality for academic writing) @xcite @xcite and reduce intrinsic defects of pre-trained language models (e.g. toxicity, repetition) @xcite . retraining models @xcite are subject to computational over- * *corresponding author. there is some research @xcite introducing reinforcement learning (rl) into controllable text generation (ctg) tasks. rl paradigms can relieve the above problems, which alleviate the overfitting issue by training on self-generated sentences, and can integrate parameter-efficient strategies with canonical llm backbones. however, rl-based methods generally update language models with sentence-level (or paragraph-level) rewards, leading to suboptimal performance and slow convergence. the coarsegrained rewards cannot provide clear guidance, since semantic in the sentence often transits with twists or progression. moreover, different parts of the sentence may contribute to different attributes. therefore, rl methods with coarse-grained feedback generally require considerable training steps to converge. our objective is to granularize the coarse-grained feedback to provide more precise guidance for llms. in this paper, we propose a novel reinforcement learning algorithm with token-level guidance named tole. we first provide an alternative perspective of bayesian factorization, which inspires us to formulate the token-level rewards as the probability shifts of attribute classifiers. to enhance the robustness of tole, we propose an exploration framework with "first quantize, then noise" procedure. moreover, tole can be extended to multi-attribute scenarios with few computational overheads. we conduct two experiments on single-attribute: sentiment control and detoxification. we also evaluate tole on multi-attribute scenarios with two settings. tole achieves superior performance compared with a wide range of baselines. multi-aspect controllable text generation. along with single-aspect controlling, most research on multi-aspect controllable text generation can also categorized into finetuning and postprocessing. some post-processing research @xcite in mctg combines multiple attribute discriminators to aggregate the controllability. however, they also inherit drawbacks of post-processing methods due to direct distribution regulations. finetuning-based research tries to connect several single controllers, e.g. connectors to combine multiple plugins @xcite , latent variables to represent the unsupervised aspects @xcite , direct combination of prompts @xcite , the boundary exploration of intersected subspaces @xcite . to the best of our knowledge, we are the first to explore how to extend single-attribute reinforcement learning algorithms to the mtcg scenario. token-level guidance for reinforcement learning. there is a series of research @xcite @xcite incorporating rl techniques into the transformer structure, trying to deconstruct the coarse-grained reward into the token level for sequential modeling. however, they are hard to extend to practical applications since their specialized token settings are not in line with current llms. concurrent with our research, some research @xcite on llm alignments tries to handle the problem of coarse-grained feedback. rlhf (reinforcement learning from human feedback) algorithms of the llm alignment generally require a large-scale reward model, which should be trained on datasets formatted as pairwise sentences with the same prefix. however, such data is unavailable when confronted with a wide variety of attribute requirements. therefore, exploring a novel reinforcement learning algorithm with tokenlevel feedback is significant for controllable text generation. ## approach we will first establish the notation, provide some background on existing rl methods in controllable text generation and model alignment, and offer an overview of our algorithm. ## conclusion to summarize, we propose an extensible reinforcement learning algorithm for controllable text generation with token-level feedback. we provide an alternative perspective of bayesian factorization, which enlightens our token-level reward design. we also introduce "quantization & noise" into rl to enhance the algorithm robustness. we also propose a small-scale module "weigher" to extend our algorithm to multiple constraints. extensive experiments demonstrates the effectiveness of our algorithm.