Text Generation with Efficient (Soft) $Q$-Learning

1 INTRODUCTION . Recent natural language generation systems have made remarkable progress in producing wellformed coherent text , especially with the massive pretrained language models ( LMs ) ( Radford et al. , 2019 ; Brown et al. , 2020 ; Lewis et al. , 2020 ; Raffel et al. , 2019 ) . Those models are typically trained using maximum likelihood estimation ( MLE ) with a large amount of data supervisions . Despite its successful outcomes , 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 ( Yin et al. , 2019 ; Shin et al. , 2020 ; Zhong et al. , 2021 ) , learning text generation from noisy or even negative data , generating adversarial text attacks for robustness study ( Wallace et al. , 2019 ; Atanasova et al. , 2020 ) , and others ( Figure 1 , right ) . Due to the failure of standard MLE , people have had to devise specialized algorithms for those problems respectively . On the other hand , reinforcement learning ( RL ) ( Sutton & Barto , 2018 ) offers an alternative principled framework for learning from arbitrary reward functions , and has achieved great advances in robotic and game control . 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 ( Williams , 1992 ) or actor-critic based ( Bahdanau et al. , 2016 ; Rennie et al. , 2017 ) algorithms , with policy gradient ( PG ) being the most prevalent example ( Ranzato et al. , 2016 ; Li et al. , 2016 ; Rennie et al. , 2017 ; Tan et al. , 2018 ; Pasunuru & Bansal , 2018 ; Paulus et al. , 2018 ) . 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 , onpolicy 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 ( Choshen et al. , 2020 ; Wu et al. , 2018 ) . 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 methods . Previous work has used either importance weighted PG ( Pang & He , 2021 ; Zhou et al. , 2017 ; Kandasamy et al. , 2017 ) or Q-learning based algorithms ( Guo , 2015 ; Jaques et al. , 2020 ; Narasimhan et al. , 2015 ) . However , off-policy methods have been considered to be less stable . For example , theQ-learning performance relies heavily on how accurate the learnedQ-function assesses the quality of intermediate subsequences – a challenging task due to the sparse reward signals . 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 ( Haarnoja et al. , 2017 ; Schulman et al. , 2017 ) . 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 ( Nachum et al. , 2017 ) 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 . 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 . 2 BACKGROUND AND CHALLENGES . The goal of text generation is to produce coherent text y = ( y0 , ... , yT ) of certain properties for a given task , where yt is a token from a vocabulary V , and T is the text length . The generation can condition on arbitrary input context , which we omit for simplicity of notations . We aim to learn a generation model pθ ( y ) which is typically decomposed autoregressively as pθ ( y ) = ∏T t=0 pθ ( yt | y < t ) , where y < t = ( y0 , ... , yt−1 ) is the prefix , and the distribution at each step t is obtained by applying the softmax function on the output logits : pθ ( yt | y < t ) = exp fθ ( yt | y < t ) ∑ y′∈V exp fθ ( y ′ | y < t ) . ( 1 ) Here fθ ( y | y < t ) is the logit of token y computed by the generation model . Given a training example y∗ , maximum likelihood training ( MLE ) updates the model with the gradient ∇θLMLE ( θ ) = ∑T t=0∇θ log pθ ( y∗t | y∗ < t ) . Despite its popularity , MLE-based training only applies when clean supervised data y∗ is available , and can not be used to optimize arbitrary task metrics ( e.g. , BLEU , entailment score ) which are typically the goal in many text generation tasks . 2.1 REINFORCEMENT LEARNING ( RL ) FORMULATIONS FOR TEXT GENERATION . Notations . Previous research has formulated text generation as an RL problem by considering the following finite-time Markov Decision Process ( MDP ) . At each time step t , let the “ state ” be st = y < t , namely the partial sequence generated so far . The model , also known as the “ agent ” , takes as input the current state st and outputs a token , also called “ action ” , at ∈ V according to a policy π ( at | st ) . The agent then receives a reward rt = r ( st , at ) and deterministically transitions to next state st+1 ( i.e. , the concatenation of the tokens in st and the new token at ) . 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 t=0 γ trt ] , where γ ∈ ( 0 , 1 ] is the discount factor . A central concept in RL is the Q-function of policy π , defined as Qπ ( st , at ) = Eπ [ ∑T t′=t γ t′rt′ | st , at ] , which is the expected future reward of taking action at ( i.e. , generating token at ) in state st and continuing with the policy π . Challenges . Text generation poses significant challenges to RL , particularly because ( 1 ) the reward signal is usually sparse , i.e. , rt = 0 , ∀t < T and the agent receives a non-zero reward rT only after it generates the full sequence , ( 2 ) the action space ( i.e. , the vocabulary V ) is extremely large , often containing millions of words . 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 πθ ( at | st ) exactly corresponds to the above generation model pθ ( yt | y < t ) . Policy gradient ( PG ) is one of the most widely used algorithms for text generation ( Ranzato et al. , 2016 ) . It optimizes the cumulative reward with the policy gradient : ∇θJ ( πθ ) = −Eτ∼πθ [ ∑T t=0 Q̂ ( st , at ) ∇θ log πθ ( at | st ) ] , ( 2 ) where Q̂ ( st , at ) is the estimated Qπθ value with sample τ . Notice that the expectation is taken w.r.t . the policy πθ , which makes PG 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 ( Wu et al. , 2018 ; Choshen et al. , 2020 ) . Value-based RL techniques , such as Q-learning , implicitly learn the policy π by approximating the value Qπ ( s , a ) directly . Specifically , let Q∗ ( s , a ) = maxπ Qπ ( s , a ) denote the optimal value over policies . Thus the optimal policy π∗ is simply taking the action of maximal Q∗ value at each state . The approximation of Q∗ is based on the well-known Bellman temporal consistency : Q∗ ( st , at ) = rt + γmaxat+1 Q ∗ ( st+1 , at+1 ) . ( 3 ) Deep Q-learning ( Mnih et al. , 2013 ) parameterizes the Q-function as Qθ ( x , a ) ( e.g. , a neural network ) , and train the parameters by minimizing the following regression objective : L ( θ ) = Eπ′ [ 0.5 · ( rt + γmaxat+1 Qθ̄ ( st+1 , at+1 ) −Qθ ( st , at ) ) 2 ] , ( 4 ) Single-Step PCL Training Multi-Step PCL Training . . . sequence reward . . . sequence reward Figure 2 : SoftQ-Learning with path consistency learning ( PCL ) objectives , where we illustrate with a vocabulary of size 3 . Left : Single-step objective ( Eq.9 ) , where for each ( st , at ) , the computation involves step t and t+1 . Dashed boxes in dark green and gray indicate the regression target , where the intermediate reward rt is often 0 due to sparsity . The gradient is applied to parameters θ at step t ( indicated by orange color ) . Right : Multi-step objective ( Eq.11 ) which aggregates from step t all the way to T . In this way , the final-step non-zero reward rT is used as the regression target . where θ̄ is the parameters of the target Q-network , which is a slow copy of θ and considered as constant for gradient computation of θ . Here π′ is an behavior policy which can be an arbitrary distribution over text , such as the data distribution or replay buffer ( Mnih et al. , 2013 ) . This makes Q-learning an off-policy algorithm because of its ability to use samples coming from other policies . After learning Qθ , one can induce a policy π from it that takes arg maxaQθ ( s , a ) at each state s. Jaques et al . ( 2017 ) instead sample tokens from the softmax function applied to Qθ . However , the training can be unstable and inefficient due to several challenges : ( 1 ) The bootstrapping nature of the above regression problem can make the training unstable . That is , the regression target rt + γmaxat+1 Qθ̄ ( st+1 , at+1 ) itself is derived from the Q-function to be learned ( Kumar et al. , 2019 ) . The problem is exacerbated in the presence of sparse reward in text generation , where the real observed signal rt is zero for all intermediate t < T ; ( 2 ) The large action space ( e.g. , 104 ) in text generation results in slow updates . In particular , notice that Eq . ( 4 ) applies the gradient update to the Qθ-value of the only one particular token at ( out of the 104 candidate tokens in the vocabulary ) , making the training inefficient ; ( 3 ) Besides , pure off-policy updates could be highly sensitive to the quality of training data , and miss the opportunity of on-policy exploration that maximizes the reward of interest in a more direct way . 3 THE SOFT Q-LEARNING FRAMEWORK In this section , we combat the difficulties of previous RL methods by introducing the softQ-learning ( SQL ) formulation of text generation . We show that the formulation is seamlessly compatible with the common architecture of text generation model ( Eq.1 ) , permitting easy implementation ( §3.1 ) . The formulation further allows us to integrate the latest advances in RL , notably path consistency learning ( Nachum et al. , 2017 ) that makes the RL training efficient and stable in practice ( §3.2 ) . Figure 2 and Algorithm 1 summarizes the resulting SQL framework for efficient training . 3.1 SOFT Q-LEARNING FORMULATION FOR TEXT GENERATION SoftQ-learning ( Haarnoja et al. , 2017 ; Schulman et al. , 2017 ; Nachum et al. , 2017 ) is an maximumentropy ( MaxEnt ) extension to the standard ( hard ) Q-learning ( Mnih et al. , 2015 ; Sutton & Barto , 2018 ) . Under this framework , the agent is encouraged to optimize the reward while staying as stochastic as possible , with the objective JMaxEnt ( π ) = Eτ∼π [ ∑T t=0 γ trt + αH ( π ( · | st ) ) ] , which augments the vanilla J ( π ) with the additional Shannon entropy term H with coefficient α.1 This is appealing because it seamlessly connects the Q-values to the familiar output logits of a text generation model , which enables straightforward implementation of the SQL formulation . Q-values as Generation Model Logits . We show the connection of the Q-values with the logits , i.e. , the model outputs right before the softmax layer . Concretely , with the SQL objective , the following relationship between optimal policy π∗ and action-value Q∗ holds ( Haarnoja et al. , 2017 ; Schulman et al. , 2017 ) : π∗ ( a | s ) = expQ ∗ ( s , a ) ∑ a′ expQ ∗ ( s , a′ ) . ( 5 ) 1WLOG , we can assume α=1 , as it can be folded into the reward function by scaling the latter with 1/α . This form is highly reminiscent of the softmax layer of the generation model in Eq. ( 1 ) . The connection suggests that we can naturally parameterize the Q-function in SQL as the generation model logit function , i.e. , Qθ ( s , a ) ≡ fθ ( a | s ) . In other words , the model output fθ ( a | s ) , originally interpretted as the “ logit ” of token a given the preceding tokens s , is now re-interpretted as the Qvalue of action a in state s. When achieving optimality , fθ∗ ( a | s ) , namely Q∗ ( s , a ) , represents the best possible future reward achievable by generating token a in state s. Similarly , the full generation model pθ ( a | s ) in Eq . ( 1 ) that applies softmax to fθ now precisely corresponds to the policy πθ induced from Qθ ( s , a ) . That is , πθ ( a | s ) = expQθ ( s , a ) ∑ a′ expQθ ( s , a ′ ) ≡ exp fθ ( a | s ) ∑ a′ exp fθ ( a ′ | s ) = pθ ( a | s ) . ( 6 ) We could further gain even more intuitive interpretation of the above generation policy π∗ from the lens of advantage function ( Sutton & Barto , 2018 ) . Specifically , in SQL , the optimal state-value function is the log-normalizer of the optimalQ-values ( Haarnoja et al. , 2017 ; Schulman et al. , 2017 ) . This allows us to rewrite Eq . ( 5 ) into a more concise form : V ∗ ( s ) = log ∑ a′ expQ∗ ( s , a′ ) , π∗ ( a | s ) = exp ( Q∗ ( s , a ) −V ∗ ( s ) ) = expA∗ ( s , a ) , ( 7 ) where A∗ is the optimal advantage function . The equation says that , in the proposed text generation SQL formulation , the optimal policy generates token a in state s according to the token ’ s advantage .

Summary: This paper considers the problem of learning text generation models using reinforcement learning. The problem is challenging in that RL algorithm becomes inefficient or unstable when dealing with large action space and the sparse reward situations in text generation. To address these problems, this paper adapts the path consistency learning approach to the text generation setting. It allows to train the model with both on- and off-policy samples, bridge the sparse reward signal to directly supervise the Q-function learning, and makes efficient updates to Q-values by considering all candidate actions together. Experiments result show that the proposed approach are effective in solving a wide range of applications where MLE are not applicable, where it achieves better performance than previous RL-based methods as well as task-specific methods. In addition, on standard MLE-based tasks, the proposed approach could also achieve competitive performance when training the models from scratch.