Efficient (Soft) Q -Learning for Text Generation with Limited Good Data (2022)

Abstract: Maximum likelihood estimation (MLE) is the predominant algorithm for training text generation models. This paradigm relies on direct supervision examples, which is not applicable to many emerging applications, such as generating adversarial attacks or generating prompts to control language models. Reinforcement learning (RL) on the other hand offers a more flexible solution by allowing users to plug in arbitrary task metrics as reward. Yet previous RL algorithms for text generation, such as policy gradient (on-policy RL) and Q-learning (off-policy RL), are often notoriously inefficient or unstable to train due to the large sequence space and the sparse reward received only at the end of sequences. In this paper, we introduce a new RL formulation for text generation from the soft Q-learning (SQL) perspective. It enables us to draw from the latest RL advances, such as path consistency learning, to combine the best of on-/off-policy updates, and learn effectively from sparse reward. We apply the approach to a wide range of novel text generation tasks, including learning from noisy/negative examples, adversarial attacks, and prompt generation. Experiments show our approach consistently outperforms both task-specialized algorithms and the previous RL methods.

## introduction recent natural language generation systems have made remarkable progress in producing wellformed text, especially with massive pretrained language models. those models are typically trained using maximum likelihood estimation (mle) with a large amount of data supervisions. despite its successes, the standard training method suffers from limited applicability to many emerging text generation problems, where little or no supervised data is available. prominent examples of such low-data problems include generating prompts to control the massive lms @xcite @xcite , learning text generation from noisy or even negative data, generating adversarial text attacks for robustness study @xcite , and others (figure 1 , right). due to the failure of standard mle, people have had to devise specialized algorithms for those problems respectively. reinforcement learning (rl) (sutton and barto, 2018) offers an alternative principled framework for learning from arbitrary reward functions. however, rl by far has made limited success for training text generation, primarily due to the key challenges of sparse reward (i.e., a single reward signal is received only after the whole text sequence is generated) and large action space (i.e., a vocabulary of millions of words). for instance, a popular family of rl algorithms studied extensively for text generation is the policy-based @xcite or actor-critic based @xcite algorithms, with policy gradient (pg) being the most prevalent example @xcite @xcite @xcite . those algorithms train the model with on-policy updates, i.e., the text samples used for estimating policy gradients are from the target model itself. due to the exponentially large space of sequences, on-policy updates often suffer from extremely high variance and low data efficiency (e.g., most model samples are not useful for learning). thus directly training with pg from scratch is usually impossible. in practice, the model has to be initialized by mle training, followed by pg as finetuning, which often leads to limited improvement @xcite . another set of work has resorted to off-policy rl. the key advantage is that samples from other sources, e.g., human-written text, can be used, making them more data efficient than on-policy meth-... text generation is challenging due to sparse reward (i.e., the rewards of all intermediate steps are 0) and large action space (i.e., large vocabulary). our sql formulation enables several key algorithmic features as highlighted with yellow color, including (1) the combined on-and off-policy updates for the best of both, (2) bridging the final non-zero reward to directly supervise the q-value estimation at intermediate steps for learning stability, and (3) simultaneously updating the q-values of all candidate actions for efficiency. right: we explore diverse applications of the text-generation rl algorithm. in this paper, we develop a new rl formulation for text generation that tackles the above issues (figure 1 , left). we reframe the text generation problem from the soft q-learning perspective originally developed in robotics @xcite . the resulting connection allows us to seamlessly take advantage of the latest successful techniques from the rl literature. in particular, we introduce and adapt the principled path consistency learning @xcite to text generation, that (1) offers a natural way to train the model with both on-and off-policy updates, hence combining the best of the two strategies, (2) bridges the sparse reward signal to directly supervise the q function learning, leading to more accurate q estimation and credit assignment, and (3) makes efficient updates to q-values by considering all candidate actions together. the generality and efficiency of the proposed method allows us to train text generation in a wide range of applications: (1) with noisy and negative training examples, our approach learns to generate accurate entailment text that greatly improves upon the data itself as well as other various training methods; (2) our approach also manages to train an effective adversarial text generator for robustness test for classifiers; (3) we train a prompt generator with our algorithm to achieve controllable generation of pretrained lms in terms of topics. foot_1 on all the three tasks, our approach consistently improves over not only previous rl algorithms for text generation, but also diverse task-specialized methods designed specifically for each of the problems, respectively. in the appendix ( §a.1.4), we also show that on standard supervised tasks where mle prevails, our approach is competitive to train text generation models from scratch, which was usually impossible for previous rl algorithms. the contributions can be summarized as follows. on the technical side, we propose a new rl formulation for text generation based on soft q-learning. this new formulation allows us to seamlessly take advantage of the rl literature's latest successful techniques (notably the path con-sistency algorithm) to overcome the longstanding challenges (e.g., sparse reward and large action space) in text generation. on the empirical side, we conduct studies on a wide variety of text generation tasks with limited data (i.e., generating from noisy/negative data, adversarial text generation, prompt generation). we propose their rl formulations, and show that our general approach consistently improves over not only previous text rl algorithms, but also diverse task-specialized methods. ## background and challenges we aim to learn a generation model p θ (y) = t @xmath0, where y t is a token from a vocabulary v. the distribution at each step t is obtained by applying softmax on the logits f θ (y|@xmath1: following the notation convention in rl, let τ be the trajectory (i.e., text sample) generated by the policy. the agent's objective is to maximize the accumulative reward, j(π) = e τ ∼π t @xmath2, where γ is the discount factor. a central concept in rl is the q-function of policy π, q π (s t , a t ) = e π t @xmath3, a t , the expected future reward of taking action a t (i.e., generating token a t ) in state s t and continuing with the policy π. challenges. text generation poses significant challenges to rl, particularly because (1) the reward signal is usually sparse, i.e., r @xmath4, ∀@xmath5, (2) the action space (i.e., the vocabulary v) is extremely large. the challenges have led to difficulties of the two major families of rl approaches applied to text generation problems, as detailed below. policy-based rl techniques directly parameterize the policy π θ with parameters θ. thus the policy π θ (a t |s t ) exactly corresponds to the above generation model p θ (y t |@xmath6. policy gradient (pg) is one of the most widely used algorithms for text generation @xcite . it optimizes the cumulative reward with the policy gradient using the estimated q π θ value based on sample τ . pg is an on-policy algorithm, meaning that the sample τ needs to come from the the current policy π θ itself. in practice, however, optimizing this objective alone from scratch is unlikely going to work because most samples τ ∼ π θ are just gibberish with zero reward, failing to provide meaningful training signals for updating the policy. previous literature either initializes the policy π θ with mle training, and/or use a combination of mle and pg updates, which often leads to marginal gains in practice @xcite . value-based rl techniques, such as q-learning, implicitly learn the policy π by approximating the value q π (s, a) directly. deep q-learning @xcite parameterizes the q-function as q θ (x, a), and train the parameters by minimizing the following regression objective l(θ) based on the bellman temporal consistency: ## the soft q-learning framework we introduce the soft q-learning (sql) formulation of text generation. it is seamlessly compatible with the common architecture of text generation model (eq.1), permits easy implementation ( §3.1), and enables efficient and stable rl training in practice ( §3.2). figure 2 and algorithm 1 summarizes the resulting sql framework. ## applications and experiments we show broad applications of the proposed rl text generation framework to a variety of problems where no clean supervision data is available. these include learning with noisy or even negative data ( §4.1), generating adversarial text attacks ( §4.2), and generating prompts to steer pretrained lms ( §4.3). we also study the performance on standard supervised generation tasks ( §a.1.4) and show that our approach is competitive to train text generation models from scratch. we provide detailed configurations in the appendix ( §a.2). ## related work standard rl algorithms can sometimes be oversensitive to the randomness in the environment. recent works have considered maximum-entropy rl extensions, such as the soft q-learning (sql) @xcite @xcite , that maximize the entropy of policy besides the rewards, and demonstrated substantial improvement in robotic and game control @xcite @xcite . our work is the first to adapt sql and its advanced variants (in particular the path consistency learning @xcite ) to the challenging text generation problem and show significant results on diverse applications. applying rl for text generation has been discussed in alleviating the exposure bias problem and optimizing task metrics @xcite @xcite @xcite @xcite . for example, ranzato et al. (2016) used the reinforce algorithm @xcite , and @xcite used the actorcritic algorithm; @xcite and @xcite tried to relieve the sparsity problem via hierarchical and inverse rl methods, resp. they are all on-policy rl algorithms with the need of pretraining their models using mle. raml @xcite implicitly relies on the quality of off-policy data; this does not necessarily apply in our experiments with limited good data. @xcite and @xcite offer a unified view of raml, rl, and other training methods. another line of work focused mostly on using only off-policy data, often for offline training of chatbots @xcite @xcite . as a result, the opportunity of directly improving the reward (as in on-policy updates) for other rich tasks is missed. our proposed framework combines onand off-policy training, and further offers solutions for efficient training from scratch in the presence of large action space and sparse sequence-level reward in text generation. ## conclusion we develop a new rl formulation for text generation based on soft q-learning and path consistency learning. we conduct experiments on learning with noisy and negative data, black box adversarial attack, prompting a pretrained language model for controllable generation, and standard supervised tasks. this formulation opens up new opportunities to integrate more advances made in the fertile rl literature to improve text generation problems.