Adaptive Correlated Monte Carlo for Contextual Categorical Sequence Generation

1 INTRODUCTION . Contextual categorical sequence generation is a core modeling component in a wide variety of machine learning tasks , such as neural program synthesis ( Bunel et al. , 2018 ; Devlin et al. , 2017b ; Si et al. , 2018 ; Chen et al. , 2019 ) and image captioning ( Vinyals et al. , 2015 ; Xu et al. , 2015 ) . Typically , an encoder-decoder framework is applied . The encoder maps a contextual input to a latent representation , conditioning on which and previously generated tokens the decoder generates categorical tokens in a consecutive manner ( Bahdanau et al. , 2014 ; Sutskever et al. , 2014 ; Cho et al. , 2014 ; Rush et al. , 2015 ; Chopra et al. , 2016 ) . It is common to train contextual sequence generation models using maximum likelihood estimation ( MLE ) , which attempts to maximize the likelihood of each token in a target sequence given its preceding tokens . Learning with MLE is often sub-optimal as it does not directly optimize the evaluation metric of the end task . It generally suffers from the exposure bias ( Bengio et al. , 2015 ; Ranzato et al. , 2016 ) , which refers to the discrepancy between training and generation using the Teacher Forcing ( Williams & Zipser , 1989 ) strategy , i.e. , during training ground truth tokens are used as inputs , while during generation , only generated tokens are available . Thus giving higher likelihoods to target sequences does not guarantee the model to generate sequences close to the target or good sequences . Moreover , MLE requires target sequences for training , while for many scenarios in task-oriented dialogue ( Williams & Young , 2007 ) and program synthesis ( Zhong et al. , 2017 ) , only the final rewards to the generated sequences are available . To overcome the aforementioned issues of MLE , it is common to refine a contextual sequence generation model pre-trained with MLE under the reinforcement learning ( RL ) framework ( Zaremba & Sutskever , 2015 ; Ranzato et al. , 2016 ; Bahdanau et al. , 2016 ; Wu et al. , 2018 ; Paulus et al. , 2017 ) . The objective becomes maximizing the expected rewards of model generated sequences . During training , only the model generated tokens are fed into the model so that the exposure bias is avoided . The reward to guide RL can be : 1 ) a task-dependent user-defined metric , such as CIDEr for image captioning ( Vedantam et al. , 2015 ) and Generalization for neural program synthesis ( Bunel et al. , 2018 ) ; and 2 ) automatically learned reward using a discriminator or language model ( Yang et al. , 2018 ; Yu et al. , 2017 ; Lamb et al. , 2016 ; Caccia et al. , 2018 ; d ’ Autume et al. , 2019 ) . The RL training Code link : https : //github.com/xinjiefan/ACMC ICLR enables direct improvement of the user-defined or learned reward . Moreover , in cases where only weak-supervision is available , e.g. , in neural program synthesis , RL may considerably improve the model performance ( Bunel et al. , 2018 ; Zhong et al. , 2017 ) . However , the gradients of the expected reward in RL often suffer from high Monte Carlo ( MC ) estimation variance , due to noisy and/or sparse rewards and the large action space that grows exponentially with the sequence length . There has been significant recent interest in variance reduction methods for MC gradient estimation ( Mohamed et al. , 2019 ) . A highly effective solution is the reparameterization trick ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) , which , however , is applicable to neither discrete variables nor non-differentiable reward functions . For variance reduction involving discrete variables , one potential solution is to combine the Gumbel-softmax trick , which relaxes the discrete variables to continuous ones , with reparameterization to produce low-variance but biased gradients ( Jang et al. , 2017 ; Maddison et al. , 2017 ) . Another common way for variance reduction is adding appropriate baselines ( Owen , 2013 ; Williams , 1992 ; Paisley et al. , 2012 ; Ranganath et al. , 2014 ; Mnih & Gregor , 2014 ) , and there exist several such methods customized for discrete variables ( Tucker et al. , 2017 ; Grathwohl et al. , 2018 ) . However , due to either the inherent biases or difficulty to learn the parameters of the baselines , it is unclear how effective these newly proposed estimators are in backpropagating the gradients through a sequence of discrete variables ( Yin et al. , 2019 ) . This is exacerbated in contextual categorical sequence generation problems , where it is common for a sequence to contain quite a few tokens/actions , each of which is selected from a set of thousands of candidates . Another practical issue is that generating a token from a large vocabulary via the softmax output layer is often computationally heavy . This prevents a categorical sequence generation model from being deployed to low-power devices . Despite significant recent efforts in addressing the computation bottleneck due to a wide softmax layer ( Shim et al. , 2017 ; Zhang et al. , 2018 ; Chen et al. , 2018 ) , for categorical sequence generation , it is so far unclear how to address the softmax computational bottleneck while at the same time providing low-variance gradient of its RL objective . This paper makes two primary contributions : 1 ) To address the high gradient variance issue , we adapt to contextual categorical sequence generation tasks the augment-REINFORCE-swap-merge ( ARSM ) estimator ( Yin et al. , 2019 ) , which provides unbiased , low-variance gradients for categorical variables , using token-level rewards from correlated MC rollouts that naturally serve as the baselines for each other . We show that the number of rollouts is adapted to the model uncertainty on the latest generated token , and can also be manually controlled to balance variance reduction and computation . 2 ) To address the high generation cost issue , we replace the generation of each categorical variable in the sequence , via a categorical softmax output layer , with the generation of a sequence of binary decisions on a binary tree from its root node to a leaf node , which is occupied by a unique term of the vocabulary . For training , we adapt the augment-REINFORCE-merge ( ARM ) estimator ( Yin & Zhou , 2019 ) , which provides unbiased , low-variance gradients for binary variables , to backpropagate the gradients through the sequence of binary sequences . Under this binary-tree construction , the cost of generating a categorical token reduces from O ( V ) to O ( log2 ( V ) ) , where V is the vocabulary size . We demonstrate our methods on two representative contextual categorical sequence generation tasks , with the number of actions ranging from 53 ( neural program synthesis ) to 9978 ( image captioning ) . 2 PRELIMINARIES ON CONTEXTUAL SEQUENCE GENERATION . For a dataset of context-output pairs D : = { xi , yi } Ni=1 , our goal is to learn the conditional distribution of output yi given its context xi , expressed as pθ ( yi |xi ) . Below we drop the data index subscript for brevity . We focus on the case that an output is a sequence of T categorical variable as y = { y1 , · · · , yT } , where yt ∈ { 1 , . . . , V } . A common way to model pθ ( y |x ) is to decompose it as pθ ( y |x ) = ∏T t=1 pθ ( yt | y1 : t−1 , x ) , where the t-th term in the product , which models the distribution of token yt conditioning on the context x and previously generated tokens y1 : t−1 , is commonly parameterized by a recurrent neural network ( Sutskever et al. , 2014 ) . MLE is a common way to train the model : θ̂MLE = argmaxθ E { x , y } ∼pdata ( x , y ) [ log pθ ( y |x ) ] . Viewing pθ ( yt | y1 : t−1 , x ) as a stochastic policy for choosing an action given the state , we can formulate contextual sequence generation as an RL problem and infer the policy parameter θ as θ̂RL = argmaxθ E { x , y } ∼pdata ( x , y ) Ez∼pθ ( · |x ) [ r ( z |x , y ) ] , where r ( z |x , y ) denotes the reward of the generated ( hypothesis ) sequence z given the context x and the reference target sequence y . For example , for image captioning , the reward could be the CIDEr score that measures the similarity between the generated caption z and the reference y ( Rennie et al. , 2017 ) . Denote σ ( · ) as the softmax function and Tθ ( · ) as a deterministic function defined by a deep neural network with parameter θ . We model pθ ( z |x ) = ∏T t=1 pθ ( zt | z1 : t−1 , x ) , where pθ ( zt |x , z1 : t−1 ) = Cat ( σ ( φt ) ) , φt : = Tθ ( x , z1 : t−1 ) . ( 1 ) For a context-target pair { x , y } , we can expand the expected reward ER under policy pθ ( z |x ) as ER = Ez∼pθ ( · |x ) [ r ( z |x , y ) ] = Ez1 : t−1∼pθ ( · |x ) Ezt∼Cat ( σ ( φt ) ) [ r ( z1 : t |x , y ) ] , ( 2 ) where the partial-sentence reward is defined as r ( z1 : t |x , y ) = Ezt+1 : T∼pθ ( · | x , z1 : t ) [ r ( z |x , y ) ] . Using the chain rule and REINFORCE ( Williams , 1992 ) estimator , we have ∇θER = ∑T t=1∇φtER∇θφt , ∇φtER = Ez1 : t−1∼pθ ( · |x ) ∇φtEzt∼Cat ( σ ( φt ) ) [ r ( z1 : t |x , y ) ] , ∇φtEzt∼Cat ( σ ( φt ) ) [ r ( z1 : t |x , y ) ] = Ezt∼Cat ( σ ( φt ) ) [ r ( z1 : t |x , y ) ∇φt ln p ( zt ; σ ( φt ) ) ] . The main challenge here is to control the variance in estimating∇θER . A variety of methods have been proposed . For example , drawing sentence z ∼ pθ ( z |x ) and using its reward r ( z |x , y ) to approximate all partial-sentence rewards { r ( z1 : t |x , y ) } 1 , T , Ranzato et al . ( 2016 ) introduce MIXER with a scheduled training iterating between MLE and RL , estimating the RL gradient as ∇̂θER = ∑T t=1 ( r ( z |x , y ) − b ( z1 : t−1 ) ) ∇φt ln p ( zt ; σ ( φt ) ) ∇θφt , where b ( z1 : t−1 ) is a baseline function . To improve MIXER , Rennie et al . ( 2017 ) introduces the self-critic ( SC ) sequence training algorithm that sets b ( z1 : t−1 ) = r ( z̃ |x , y ) for all t , where z̃ is a greedy sequence rollout under pθ ( z |x ) . As using sentence-level reward r ( z |x , y ) to guide the learning is found to be sensitive to the algorithm parameters such as the learning rate , Liu et al . ( 2017 ) follow Yu et al . ( 2017 ) to use token-level rewards that approximate each partial-sentence reward with K independent MC rollouts ( MC-K ) as r̂ ( z1 : t |x , y ) = 1K ∑K k=1 r ( z1 : t , z ( k ) ( t+1 ) : T |x , y ) , z ( k ) ( t+1 ) : T ∼ pθ ( · |x , z1 : t ) . With these token-level rewards r̂ ( z1 : t |x , y ) , Liu et al . ( 2017 ) estimate the gradient as ∇̂θER = ∑T t=1 ( r̂ ( z1 : t |x , y ) − b ( z1 : t−1 ) ) ∇φt ln p ( zt ; σ ( φt ) ) ∇θφt . ( 3 ) Following SC , for token t , one may choose baseline b ( z1 : t−1 ) = r ( z1 : t−1 , z̃t : T |x , y ) , where z̃t : T is a greedy rollout following partial sentence z1 : t−1 . In addition to being more robust , using token-level rewards r̂ ( z1 : t ) to guide the learning is often necessary when the reward signal is sparse , e.g. , in neural program synthesis , it is common that a full MC roll receives zero reward with high probability . In addition to these methods mentioned above , the actor-critic method ( Sutton , 1988 ) is used to reduce gradient variance at the expense of introducing bias ( Bahdanau et al. , 2016 ; Zhang et al. , 2017 ) . Several approaches explore beam search instead of sampling based methods ( Wiseman & Rush , 2016 ; Bunel et al. , 2018 ) , also at the expense of introducing bias . Several other methods combine the MLE and RL objectives for training ( Norouzi et al. , 2016 ; Ding & Soricut , 2017 ) .

Summary: The authors propose a new algorithm for unbiased stochastic gradient estimation for use in reinforcement learning of sequence generation tasks (specifically neural program synthesis and image captioning). The method consists in performing correlated Monte Carlo rollouts starting from each token in the generated sequence, and using the multiple rollouts to reduce gradient variance. An interesting property of the proposed algorithm is that the number of rollouts automatically scales with the uncertainty of the policy.