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 paper presents a novel reinforcement learning-based algorithm for contextual sequence generation. The algorithm builds on the previously proposed MIXER algorithm and improves it by integrating gradient estimates with lower variance (augment-REINFORCE-swap-merge). To further improve the runtime complexity of the proposed algorithm, binary tree-based hierarchical softmax is applied. The algorithm is evaluated on the Karel dataset for neural program synthesis and the MS COCO dataset for image captioning.