Connecting the Dots Between MLE and RL for Sequence Prediction

1 INTRODUCTION . Sequence prediction problem is ubiquitous in many applications , such as generating a sequence of words for machine translation ( Wu et al. , 2016 ; Sutskever et al. , 2014 ) , text summarization ( Hovy & Lin , 1998 ; Rush et al. , 2015 ) , and image captioning ( Vinyals et al. , 2015 ; Karpathy & Fei-Fei , 2015 ) , or taking a sequence of actions to complete a task . In these problems ( e.g. , Mnih et al. , 2015 ; Ho & Ermon , 2016 ) , we are often given a set of sequence examples , from which we want to learn a model that sequentially makes the next prediction ( e.g. , generating the next token ) given the current state ( e.g. , the previous tokens ) . A standard training algorithm is based on supervised learning which seeks to maximize the loglikelihood of example sequences ( i.e. , maximum likelihood estimation , MLE ) . Despite the computational simplicity and efficiency , MLE training can suffer from compounding error ( Ranzato et al. , 2016 ; Ross & Bagnell , 2010 ) in that mistakes at test time accumulate along the way and lead to states far from the training data . Another line of approaches overcome the training/test discrepancy issue by resorting to the reinforcement learning ( RL ) techniques ( Ranzato et al. , 2016 ; Bahdanau et al. , 2017 ; Rennie et al. , 2017 ) . For example , Ranzato et al . ( 2016 ) used policy gradient ( Sutton et al. , 2000 ) to train a text generation model with the task metric ( e.g. , BLEU ) as reward . However , RL-based approaches can face challenges of prohibitively poor sample efficiency and high variance . To this end , a diverse set of methods has been developed that is in a middle ground between the two paradigms of MLE and RL . For example , RAML ( Norouzi et al. , 2016 ) adds reward-aware perturbation to the MLE data examples ; SPG ( Ding & Soricut , 2017 ) leverages reward distribution for effective sampling of policy gradient . Other approaches such as data noising ( Xie et al. , 2017 ) also show improved results . In this paper , we establish a unifying perspective of the above distinct learning algorithms . Specifically , we present a generalized entropy regularized policy optimization framework , and show that the diverse algorithms , such as MLE , RAML , data noising , and SPG , can all be re-formulated as special cases of the framework , with the only difference being the choice of reward and the values of two weight hyperparameters ( Figure 1 ) . In particular , we show MLE is equivalent to using a RAML . Delta-function reward which returns 1 to model samples that match training examples exactly , and −∞ to any other samples . Such extremely restricted reward has literally disabled any exploration of the model beyond training data , yielding brittle prediction behaviors . Other algorithms essentially use various locally-relaxed rewards , joint with the model distribution , for broader ( and more costly ) exploration during training . Besides the new views of the existing algorithms , the unifying perspective also leads to new algorithms for improved learning . We develop interpolation algorithm , which , as training proceeds , gradually expands the exploration space by annealing both the reward function and the weight hyperparameters . The annealing in effect dynamically interpolates among the existing algorithms from left to right in Figure 1 . We conduct experiments on the tasks of text generation including machine translation and text summarization , and game imitation learning . The interpolation algorithm shows superior performance over various previous methods . 2 RELATED WORK . Given a set of data examples , sequence prediction models are usually trained to maximize the loglikelihood of the next label ( token , action ) conditioning on the current state observed in the data . Reinforcement learning ( RL ) addresses the discrepancy between training and test by also using models ’ own predictions at training time . Various RL approaches have been applied for sequence generation , such as policy gradient ( Ranzato et al. , 2016 ) and actor-critic ( Bahdanau et al. , 2017 ) . Reward augmented maximum likelihood ( RAML ) ( Norouzi et al. , 2016 ) is an algorithm in between MLE and policy gradient . Mathematically , RAML shows that MLE and maximum-entropy policy gradient are respectively minimizing KL divergences in opposite directions . Koyamada et al . ( 2018 ) thus propose to use the more general α-divergence as a combination of the two paradigms . Our framework is developed in a different perspective , reformulates a different and more comprehensive set of algorithms , and leads to new insights in terms of exploration and learning efficiency of the various algorithms . Besides the algorithms discussed in the paper , there are other learning methods for sequence models . For example , Hal Daumé et al . ( 2009 ) ; Leblond et al . ( 2018 ) ; Wiseman & Rush ( 2016 ) use a learning-to-search paradigm for sequence generation or structured prediction . Scheduled Sampling ( Bengio et al. , 2015 ) and variants ( Zhang et al. , 2019 ) adapt MLE by randomly replacing ground-truth tokens with model predictions as the input for decoding the next-step token . Policy optimization for reinforcement learning is studied extensively in robotic and game environment . For example , Peters et al . ( 2010 ) introduce a relative entropy regularization to reduce information loss during learning . Schulman et al . ( 2015 ) develop a trust-region approach for monotonic improvement . Dayan & Hinton ( 1997 ) ; Levine ( 2018 ) ; Abdolmaleki et al . ( 2018 ) study the policy optimization algorithms in a probabilistic inference perspective . Zhu et al . ( 2018 ) combine imitation learning with RL , whose approach is orthogonal to ours and can be plugged into our framework to incorporate imitation reward . The entropy-regularized policy optimization formulation presented here can be seen as a generalization of many of the previous policy optimization methods . Besides , we formulate the framework primarily in the sequence generation context . 3 CONNECTING THE DOTS . We first present a generalized formalism of entropy regularized policy optimization . The formulation contains a reward function and two weight hyperparameters that define the learning procedure . Therefore , varying the values of the reward and weights result in a large space of algorithms . We show that several existing popular algorithms , which were originally proposed in distinct perspectives , can all be seen as members in the space . In particular , we reformulate the MLE algorithm in the same policy optimization form , which enables side-by-side comparison between the broad spectrum of algorithms . The resulting unifying view provides new insights into the exploration and computation efficiency , and creates improved learning approaches for sequence prediction . For clarity , we present the framework in the sequence generation context . The formulations can straightforwardly be extended to other settings such as imitation learning in robotic and game environments , as discussed briefly at the end of this section and also shown in the experiment . We first establish the basic notations . Let y = ( y1 , . . . , yT ) be the sequence of T tokens . Let y∗ be a training example drawn from the empirical data distribution . From the sequence examples , we aim to learn a sequence generation model pθ ( y ) = ∏ t pθ ( yt|y1 : t−1 ) with parameters θ . Note that generation of y can condition on other factors . For example , in machine translation , y is the sentence in target language and depends on an input sentence in source language . For simplicity of notations , we omit the conditioning factors . 3.1 ENTROPY REGULARIZED POLICY OPTIMIZATION ( ERPO ) . Policy optimization is a family of reinforcement learning ( RL ) algorithms . Assume a reward functionR ( y|y∗ ) ∈ R that evaluates the quality of generation y against the true y∗ . For example , BLEU score ( Papineni et al. , 2002 ) can be a reward in machine translation . The general goal of policy optimization is to learn the model pθ ( y ) ( a.k.a policy ) 1 to maximize the expected reward . Previous work develops entropy regularized approaches , which augment the objective with information theoretic regularizers for stabilized training . We present a generalized variational formulation of ERPO , which , as we show shortly , has the power of subsuming an array of other popular algorithms . Specifically , we introduce a non-parametric variational distribution q ( y ) w.r.t the model pθ ( y ) . The objective to maximize is as follows : L ( q , θ ) = Eq [ R ( y|y∗ ) ] − αKL ( q ( y ) ‖pθ ( y ) ) + βH ( q ) , ( 1 ) where KL ( ·‖· ) is the Kullback–Leibler divergence forcing q to stay close to pθ ; H ( · ) is the Shannon entropy imposing maximum entropy assumption on q ; and α and β are balancing weights of the respective terms . Intuitively , the objective is to maximize the expected reward under the variational distribution q while minimizing the distance between q and the model pθ , with maximum entropy regularization on q . The above formulation is relevant to and can be seen as a variant of previous policy optimization approaches in RL literature , such as relative entropy policy search ( Peters et al. , 2010 ) , maximum entropy policy gradient ( Ziebart , 2010 ; Haarnoja et al. , 2017 ) , and other work where the variational distribution q is formulated either as a non-parametric distribution as ours ( Abdolmaleki et al. , 2018 ; Peters et al. , 2010 ) or parametric one ( Schulman et al. , 2015 ; 2017a ; Teh et al. , 2017 ) . The objective can be maximized with a standard EM procedure ( Neal & Hinton , 1998 ) that iterates two coordinate ascent steps optimizing q and θ , respectively . At iteration n : E-step : qn+1 ( y ) ∝ exp { α log pθn ( y ) +R ( y|y∗ ) α+ β } , M-step : θn+1 = arg maxθ Eqn+1 [ log pθ ( y ) ] . ( 2 ) In the E-step , q has a closed-form solution , which is an energy-based distribution . We can have an intuitive interpretation of its form . First , it is clear to see that if α → ∞ , we have qn+1 = pnθ . This is also reflected in the objective Eq . ( 1 ) where a larger weight α encourages q to be close to pθ . Second , the weight β serves as the temperature of the q softmax distribution . In particular , a large temperature β → ∞ makes q a uniform distribution , which is consistent with the outcome of an infinitely large maximum entropy regularization in Eq. ( 1 ) . In the M-step , the update rule can be interpreted as maximizing the log-likelihood of samples from the distribution q . 1In the following , we will use the term “ model ” and “ policy ” exchangeably . Token-level Formulation In the context of sequence generation , it is sometimes more convenient to express the equations at token level ( instead of the sequence level ) , as shown when we devise a new algorithm in the next section . To this end , we decompose R ( y|y∗ ) along the time steps : R ( y|y∗ ) = ∑ t R ( y1 : t|y∗ ) −R ( y1 : t−1|y∗ ) : = ∑ t ∆R ( yt|y1 : t−1 , y∗ ) , ( 3 ) where ∆R ( yt|y∗ , y1 : t−1 ) measures the reward contributed by token yt . The solution of q in Eq . ( 2 ) can then be re-written as : qn+1 ( y ) ∝ ∏ t exp { α log pθn ( yt|y1 : t−1 ) + ∆R ( yt|y1 : t−1 , y∗ ) α+ β } . ( 4 ) The Algorithm Space The above ERPO formalism includes three key components , namely , the reward R and the weight hyperparameters α and β > 0 . Variation in these components can result in different procedures of updating the model . In other words , different algorithms in the ERPO family correspond to a point ( or a region ) in the space spanned by the three components . The following sections visit a set of existing approaches , and connect them to the unifying picture by reformulating their seemingly distinct objectives . Figure 1 illustrates the particular algorithms in the space , clustered by the exploration behavior in learning , of which we will discuss more . Softmax Policy Gradient ( SPG ) We first briefly discuss the previous RL algorithms for sequence prediction that fit in the ERPO formalism . SPG ( Ding & Soricut , 2017 ) was originally developed in the perspective of combining the rewardR and policy pθ to improve sampling quality . The algorithm is equivalent to setting β = 0 and treating α > 0 as the temperature of the energy-based distribution q ( y ) . That is , q ( y ) in the E-step of Eq . ( 2 ) is now in the form q ( y ) ∝ pθ ( y ) exp { R ( y|y∗ ) /α } . The reward R is set to any normal task-specific reward . Note that sampling from q ( y ) ( e.g. , in the M-step ) is typically difficult due to its energy-based form and the fact that the task reward R often does not have particular structures amenable for sampling . We will see in the next section that the MLE algorithm in contrast uses a special reward to avoid the computational difficulty in sampling , at the cost of restricted exploration during training . We also note the previous work of Sequence Tutor ( Jaques et al. , 2017 ) , which was motivated by the idea of using an MLE-trained policy as a prior to guide the learning of the target policy in an RL framework . The formalism closely resembles SPG , namely ( α > 0 , β = 0 ) , with the exception that the variational distribution q ( y ) in Sequence Tutor is a parameterized model instead of a nonparametric one as in SPG and our more general ERPO formulation .

Summary: This paper presents a formalism of entropy regularized policy optimization. They also show that various policy gradients algorithms can be reformulated as special instances of the presented novel formalism. The only difference between them being the reward function and two weight hyperparameters. Further, the paper proposes an interpolation algorithm, which, as training proceeds, gradually expands the exploration of space by annealing the reward function and the weight hyperparameters. Experiments on text generation tasks and game imitation learning show superior performance over previous methods.