Title: ADAPTIVE CORRELATED MONTE CARLO FOR CON-TEXTUAL CATEGORICAL SEQUENCE GENERATION

Abstract: Sequence generation models are commonly refined with reinforcement learning over user-defined metrics. However, high gradient variance hinders the practical use of this method. To stabilize this method, we adapt to contextual generation of categorical sequences a policy gradient estimator, which evaluates a set of correlated Monte Carlo (MC) rollouts for variance control. Due to the correlation, the number of unique rollouts is random and adaptive to model uncertainty; those rollouts naturally become baselines for each other, and hence are combined to effectively reduce gradient variance. We also demonstrate the use of correlated MC rollouts for binary-tree softmax models, which reduce the high generation cost in large vocabulary scenarios by decomposing each categorical action into a sequence of binary actions. We evaluate our methods on both neural program synthesis and image captioning. The proposed methods yield lower gradient variance and consistent improvement over related baselines.

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 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(log 2 (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 := {x i , y i } N i=1 , our goal is to learn the conditional distribution of output y i given its context x i , expressed as p θ (y i | x i ). 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 = {y 1 , • • • , y T }, where y t ∈ {1, . . . , V }. A common way to model p θ (y | x) is to decompose it as p θ (y | x) = T t=1 p θ (y t | y 1:t-1 , x), where the t-th term in the product, which models the distribution of token y t conditioning on the context x and previously generated tokens y 1: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 θ (y t | y 1: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) E z∼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 θ (z t | z 1:t-1 , x), where p θ (z t | x, z 1:t-1 ) = Cat(σ(φ t )), φ t := T θ (x, z 1:t-1 ).

(1)

For a context-target pair {x, y}, we can expand the expected reward ER under policy p θ (z | x) as

EQUATION

where the partial-sentence reward is defined as r(z

1:t | x, y) = E z t+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 ∇ φ t ER∇ θ φ t , ∇ φ t ER = E z1:t-1∼p θ (• | x) ∇ φ t E zt∼Cat(σ(φ t )) [r(z 1:t | x, y)], ∇ φ t E zt∼Cat(σ(φ t )) [r(z 1:t | x, y)] = E zt∼Cat(σ(φ t )) [r(z 1:t | x, y)∇ φ t ln p(z t ; σ(φ 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(z 1: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(z 1:t-1 ))∇ φ t ln p(z t ; σ(φ t ))∇ θ φ t ,

where b(z 1:t-1 ) is a baseline function. To improve MIXER, Rennie et al. ( 2017) introduces the self-critic (SC) sequence training algorithm that sets b(z 1: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(z

1:t | x, y) = 1 K K k=1 r(z 1:t , z (k) (t+1):T | x, y), z (k) (t+1):T ∼ p θ (• | x, z 1:t ).

With these token-level rewards r(z 1:t | x, y), Liu et al. (2017) estimate the gradient as

∇θ ER = T t=1 (r(z 1:t | x, y) -b(z 1:t-1 ))∇ φ t ln p(z t ; σ(φ t ))∇ θ φ t .

(3)

Following SC, for token t, one may choose baseline b(z 1:t-1 ) = r(z 1:t-1 , zt:T | x, y), where zt:T is a greedy rollout following partial sentence z 1:t-1 . In addition to being more robust, using token-level rewards r(z 1: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) .

3. POLICY GRADIENT WITH ADAPTIVE CORRELATED MC ROLLOUTS:

Correlated MC samples, if well designed, can be combined to achieve much greater variance reduction than using the same number of independent ones (Owen, 2013) . To reduce gradient variance for contextual categorical sequence generation and remove the need to construct explicit baselines, we adapt the augment-REINFORCE-swap (ARS) and ARS-merge (ARSM) estimators of Yin et al. (2019) to generate correlated MC rollouts. The number of correlated MC rollouts, used to estimate each token-level partial-sequence reward, is adaptive according to the uncertainty on token generation.

The key idea here is to rewrite the gradient as differently expressed but equivalent expectations, whose MC samples, generated by sharing the same set of random numbers and hence correlated, are subsequently merged for variance reduction, without the need of learnable baseline functions.

Denote z ∼ Cat(σ(φ)) as a univariate categorical variable such that

P (z = v | φ) = σ(φ) v = e φv V i=1 e φi . Denote π = (π 1 , . . . , π V ) ∼ Dir(1 V

) as a random probability vector drawn from the Dirichlet distribution whose V parameters are all ones. Denoting m j as the operation of swapping the mth and jth elements of a vector, we have π m j m = π j , π m j j = π m , and

π m j i = π i , ∀i / ∈ {m, j}. With z := argmin i∈{1,••• ,V } π i e -φi , E(φ) = E c∼Cat(σ(φ)) [r(c)] can be reexpressed as E(φ) = E π∼Dir(1 V ) [r(z)],

whose gradient under the ARS estimator can be expressed as

∇ φv E(φ) = E π∼Dir(1V ) [g ARS (π, j) v ], g ARS (π, j) v := [r(z v j ) -1 V V m=1 r(z m j )](1 -V π j ), (4)

where j is a reference category randomly selected from {1, . . . , V } and z m j := argmin i∈{1,••• ,V } π m j i e -φi . ARSM further improves ARS by adding a merge step as

∇ φv E(φ) = E π∼Dir(1 V ) [g ARSM (π) v ], g ARSM (π) v := 1 V V j=1 g ARS (π, j) v .

(5)

We refer to z as the true action and z m j as pseudo actions. These actions are correlated to each other as they are transformed from the same Dirichlet distributed π vector under different pairwise index swaps. While there are V (V -1)/2 unique pairwise swaps, after MLE pre-train with V ∼ 10 4 , the number of unique pseudo actions that differ from the true action is random and often stays below 10, and becomes 0 more and more frequently as the progress of RL training reduces model uncertainty.

3.1. ADAPTIVE CORRELATED MC BASED POLICY GRADIENT FOR CATEGORICAL SEQUENCE:

Applying the ARSM estimator in (5) to the expected reward shown in (2), we have

∇ φ tv ER = E z 1:t-1 ∼p θ (• | x) E π t ∼Dir(1 V ) [g ARSM (π t ) v ], g ARSM (π t ) v := 1 V V j=1 g ARS (π t , j) v , g ARS (π t , j) v := r(z 1:t-1 , z v j t | x, y) -1 V V m=1 r(z 1:t-1 , z m j t | x, y) (1 -V π tj ),

where

z m j t := argmin i∈{1,••• ,V } π m j ti e -φti , π t ∼ Dir(1 V ), and r(z 1:t-1 , z m j t | x, y) = E z (t+1):T ∼p θ (• | x, z1:t-1,z m j t ) [r(z | x, y)].

We can therefore estimate each expectation using regular MC in the augmented space. Detailed formulations are deferred to Appendix B.1. Note if given π t , all pseudo actions z m j t are equal to true action z t , then g ARSM (π

t ) v = g ARS (π t , j) v = 0.

We note while the notation appears cumbersome, its implementation is not difficult, as described in Algorithm 2, Appendix C. The intuitive explanation of ARSM is that given π 1:T , it first generates true action sequence z 1:T (main trajectory) with z t = argmin i π ti e φti ; it then performs embarrassingly parallel MC rollouts for all unique pseudo actions that differ from their corresponding true actions: at step t, given the true actions z 1:t-1 , it generates pseudo actions z m j t , and for each unique value of them that differs from z t , it estimates its expected reward by rolling out a full sequence of length T ; and finally it combines the sampled rewards of the true action sequence and unique pseudo action sequences, which are correlated to each other, to achieve significant variance reduction. Despite significant gradient variance reduction, ARSM may become less efficient in computation when V becomes large (e.g., ∼ 10, 000). In the worst case, for each gradient estimate, it needs to generate as many as V -1 unique pseudo action sequences at each token; while in practice, the actual number is much smaller, it still could be large enough to cause computational issues, especially if the policy parameter is far from convergence. We note while in theory ARSM enjoys embarrassingly parallel computation for rolling out all unique pseudo action sequences, the acceleration via parallelization in practice is constrained by the capacity of our own computation platform. This motivates the following remedy. For large V , we choose K reference categories γ 1 , . . . , γ K , randomly sampled from {1, . . . , V } without replacement, to perform the swapping operations for pseudo action generation, and averaging over their corresponding ARS estimators as

EQUATION

We refer to this gradient estimator as the ARS-K gradient estimator. Whether this remedy could be successful depends on how large K needs to be as V increases. We find via experiments that the sufficient size of K grows slowly as V increases. For example, we will show in Section 4.2 that for the image captioning task with V = 9, 788, setting K = 5 already leads to competitive results.

Note during testing, regardless of whether using ARSM, ARS-K, or some other estimators, the categorical softmax output layer could become the computation bottleneck for random sequence generation. This motivates us to provide an algorithm to considerably reduce the generation cost during testing, though at the expense of reduced performance. We describe such a solution below.

3.2. BINARY-TREE-ARSM FOR COMPUTATIONAL RESOURCE LIMITED APPLICATIONS:

The conventional way to generate a word token is to sample from a V -way categorical distribution, whose probability parameters are obtained via a softmax output layer. This softmax output layer often becomes the computation bottleneck when V is large, making it difficult to be applied to resource-constrained environments, such as mobile devices. To mitigate this issue, related to the hierarchical softmax idea (Morin & Bengio, 2005; Grave et al., 2017; Goodman, 2001) , we first construct a binary tree to allocate each word of the vocabulary to one and only one leaf node of this tree. A simple solution is to perform binary hierarchical clustering of the words.

Denote e v as the word embedding vector of word v. In this paper, we use agglomerative clustering (Sibson, 1973 ) on e 1 , . . . , e V to recursively merge two closest clusters at a time until there is only one cluster. The root is linked to V leaf nodes via V overlapping root-to-leaf paths, each of which can be represented by a unique binary code b v of length D, where D = O(log 2 V ) is the depth of the tree. Note the V paths are not restricted to travel through the same number of nodes, but for simplicity we zero pad them to the same length. Both off-the-shelf embedding vectors (Pennington et al., 2014) and task-specific ones can be utilized. They provide useful prior information about the structure of the vocabulary, which we can exploit to facilitate our search within the action space.

With the binary tree, we transform the problem of choosing one out of V categories into that of making a There is a one-to-one mapping between the V root-to-leaf paths and V vocabulary words. We denote ν(b) ∈ {1, . . . , V } as the word that path b is mapped to, and β(v) ∈ {0, 1} D as the path that word v is mapped to. Note for a binary tree with V leaves, there will be V -1 non-leaf nodes, each of which needs a logit φ for its Bernoulli probability. Thus in total we need V -

1 logits φ 1 , • • • , φ V -1 .

The computational saving in generating categorical sequences comes from the fact that to generate a word token we need D φ's at most rather than all V -1 φ's. Therefore, with the binary tree, the computation for the softmax output layer to generate a token decreases from O(V ) to O(log 2 V ), which is significant especially for mobile applications. The binary-tree softmax model can be trained with MLE, or with the binary-tree-ARSM (BT-ARSM) gradient estimator introduced below.

For the binary case, both ARS and ARSM reduce to augment-REINFORCE-merge (ARM) (Yin & Zhou, 2019) , which expresses the gradient of

EQUATION

where b true := 1 [π<σ(φ)] and b sudo := 1 [π>σ(-φ)] are referred to as the true and pseudo actions, respectively. We note if we represent a V -way categorical variable as a sequence of D = O(log 2 V ) binary variables, the number of unique pseudo actions that differ from the true actions is at most D.

In the binary-tree setting, the conditional probability of generating token z t is changed from (1) to

p θ (z t | x, z 1:t-1 ) = Dz t l=1 Bernoulli b tl ; σ(φ t,b t(1:l-1)

) , (φ t1 , . . . , φ t(V -1) ) := T θ (x, z 1:t-1 ), (8) where (b t1 , . . . , b tDz t ) := β(z t ), φ t,b t(1:l-1) is the parameter of the non-leaf node at the end of the path defined by b t(1:l-1) , and D zt is the number of non-leaf nodes in the root-to-leaf path that leads to z t . Similar to the derivation in Section 3.1, we apply the ARSM gradient estimation to the decomposed binary sequences (BT-ARSM).

We provide the detailed formulations in Appendix B.2 and pseudo code in Algorithm 3, Appendix C. Intuitively, it first samples the true sequence of binary sequences {(b t1 , . . . , b tDz t )} T t=1 ; it then performs embarrassingly parallel MC rollouts for all pseudo actions that differ from their corresponding true actions: at step t, and depth l, given the previous true tokens z 1:t-1 , and true binary code b t1 , . . . , b t,l-1 , it generates pseudo binary code b (sudo) tl , and if it differs from b tl , then we estimate its expected reward by first rolling out a full binary code up to depth D and rolling out a full sequence up to length T ; and finally it combines the sampled rewards of the true action sequence and pseudo action sequences, which are correlated to each other, to achieve significant variance reduction. Note that if b (sudo) tl is the same as b tl , then the corresponding ARM gradient is zero.

4. EXPERIMENTS:

We evaluate our models with both neural program synthesis (NPS) and image captioning. 

m i ) = O m i ) for all m ∈ {1, . . . , M i }.

The evaluation metric is Generalization (Bunel et al., 2018) , defined as the proportion of the test instances

{I m i , O m i } m=1,Mi that satisfy λ i (I m i ) = O m i

for all m ∈ {1, ..., M i }. We evaluate on the Karel dataset (Devlin et al., 2017a) , consisting of 10, 000 training reference Karel programs1 with 2, 500 validation and 2, 500 test samples. Each program consists of a sequence of actions to move an agent inside a grid-world from one starting grid (input) to an end grid (output). The size of the action space V is 53 and average program length is around 20.

Baselines We incorporate five baseline algorithms in our evaluation. (i) MC-2 (Eq 3), using tokenlevel rewards and greedy baselines. (ii) MC-0, using sentence-level reward and token-level greedy baselines, which corresponds to the TD-SCST in Rennie et al. ( 2017). (iii) REINFORCE, using sentence-level reward and with mini-batch mean as the baseline. (iv) Self-Critic (SC) as in Rennie et al. ( 2017). (v) RL beam, the state-of-the-art method for NPS proposed by Bunel et al. (2018) to reduce the gradient variance while sacrificing the unbiasedness. The objective of RL beam is to maximize the expected reward under a distribution defined on a space constructed with beam search BS(p θ , S), where S = 64 is the beam size. Since the vocabulary size of V = 53 is not that large, we directly apply ARSM (i.e., ARS-53) policy gradient and compare it with the other methods.

We use the code of Bunel et al. (2018) as basis and use the same model architecture except for the exclusion of the optional grammar checker. The grammar checker, not available for all NPS tasks, helps adaptively reduce the search (action) space and hence simplifies optimization. Excluding the optional grammar checker eliminates its confounding influence on the core NPS task, making the comparison more generic and fair. We use greedy search for both testing and validation. All policy gradient based methods are fine-tuning a pre-trained (and converged) MLE model.

Results and analysis Figs. 1a and 1b plot against iteration the log variance, and average number of rollouts (including greedy rollouts used to construct baselines) per step for each method. We observe that ARSM overall has the smallest gradient variance, and at the beginning ARSM has more MC rollouts (unique pseudo actions) and hence takes relative longer time per iteration, but soon it becomes more and more confident (reflected as fewer and fewer pseudo actions per iteration) and turns faster. We note that the gradient variance at a given iteration is related to both the property of the gradient estimator and the parameter value at that iteration. Thus having smaller gradient variance may not necessarily imply better performance if different learning algorithms are not moving their parameters towards the same solution. This could help explain why SC has lower gradient variance than both MC-0 and MC-2 do but worse validation and test Generalization scores. Examining the performance on the training and validation sets suggests that REINFORCE and SC both diverge quickly; MC-0 stays around the starting point; MC-2 improves upon MLE initially, but then gradually diverges; RL beam reaches a good solution very fast but then gradually degrades towards worse solutions; and ARSM is the only one that makes steady improvement as the training progresses. Observing how the gap between training and testing evolves, we see evidence suggesting that RL beam overfits the training data, possibly due to the use of biased gradients, while ARSM does not. We note that there is no explicit regularization in ARSM. However, since ARSM tends to generate fewer and fewer unique pseudo actions as the policy becomes more and more confident, this adaptive characteristic may serve as an implicit regularization during the training process. Moreover, as the policy becomes more confident, the ARSM estimator has an increasing probability to yield MC gradient estimates that are exactly zeros, which may also help prevent overfitting as zero gradients will freeze the update of model parameters.

We summarize the validation and test Generation scores in Table 1 . Both MLE and RL beam (Bunel et al., 2018) perform reasonably well, but are outperformed by ARSM with a large margin. Even though MC-2 seems to improve upon MC-0 and SC, indicating the importance of using token-level rewards rather than sentence-level reward to guide the learning in this sparse reward scenario, it still clearly underperforms ARSM, which on average uses much fewer rollouts to estimate token-level rewards. This demonstrates the advantage of using an adaptive number of correlated MC rollouts over a fixed number of independent MC rollouts.

4.2. IMAGE CAPTIONING:

Image captioning, mapping an image x to a summary sentence y = (y 1 , . . . , y T ), has become a standard task to compare different policy gradient based RL methods. We conduct our experiments on the MS COCO dataset (Lin et al., 2014) , following the standard data split from Karpathy & Fei-Fei (2015) . We fine-tune a pre-trained MLE model using CIDEr score as the reward. Our implementation is based on Luo et al. (2018) . Details about the experimental setup can be found in Appendix D. ARS-K for computation-sufficient deployments We first investigate the effectiveness of the proposed method when it is computationally feasible to use the categorical softmax output layer at the test time. We consider MLE, REINFORCE, SC, and ARS-K with the same vocabulary size of V = 9788. For ARS-K, we experiment with several different K values. We report CIDEr score, and other commonly used metrics for the test set in Table 2 . We observe that while ARS-1 underperforms SC, ARS-K quickly improves as K increases: ARS-5 becomes comparable to SC in performance; ARS-10 and ARS-20 outperform SC by a large margin with statistical significance (standard error is about 0.2). The superior performance of ARS-K with large K is also evidenced by Fig. 3 . The gradient variance of ARS-20 is significantly lower than other algorithms (Fig. 3b upper ). In Fig. 3c (upper), we compare the average number of correlated MC rollouts of ARS-K for different K. While in theory the number of unique pseudo actions in ARS-K could be as many as V -1 at each step, it can be seen that after MLE pre-training, for each ARS-K (K = 1, 5, 10, 20), on average that number is small (fewer than 10 for V = 9488 and K = 20) and has an evident decreasing trend during training. Moreover, it increases slowly as K increases (clearly below a linear increasing rate).

Fig. 2 shows two pictures (see more plots in Appendix E.1) with their main sentences, which are greedily generated, and pseudo sentences generated by ARS-5 starting from the 7th token and 3rd token, respectively. Both greedily generated captions contain incorrect information about given images, while the pseudo sentences are semantically close to the greedy generations, however with interpretable variations in some details. Some pseudo sentences are better than the greedily generated captions. These pseudo-captions assembled together capture the nuance variations of the neighborhood of the generation, thus can serve as a good baseline to reduce the variance of the policy gradient. Note that there are more pseudo actions in the second plot, because the image is more complex and also there is more uncertainty at the beginning stage of generation (3rd token) compared to the latter stage of generation (7th token).

BT-ARSM for computation-limited deployments We evaluate the binary-tree ARSM (BT-ARSM) described in Section 3.2, which decomposes the action space to a sequence of binary actions.

(a) Binary tree constructions and MLE pretraining. We explore three different ways to construct binary trees over the action space: (i) tree WV : We apply agglomerative clustering to off-the-shelf pre-trained Word-to-Vector (WV) embeddings (Mikolov et al., 2013) to get a binary tree with a depth of 25; (ii) tree DIS: We follow tree WV except that we use the word embeddings pre-trained for image captioning with standard MLE objective and full vocabulary to distill (DIS) the full-softmax model's knowledge to the binary tree; (iii) tree RD : We randomly permute the leaves of tree DIS to produce a tree with no meaningful structure, referred as tree RD. We pre-train all these three models with MLE, and report the CIDEr score in Table 2 . Among these three binary trees, tree DIS performs the best, indicating that the tree structure has impact on its performance, and a task-specific pre-trained embedding is preferable when constructing a binary tree. As expected, comparing with the models trained using regular softmax layer (Table 2 ), the performance of the models with binary-tree softmax layers drop. However, the Multi-Adds softmax operations needed for generating a token is reduced by V /D (∼ 380 in our case) times, leading to a significant improvement in efficiency especially for deployment in computing resource limited scenarios at the cost of moderately degraded accuracy.

(b) Fine-tuning using BT-ARSM. We further fine-tune the pre-trained tree DIS model, with binarytree-REINFORCE (BT-RF), binary-tree-Self-Critic (BT-SC), and binary-tree-ARSM (BT-ARSM) respectively. Table 2 shows that BT-ARSM significantly outperforms the other two, which can be explained by the considerable variance reduction of BT-ARSM as is shown in Fig 3b (lower) . Notably, the performance of BT-ARSM is superior to vanilla softmax model trained with MLE even though it has been injected with strong inductive bias via binary-tree softmax to reduce its generation cost.

Adaptiveness of ARS-K and BT-ARSM As shown in Fig. 3 

5. CONCLUSION:

In this paper, we demonstrate the adaptation of ARSM policy gradient estimator, utilizing tokenlevel rewards of correlated Monte Carlo (MC) rollouts, to optimize contextual categorical sequence generation model. We apply the gradient estimators based on this idea to both the regular softmax model and binary-tree softmax model. The binary-tree softmax model has low cost for generating categorical tokens and hence is suited for computation-limited scenarios. We conduct empirical study on two challenging tasks: neural program synthesis and image captioning. Our observations verify that fewer and fewer correlated MC rollouts are conducted as the model becomes increasingly more certain during training. In addition, we show with correlated MC rollouts serving as baselines for each other, our methods show significant reduction of gradient variance and consistently outperform related baselines. We note that in a cold-start setting where we start from a complete random policy, it is still challenging to make our methods work efficiently as the number of pseudo actions may be too large if V is large. We consider it as future work to adapt our methods to this more challenging setting, where, to our best knowledge, little work has been done except for Ding & Soricut (2017) and d'Autume et al. (2019) .

D EXPERIMENTAL SETUP DETAILS D.1 IMAGE CAPTIONING:

Image captioning maps an image x to a sentence y = (y 1 , . . . , y T ) that summarizes the image information. It has become a standard task to compare different RL methods using policy gradient. A popular evaluation metric for this task is the CIDEr score (Vedantam et al., 2015) , which measures the similarity between the generated caption y and some reference ones. (Rennie et al., 2017; Anderson et al., 2018; Xu et al., 2015) . We conduct our experiments on the MS COCO dataset (Lin et al., 2014) that consists of 123,287 images. Each image has at least five captions. We use the standard data split from Karpathy & Fei-Fei (2015) , with 113,287 training, 5000 validation, and 5000 testing images.

The vocabulary size V is 9488 and the max caption length T is 16. For the model architecture, we employ a Fully-Connected (FC) model without attention (Rennie et al., 2017) . Image features are extracted from a pre-trained ResNet (He et al., 2016) . Our implementation is based on Luo et al. (2018) . We pre-train a model with MLE until convergence and use it for initialization. The CIDEr scores between the generated captions and references are used as rewards.

E QUALITATIVE RESULTS:

E.1 PSEUDO SENTENCES PRODUCED BY ARS-K ALGORITHM.

Main sentence: a man riding a motorcycle with a dog. Pseudo sentence 1: a man riding a motorcycle with hay on it. Pseudo sentence 2: a man riding a motorcycle with pack of sheep.

Main sentence: a group of people flying kites in a field. Pseudo sentence 1: a group of teenagers standing in a field flying kites.

:

The original dataset contains 1 million training instances.Bunel et al. (2018) proposed to reduce the dataset to 10, 000 examples and observed significant improvement of RL upon MLE when the reference program data is limited. Our experiments are based on the same reduced dataset.

Figures and Tables:

sequence of binary decisions b = (b 1 , . . . , b D ). If making l < D binary decisions (b 1 , . . . , b l ) has already led to a leaf node, then the sequence is terminated and b l+1 , . . . , b D all become zeros.

Figs. 1c and 1d plot the Generalization scores against training time on the training and validation sets. Due to large gradient variance, all methods except ARSM and RL beam either diverge or fail to improve the training objective. Examining the performance on the training and validation sets suggests that REINFORCE and SC both diverge quickly; MC-0 stays around the starting point; MC-2 improves upon MLE initially, but then gradually diverges; RL beam reaches a good solution very fast but then gradually degrades towards worse solutions; and ARSM is the only one that makes steady improvement as the training progresses. Observing how the gap between training and testing evolves, we see evidence suggesting that RL beam overfits the training data, possibly due to the use of biased gradients, while ARSM does not. We note that there is no explicit regularization in

Figure 1: From left to right are the comparisons of various methods in terms of gradient variance, number of sequence rollouts, training Generalization score, and validation Generalization score.

Figure 2: Main and pseudo trajectories for image captioning.

Figure 3: Comparison of different gradient estimators for image captioning task. "RF" denotes REINFORCEand "SC" denotes Self-Critic. Upper (lower) row: models using a regular softmax (binary-tree softmax).

and Fig.4(in Appendix A), our proposed methods can adaptively choose the number of correlated MC rollouts in four aspects: (i)(adapt across samples) in Fig.3d, we show the 2-D density estimation for the numbers of rollouts and the CIDEr scores of different samples during the later stage of training. We observe a statistically significant negative correlation (p < 0.05) between the numbers of rollouts and CIDEr scores, indicating that our algorithms can adaptively generate more rollouts for harder samples (lower CIDEr scores) and less rollouts for easier ones (higher CIDEr scores); (ii)(adapt across iterations) as shown in Fig.3c, during the training, the number of correlated MC rollouts decreases as the model improves and converges; (iii)(adapt across sentence positions) as shown inFigs. 4a, 4b, 4d, and 4e, more MC

Comparison of various algorithms in terms of the Generalization score on the Karel dataset. Mi , the goal is to learn a synthesizer parameterized by θ to generate a program λ i , which will produce a sequence of categorical actions to map input state I m i to output state O m

Table content: <table><tr><td>Methods</td><td>MLE</td><td>SC</td><td colspan="3">MC-0 MC-2 RL beam ARSM</td></tr><tr><td>Generalization (validation)</td><td>13.6</td><td colspan="2">12.51 12.64 13.56</td><td>14.76</td><td>17.07</td></tr><tr><td>Generalization (test)</td><td colspan="3">12.76 12.12 12.56 12.76</td><td>14.92</td><td>16.28</td></tr><tr><td>4.1 NEURAL PROGRAM SYNTHESIS</td><td/><td/><td/><td/></tr><tr><td colspan="6">NPS is a challenging representative task in contextual categorical sequence generation. First, the</td></tr><tr><td colspan="6">reward is only available after finishing the whole sequence. Second, the initial reward signals are</td></tr><tr><td colspan="6">often sparse because the generated programs rarely succeed in the beginning of training. We follow</td></tr><tr><td colspan="6">Bunel et al. (2018) to investigate an NPS task: for data sample i consisting of a set of input-output</td></tr><tr><td>states {I m i , O m i } m=1,</td><td/><td/><td/><td/></tr></table>

Performance comparison on the test set of COCO-caption dataset. rollouts appear at the timesteps in the middle range of the generated sequence, as they are associated with higher uncertainty; (iv)(adapt across depths) as shown in Figs.4c and 4f, for binary-tree softmax, the top layers (close to the root) of the tree are associated with more MC rollouts, since they are more uncertain about what to predict. More details are provided in Appendix A.

Table content: <table><tr><td>Method</td><td colspan="7">CIDEr BLEU-4 BLUE-3 BLEU-2 BLEU-1 ROUGE METEOR</td></tr><tr><td>Soft Attention (Xu et al., 2015)</td><td>-</td><td>24.3</td><td>34.4</td><td>49.2</td><td>70.7</td><td>-</td><td>23.9</td></tr><tr><td>Hard Attention (Xu et al., 2015)</td><td>-</td><td>25.0</td><td>35.7</td><td>50.4</td><td>71.8</td><td>-</td><td>23.0</td></tr><tr><td>Show &amp; Tell (Vinyals et al., 2015)</td><td>85.5</td><td>27.7</td><td>-</td><td>-</td><td>-</td><td>-</td><td>23.7</td></tr><tr><td>ATT-FCN (You et al., 2016)</td><td>-</td><td>30.4</td><td>40.2</td><td>53.7</td><td>70.9</td><td>-</td><td>24.3</td></tr><tr><td>SCN-LSTM (Gan et al., 2017)</td><td>101.2</td><td>33.0</td><td>43.3</td><td>56.6</td><td>72.8</td><td>-</td><td>25.7</td></tr><tr><td>Vanilla Softmax with MLE</td><td>93.3</td><td>30.4</td><td>40.2</td><td>53.6</td><td>70.7</td><td>52.2</td><td>24.7</td></tr><tr><td>REINFORCE</td><td>103.6</td><td>31.6</td><td>42.9</td><td>57.8</td><td>74.8</td><td>54.0</td><td>25.1</td></tr><tr><td>Self-Critic (Rennie et al., 2017)</td><td>106.5</td><td>32.3</td><td>43.8</td><td>58.7</td><td>75.6</td><td>54.6</td><td>25.6</td></tr><tr><td>ARS-1</td><td>103.6</td><td>30.8</td><td>42.4</td><td>57.8</td><td>75.4</td><td>54.0</td><td>25.2</td></tr><tr><td>ARS-5</td><td>106.1</td><td>31.6</td><td>43.3</td><td>58.6</td><td>76.1</td><td>54.6</td><td>25.5</td></tr><tr><td>ARS-10</td><td>107.7</td><td>31.8</td><td>43.5</td><td>58.8</td><td>76.0</td><td>54.7</td><td>25.7</td></tr><tr><td>ARS-20</td><td>108.4</td><td>32.1</td><td>43.8</td><td>59.0</td><td>76.4</td><td>54.8</td><td>25.8</td></tr><tr><td>tree RD with MLE</td><td>77.5</td><td>23.1</td><td>33.7</td><td>48.6</td><td>67.1</td><td>49.2</td><td>22.4</td></tr><tr><td>tree WV with MLE</td><td>80.2</td><td>23.3</td><td>34.1</td><td>49.3</td><td>68.0</td><td>49.9</td><td>22.9</td></tr><tr><td>tree DIS with MLE</td><td>84.2</td><td>24.9</td><td>35.1</td><td>49.7</td><td>67.6</td><td>50.5</td><td>23.7</td></tr><tr><td>tree DIS with BT-RF</td><td>93.6</td><td>28.8</td><td>39.7</td><td>55.0</td><td>72.5</td><td>52.2</td><td>23.7</td></tr><tr><td>tree DIS with BT-SC</td><td>96.5</td><td>29.4</td><td>40.5</td><td>55.4</td><td>72.7</td><td>52.8</td><td>24.1</td></tr><tr><td>tree DIS with BT-ARSM</td><td>99.2</td><td>29.9</td><td>41.2</td><td>56.2</td><td>73.3</td><td>52.7</td><td>24.3</td></tr></table>

Peer review (iclr2020):
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. The presentation of the paper must be improved (my score assumes that this will have been done). It is nice to have the detailed derivations when trying to dive deeper into the problem, but it hinders understanding of the main concepts at the first reading. Therefore, I would highly recommend to move most of the formulas to the appendix and keep instead only key ideas with intuitive explanations. It would also free some space in the main paper for the experiments from the appendix. Several questions on the technical side: 1. Is there any intuition why pseudo actions tend to be equal to the true one when learning progresses? What causes this? Might it enforce any structure (like uniform)? 2. When all pseudo actions are the same, the gradient is zero. In theory, zero gradient of a function corresponds to its extremum. Does it mean that when all pseudo actions are the same, an extremum is reached or is it just an artifact of this particular estimator? Can one prove any results of this kind? 3. I understand that the ARSM estimator should be unbiased for V = 2. Does the estimator remain unbiased when V > 2? 4. In the experiments, the variance is shown to reduce significantly which is nice. However, in theory, does the ARMS guarantee non-increasing variance or can it potentially go up in some cases? If it can, have it ever been observed in practice? 5. How does the runtime of the proposed algorithm compare to the competitors? Experiments: 1. Bunel et al (2018) report higher generalization on the Karel dataset. Is the difference due to the removal of the optional grammar checker? Can the same experiments be performed with this checker on or are there any constraints of the ARSM-based method? 2. The submitted code for the NPS experiment is actually the one by Bunel et al with their comments. I could not find any instructions or scripts reproducing the results of this paper (and I didn’t have much time to figure that out). One thing I wanted to check in the code is how the variance was computed for the plots? Minor: 1. p. 4, “expected award” >> “expected reward” 2. g_{ARSM} defined twice in (5) and in the beginning of p. 4 3. “Fig. 1 (left two) plots” and “Fig. 1 (right two) plots” not good