Text Generation by Learning from Demonstrations

1 INTRODUCTION . A dominant approach to text generation is to use autoregressive models learned by maximum likelihood estimation ( MLE ) on supervised data . However , this approach introduces two well-known discrepancies between training and evaluation objectives that lead to undesired generations . First , the training loss is negative log-likelihood , whereas the evaluation is based on human judgment of the output quality . Under model misspecification , MLE tends to over-generalize , assigning large probability mass to both high-quality and low-quality sequences ( Huszár , 2015 ; Simon et al. , 2019 ) . Therefore , in practice , we must carefully select the decoding algorithms to produce high-quality outputs . Second , during training , the autoregressive model conditions on the gold history/prefix ; however , at inference time it conditions on model-generated history . This is known as the exposure bias problem ( Ranzato et al. , 2016 ; Bengio et al. , 2015 ) . In the worst case , one incorrect prediction can produce a low-probability prefix under the gold data distribution , and errors compound in each of the following steps ( Ross et al. , 2011 ) . In practice , prior work has observed problems such as repetition and hallucination partly due to exposure bias ( Holtzman et al. , 2020 ; Wang & Sennrich , 2020 ) . We aim to bridge the gap between training and evaluation in this paper . To match training and evaluation objectives , ideally we should maximize output quality given model-generated histories . This corresponds to the reinforcement learning ( RL ) objective : maximizing the expected reward ( quality ) over trajectories ( sequences ) induced by the policy ( model ) . However , optimizing this objective is notoriously difficult . Prior RL approaches mainly focus on fine-tuning a learned model to optimize sequence-level metrics such as BLEU ( Papineni et al. , 2002 ) , but empirically it remains unclear if RL is beneficial to text generation ( Wu et al. , 2018 ; Choshen et al. , 2020 ) . Note that many challenges in RL arise from exploring an exponentially large space of sequences , with sparse rewards only on those close to the reference . We thus propose to learn from only the reference sequences without interaction ( i.e. , the offline setting ) . Specifically , we use off-policy policy gradient with importance weighting ( Hastings , 1970 ; Hachiya et al. , 2009 ; Parshakova et al. , 2019 ) , where training examples with higher probability under the model are weighted higher . Further , our reward functions approximate human judgment of the output quality by estimating how likely a human would have generated a sequence . We call our algorithm GOLD ( Generation by Off-policy Learning from Demonstrations ) . Results on news summarization , question generation , and machine translation show that GOLD leads to better model performance than MLE and RL fine-tuning by both task metrics and human-rated quality . Further , our analysis shows that GOLD learns high-precision models that are less sensitive to decoding algorithms . In addition , it alleviates exposure bias : the output quality does not degrade much as generation length increases . 2 FROM MLE TO RL FRAMEWORK . MLE training . Given a context x such as a document , we want to generate a sequence of tokens y = ( y0 , . . . , yT ) , where yi comes from a vocabulary V . The generator is modeled by a conditional probability distribution parametrized by θ : pθ ( y | x ) = ∏T t=0 pθ ( yt | y0 : t−1 , x ) , where y0 : t−1 denotes the prefix y0 , . . . , yt−1 . Let phuman ( y | x ) denote the data-generating distribution . Using MLE , the loss function is L ( θ ) = −Ey∼phuman [ T∑ t=0 log pθ ( yt | y0 : t−1 , x ) ] . ( 1 ) At inference time , we generate tokens sequentially according to pθ . Evaluation . In practice , the quality of an output often relies on task-specific metrics such as fluency , correctness , and interestingness . Here for generality we consider perceptual quality ( Huszár , 2015 ; Hashimoto et al. , 2019 ) which measures how likely a human would have generated the output given the context , i.e. , phuman ( y | x ) . Thus the evaluation metric is Ey∼pθ [ T∑ t=0 log phuman ( yt | y0 : t−1 , x ) ] . ( 2 ) Comparing ( 1 ) and ( 2 ) , we see that the training objective encourages high recall : the model must put probability mass on all human-generated sequences . In contrast , the evaluation metric encourages high precision : all outputs from the model must be of high quality . Unfortunately , directly optimizing the evaluation metric is impossible because phuman is unknown and the expectation is difficult to estimate . We therefore develop a training objective that closely approximates ( 2 ) in the RL framework . RL formulation . Let ’ s consider generation as a sequential decision-making process . At each time step t , the policy πθ takes an action at ∈ V , transits to the next state st+1 = ( y0 : t , x ) , and receives a reward rt . The policy corresponds to the generation model : πθ ( at | st ) = pθ ( at | y0 : t−1 , x ) . We can thus represent a sequence as a trajectory τ = ( s0 , a0 , r0 , . . . , sT , aT , rT ) . The set of trajectories derived from the training data is called demonstrations which show the desired behavior of a policy . The RL objective is to maximize J ( θ ) = Eτ∼πθ [ ∑T t=0 γ trt ] , where γ ∈ ( 0 , 1 ] is the discount factor , and πθ ( τ ) denotes the distribution of τ induced by πθ . If we knew oracle rewards rt = phuman ( at | st ) , then this objective would be exactly the evaluation metric we want to optimize . Next , we describe how to optimize J ( θ ) with reward functions that approximate phuman . 3 APPROACH . 3.1 OFF-POLICY POLICY GRADIENT . Policy gradient . A straightforward way to optimize J ( θ ) is policy gradient ( PG ) ( Williams , 1992 ; Sutton et al. , 2000 ) . The gradient is given by ∇θJ ( θ ) = Eτ∼πθ [ ∑ t ∇θ log πθ ( at | st ) Q̂ ( st , at ) ] , ( 3 ) where Q̂ ( st , at ) = ∑T t′=t γ t′−trt′ is the estimated return from state st . The expectation is estimated by Monte Carlo samples from πθ . In text generation , the return Q̂ ( st , at ) is often a sequence-level reward such as BLEU . In practice , the policy is likely to get stuck in a region of zero reward during training , generating gibberish without receiving any learning signal ( Li et al. , 2018 ; Keneshloo et al. , 2019 ) . A common remedy is to initialize the policy with the MLE solution and/or interleave with MLE gradient update during PG . However , this would bias the parameters towards the MLE solution , thus often leads to marginal gains in practice ( Wu et al. , 2018 ; Choshen et al. , 2020 ) . Offline learning . To avoid zero-reward regions , we would like to reduce interaction with the environment and stay close to the demonstrated trajectories . In the extreme case , the policy is learned solely from the static demonstrations without additional interaction with the environment , which is referred to as the offline setting . While it is in general a more challenging problem , we argue that the offline setting is appropriate for text generation ( Serban et al. , 2017 ; Jaques et al. , 2019 ) . First , the environment dynamics is known : once a token is generated , we deterministically transition to the next state with the additional token appended to the prefix ; no interaction is needed to learn the environment . Second , while exploration may lead to high-quality sequences different from the reference , we lack a good reward function to identify them ( Novikova et al. , 2017 ; Aharoni & Goldberg , 2018 ; Clark et al. , 2019 ) . Therefore , the benefit of exploration in text generation is limited . In the offline setting , we can not estimate the expected return of πθ by sampling trajectories from it , and must use trajectories from a different behavioral policy πb , known as off-policy learning in RL . A common technique to estimate expectations under one distribution πθ given samples from a different distribution πb is importance sampling , which leads to the following unbiased estimator of the gradient ( Precup et al. , 2000 ) : Eτ∼πb [ ∑ t wt∇θ log πθ ( at | st ) Q̂ ( st , at ) ] , with importance weights wt = ∏t t′=0 πθ ( at′ |st′ ) πb ( at′ |st′ ) . Approximations . Computing the importance weights above requires multiplying per-action importance weight over multiple time steps . In practice , we have found that it is sensitive to optimization hyperparameters and takes longer to converge . Therefore , we use the per-action approximation : wt ≈ πθ ( at|st ) πb ( at|st ) . This corresponds to optimizing the expected return under the off-policy state distribution induced by πb and the on-policy action distribution of πθ . Although this estimator is biased , empirically it has been shown to reduce variance and work reasonably well if πb and πθ are close ( Serban et al. , 2017 ; Levine et al. , 2020 ) . Another obstacle is that we do not know πb which produced the demonstrationsD = { ( x ( i ) , y ( i ) ) } Ni=1 . One option is to estimate πb onD . Here we take a simpler approach that uses the empirical distribution : πb ( τ ) ≈ 1/N for τ ∈ D and 0 otherwise . As a result , the denominator in wt is a constant and can be ignored in optimization . Our final approximated gradient has the form : ∇θJ ( θ ) ≈ N∑ i=1 T∑ t=0 πθ ( a i t | sit ) ∇θ log πθ ( ait | sit ) Q̂ ( sit , ait ) , ( 4 ) where the superscript i represents the ith trajectory . Compared with the MLE gradient : ∑N i=1 ∑T t=0∇θ log πθ ( ait | sit ) , our gradient ( 4 ) upweights actions with high return and actions preferred by the current policy πθ . Intuitively , it encourages the learning algorithm to focus on “ easy ” examples ( high likelihood under the model ) which improves precision . 3.2 REWARD . Let R be the reward function such that rt = R ( st , at ) . To optimize the perceptual quality of a sequence ( see ( 2 ) ) , we want R ( s , a ) to approximate phuman ( a | s ) , i.e. , how likely humans would have generated a given s. In general , it is hard to develop a reliable reward function for text generation tasks because it must work well for a large set of possible generations . In the offline setting , however , we can restrict the domain of R to state-action pairs on the demonstrations . Next , we propose three reward functions . δ-reward . An obvious choice is a sequence-level reward , which considers all demonstrations to be equally good and assigns zero reward to any other outputs . Formally , Rδ ( st , at ) def = { 1 , if t = T and ( s0 : T , a0 : T ) ∈ D 0 , otherwise ( 5 ) where a reward of one is received in the terminal state for any trajectory in the demonstrations . Estimated phuman . In text generation tasks , an input often has many correct outputs and the reference may be an uncommon output that contains rare words or has complex syntax . To account for different likelihood of the references , we estimate the probability of each reference by minimizing KL ( phuman‖q ) , where q ( a | s ) approximates phuman ( a | s ) . This is equivalent to finding the MLE solution ( denoted by pMLE ) .1 Importantly , pMLE is a reasonable approximation to phuman when restricted to the demonstrations . It is not a good reward function in general , however ; it can assign large probability mass to low-quality outputs . Given the estimated perceptual quality pMLE ( a | s ) , we define two reward functions . Our first reward function corresponds to a product of probabilities when summed over the trajectory : Rp ( s , a ) def = log pMLE ( a | s ) . ( 6 ) Assuming γ = 1 , the return at time step t is Q̂t ( st , at ) = ∑T t′=t log pMLE ( at | st ) . Thus a sequence has high reward only if every word has high likelihood under pMLE . To allow for partial credits even if bad actions are taken at certain steps , we define another reward function corresponding to the sum of probabilities : Rs ( s , a ) def = pMLE ( a | s ) . ( 7 ) The return Q̂ ( st , at ) = ∑T t′=t pMLE ( at | st ) , thus the policy can recover from bad decisions if the subsequent actions receive high reward . 3.3 THE GOLD ALGORITHM Algorithm 1 : GOLD 1 πθ ← pMLE , π̃θ ← pMLE 2 for step = 1 , 2 , . . . , M do 3 Sample a minibatch B = { ( xi , yi ) } |B|i=1 4 foreach ( sit , ait ) do 5 Compute importance weights max ( u , π̃θ ) , and compute returns Q̂ ( sit , a i t ) − b 6 Update θ by ( 4 ) using gradient descent 7 if step % k = 0 then π̃θ ← πθ 8 Return : πθ Our full algorithm based on off-policy PG is shown in Algorithm 1 . For importance weights πθ ( a | s ) , to avoid drastic changes , we initialize πθ with the MLE solution . In addition , we compute the importance weights by a weighting policy π̃θ that synchronizes with πθ periodically so that the weights do not change frequently between updates . We also lower-bound the importance weight by a small number u . Another source of variance comes from policy gradients . Since our return is computed from a sum or product of probabilities ( ( 6 ) and ( 7 ) ) , we truncate the future trajectory after five steps . We follow the common practice to subtract a baseline b from the return to reduce variance ; moreover , to avoid negative reward on the demonstrations ( after subtracting baseline ) , we lower-bound pMLE in ( 6 ) and ( 7 ) by a small number c. In practice , GOLD is easy to implement ; further , given an existing pMLE , the GOLD-training stage usually takes less time than MLE . The code is available.2 .

Summary: This paper introduces a new optimization method for text generation that improves upon directly optimizing MLE. It frames text generation learning as an off-policy reinforcement learning (RL) problem using demonstrations (the training examples). The authors also discuss why off-policy learning is more suitable for text generation than on-policy learning. After simplifications and approximations, the proposed optimization objective comes down to a form that is similar to MLE, but upweighs training examples that are more likely under the learned policy $\pi_\theta$ ("easy" examples) and having higher estimated rewards (considering the future).