Year: 2020

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.

heading: 1 INTRODUCTION
text: 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).

heading: 2 PRELIMINARIES ON CONTEXTUAL SEQUENCE GENERATION
text: 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).

heading: 3 POLICY GRADIENT WITH ADAPTIVE CORRELATED MC ROLLOUTS
text: 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(1V ) 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 jm = πj , πm jj = πm, and π m j i = πi,∀i /∈ {m, j}.With z := argmini∈{1,··· ,V } πie−φi , E(φ) = Ec∼Cat(σ(φ))[r(c)] can be reexpressed as E(φ) = Eπ∼Dir(1V )[r(z)], whose gradient under the ARS estimator can be expressed as
∇φvE(φ) = Eπ∼Dir(1V )[gARS(π, j)v], gARS(π, j)v := [r(zv j )− 1V ∑V m=1 r(z m j )](1− V πj), (4)
where j is a reference category randomly selected from {1, . . . , V } and zm j := argmini∈{1,··· ,V } π m j i e −φi . ARSM further improves ARS by adding a merge step as
∇φvE(φ) = Eπ∼Dir(1V )[gARSM(π)v], gARSM(π)v := 1V ∑V j=1 gARS(π, j)v. (5)
We refer to z as the true action and zm 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 ∼ 104, 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.

heading: 3.1 ADAPTIVE CORRELATED MC BASED POLICY GRADIENT FOR CATEGORICAL SEQUENCE
text: Applying the ARSM estimator in (5) to the expected reward shown in (2), we have ∇φtvER = Ez1:t−1∼pθ(· |x)Eπt∼Dir(1V )[gARSM(πt)v], gARSM(πt)v := 1V ∑V j=1 gARS(πt, j)v,
gARS(πt, j)v := [ r(z1:t−1, z v j t |x,y)− 1V ∑V m=1 r(z1:t−1, z m j t |x,y) ] (1− V πtj),
where zm jt := argmini∈{1,··· ,V } π m j ti e −φti , πt ∼ Dir(1V ), and r(z1:t−1, zm jt |x,y) = Ez(t+1):T∼pθ(· |x, z1:t−1,zm jt )[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 zm jt are equal to true action zt, then gARSM(πt)v = gARS(π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 z1:T (main trajectory) with zt = argmini πtie
φ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 z1:t−1, it generates pseudo actions zm jt , and for each unique value of them that differs from zt, 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
gARS-K(π)v = 1 K ∑K j=1 gARS(π, γj)v. (6)
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.

heading: 3.2 BINARY-TREE-ARSM FOR COMPUTATIONAL RESOURCE LIMITED APPLICATIONS
text: 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 ev as the word embedding vector of word v. In this paper, we use agglomerative clustering (Sibson, 1973) on e1, . . . , eV 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 bv of length D, where D = O(log2 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 sequence of binary decisions b = (b1, . . . , bD). If making l < D binary decisions (b1, . . . , bl) has already led to a leaf node, then the sequence is terminated and bl+1, . . . , bD all become zeros. 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(log2 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 Eb(φ) = Ez∼Ber(σ(φ))[r(z)], σ(φ) = 1/(1 + e−φ), as
∇φEb(φ) = Eπ∼Uniform(0,1)[gARM(π)], gARM(π) := [r(btrue)− r(bsudo)](1/2− π), (7) where btrue := 1[π<σ(φ)] and bsudo := 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(log2 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 zt is changed from (1) to pθ(zt |x, z1:t−1) = ∏Dzt l=1 Bernoulli ( btl;σ(φt,bt(1:l−1)) ) , (φt1, . . . , φt(V−1)) := Tθ(x, z1:t−1), (8)
where (bt1, . . . , btDzt ) := β(zt), φt,bt(1:l−1) is the parameter of the non-leaf node at the end of the path defined by bt(1:l−1), and Dzt is the number of non-leaf nodes in the root-to-leaf path that leads to zt. 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 {(bt1, . . . , btDzt )} 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 z1:t−1, and true binary code bt1, . . . , bt,l−1, it generates pseudo binary code b (sudo) tl , and if it differs from btl, 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 btl, then the corresponding ARM gradient is zero.

heading: 4 EXPERIMENTS
text: We evaluate our models with both neural program synthesis (NPS) and image captioning.

heading: 4.1 NEURAL PROGRAM SYNTHESIS
text: NPS is a challenging representative task in contextual categorical sequence generation. First, the reward is only available after finishing the whole sequence. Second, the initial reward signals are often sparse because the generated programs rarely succeed in the beginning of training. We follow Bunel et al. (2018) to investigate an NPS task: for data sample i consisting of a set of input-output states {Imi , Omi }m=1,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 Imi to output state O m i (i.e., λi(Imi ) = O m i ) for all m ∈ {1, . . . ,Mi}. The evaluation metric is Generalization (Bunel et al., 2018), defined as the proportion of the test instances {Imi′ , Omi′ }m=1,Mi that satisfy λi′(Imi′ ) = Omi′ for all m ∈ {1, ...,Mi}. 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.
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
1The 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.
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.

heading: 4.2 IMAGE CAPTIONING
text: Image captioning, mapping an image x to a summary sentence y = (y1, . . . , yT ), 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 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 in Figs. 4a, 4b, 4d, and 4e, more MC
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.

heading: 5 CONCLUSION
text: 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).

heading: ACKNOWLEDGMENTS
text: This research was supported in part by the U.S. National Science Foundation under Grant IIS-1812699 and the McCombs Research Excellence Grant. The authors acknowledge the support of NVIDIA Corporation with the donation of the Titan Xp GPU used for this research, and the computational support of Texas Advanced Computing Center.

heading: A ADAPTIVENESS OF ARS-K/BT-ARSM
text: 1. Adaptiveness across samples: In Fig. 3d, we show the 2-D density estimation for the numbers of rollouts and the CIDEr scores for different samples during the late stage of training. We observe a statistically significant negative correlation (p < 0.05) between the number 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 samples (higher CIDEr scores).
2. Adaptiveness across time steps: As shown in Figs. 4a, 4b, 4d, and 4e, it is reasonable that there will be more MC rollouts at the middle time step of the generated sequence, since the initial words and end words are more easily to be learned due to their cardinality, and ARSM model will be more confident about the choices for these words, leading to fewer MC rollouts.
3. Adaptiveness across depths: For binary softmax model, the terms in the vocabulary are represented as leaf nodes in the tree. As shown in Figs. 4c and 4f, ASRM successfully captures the adaptiveness of MC rollouts across depths in the tree. 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.
4. Adaptiveness across iterations: In Figs. 1b and 3c, we observe that, during training, the number of correlated MC rollouts decreases as the model improves and converges.

heading: B DETAILED FORMULATIONS
text: B.1 ARS/M GRADIENT ESTIMATORS
Applying the ARSM estimator in (5) to the expected reward shown in (2), we have ∇φtvER = Ez1:t−1∼pθ(· |x)Eπt∼Dir(1V )[gARSM(πt)v], gARSM(πt)v := 1V ∑V j=1 gARS(πt, j)v,
gARS(πt, j)v := [ r(z1:t−1, z v j t |x,y)− 1V ∑V m=1 r(z1:t−1, z m j t |x,y) ] (1− V πtj),
where zm jt := argmini∈{1,··· ,V } π m j ti e −φti . Thus we can approximate the gradient∇θER as
∇̂θER = ∑T t=1 ∑V v=1 ĝARSM(πt)v∇θφtv, ĝARSM(πt)v := 1 V ∑V j=1 ĝARS(πt, j)v,
ĝARS(πt, j)v := [ r̂(z1:t−1, z v j t |x,y)− 1V ∑V m=1 r̂(z1:t−1, z m j t |x,y) ] (1− V πtj), (9)
where π1, . . . ,πT iid∼ Dir(1V ); r̂(z1:t−1, zm jt |x,y) is an approximation of r(z1:t−1, z m j
t |x,y) = Ez(t+1):T∼pθ(· |x, z1:t−1,zm jt )[r(z |x,y)], which can be estimated with r(z1:t−1, zm jt , z̃t+1:T |x,y), where z̃t+1:T ∼ pθ(· |x, z1:t−1, z m j t ) is an MC rollout.
B.2 BT-ARSM GRADIENT ESTIMATORS
∇φt,bt(1:l−1) ER = Ez1:t−1∼pθ(· |x)Ebt(1:l−1)∼pθ(· |x,z1:t−1)Eπtl∼Unif(0,1)[gARM(πtl)], gARM(πtl) = [ r ( z1:t−1, bt(1:l−1), b (true) tl |x,y ) − r ( z1:t−1, bt(1:l−1), b (sudo) tl |x,y )] (1/2− πtl),
b (true) tl := 1[πtl<σ(φt,bt(1:l−1) )], b (sudo) tl := 1[πtl>σ(−φt,bt(1:l−1) )], (10)
where r(z1:t−1, bt(1:l) |x,y) := Ebt(l+1:Dzt ),z(t+1):T∼pθ(· |x, z1:t−1, bt(1:l))[r(z |x,y)]. Thus we have ∇θER = ∑T t=1 ∑Dzt l=1 ∇φt,bt(1:l−1) ER∇θφt,bt(1:l−1) , which can be approximated with ∇̂θER =∑T
t=1 ∑Dzt l=1 ĝARM(πtl)(1− 2πtl), where ĝARM(πtl) approximates gARM(πtl) shown in (10) via MC
integration; note if given πtl∼Unif(0,1), b (sudo) tl = b (true) tl , then gARM(πtl) = 0.

heading: C ALGORITHMS
text: Efficient Pseudo-Action Computation: To implement ARS-K, we need to compute the corresponding pseudo actions for each reference category j in a reference set J . In other words, given π,φ, we need to compute zm j = argmin{i=1,...,V } lnπ m j i − φi, for all m ∈ {1, . . . , V }, j ∈ J . A naive implementation would involve taking the minimum over V -dimensional vectors for V · |J | times, which is computationally expensive when V is large. In the following, we take advantage of the correlation among pseudo actions and propose an efficient algorithm which only involves taking the minimum over V -dimensional vectors |J | times. We first explain the notations and the basic ideas, and then present the algorithm in Algorithm 1.
Let oi,j = lnπi−φj . Denote m1,m2 as the indexes for the top 2 smallest in {oi,i}i=1:V respectively (this is the only place where we need to take the minimum over whole V -dimensional vectors). Let ID(oi,j) denote the function returning the second index of oi,j , which means
ID(oi,j) = j.
In the following, we show that it is not necessary to take minimum over whole V -dimensional vectors for each pseudo actions. We can compute zm j efficiently by observing:
1. If m1 /∈ {m, j}, the smallest value in lnπm j − φ can only be among the updated two values (oj,m, om,j) and the original smallest value om1,m1 . Hence, the index of the smallest value is
zm j = ID(min(min(oj,m, om,j), om1,m1)).
2. If m1 ∈ {m, j}, om1,m1 will be updated to a new value, therefore, the above equation does not hold. But, in the following, we show we can use the second smallest value instead. The index of the smallest value will be
zm j = ID(min(min(oj,m, om,j), om2,m2))
Proof. Assume m1 ∈ {m, j}. If m2 /∈ {m, j}, since om1,m1 has been changed, om2,m2 will be the smallest value in the vector excluding the two updated values. Hence, the smallest value will be among the three values {oj,m, om,j , om2,m2}, among which we can find the index of the smallest item. If m2 ∈ {m, j}, then {m, j} becomes {m1,m2}. Hence, om1,m1 and om2,m2 will be updated. For any m /∈ {m1,m2}, min(om2,m1 , om1,m2) ≤ om2,m2 ≤ om,m. Therefore, the smallest value will be among {om2,m1 , om1,m2}. The index of the smallest value will be ID(min(om2,m1 , om1,m2)), which is equivalent to ID(min(min(oj,m, om,j), om2,m2)).
Algorithm 1: Compute Pseudo-Action Matrix for Reference Category Set J in Parallel input : Batched π and φ, Ref-Cat Set J , ID Function output : Pseudo-Action Matrix P ; Compute om1 ,m1, om2 ,m2 = Top2(lnπ − φ); for j ∈ J,m ∈ {1, . . . , V } (in parallel) do
Compute oj,m = lnπj − φm Compute om,j = lnπm − φj
end for Initialize P with size (|J |, V ) for j ∈ J,m ∈ {1, . . . , V } (in parallel by using index matrix) do
if m1 ∈ {j,m} then P [j,m] = ID(min(min(oj,m, om,j), om1,m1)) else P [j,m] = ID(min(min(oj,m, om,j), om2,m2))
end if end for
Algorithm 2: ARS-K/ARSM(K = V ) policy gradient for fine-tuning a contextual categorical sequence generation model with a discrete-action space of V actions. input : MLE pre-trained policy parameter θ, number of reference category K, main trajectory sample type mt, pseudo trajectory sample type pt output : Fine-tuned policy parameter θ while not converged do
Receive random sample x,y; First, we sample a main trajectory (z1, . . . , zT ): if mt = ’greedily sample’ then
for t = 1 : T , let zt = argmini∈{1,...,V }(−φti), where φt = Tθ(z1:t−1,x); else
for t = 1 : T , let zt = argmini∈{1,...,V }(lnπti − φti), where πt ∼ Dirichlet(1V ) (or let πti = − ln(Unif(0, 1))), and φt = Tθ(z1:t−1,x);
end if Second, we compute pseudo actions: for t = 1 : T do
Let πt ∼ Dirichlet(1V ) (or let πti = − ln(Unif(0, 1)) for i = 1, . . . , V and then normalize them to have a unit norm); Let j1, ..., jK be K reference categories randomly sampled from {1, . . . , V } without replacement; for k = 1, . . . ,K, v = 1, . . . , V (in parallel) do Let zv jkt := argmini∈{1,...,V }(lnπ v jk ti −φti) as the (v, k)th pseudo action; end for Let St = unique({zv jt }v,j) which means St is the set of all unique values in {z v j t }v,j .
Denote the cardinality of St as |St|, where 1 ≤ |St| ≤ V ; end for Third, we complete sentences and evaluate the reward for the unique set of pseudo actions: for t = 1 : T (in parallel) do
if |St| = 1 then continue end if for z̃ts ∈ St (in parallel) do
if t < T then if pt = ’greedily sample’ then
greedily sample zst+1:T ∼ pθ(zt+1:T | z1:t−1, z̃ts,x) else
randomly sample zst+1:T ∼ pθ(zt+1:T | z1:t−1, z̃ts,x) end if
end if end for for v = 1 : V, k = 1 : K (in parallel) do
Let f(zv jkt ) = r(z1:t−1, z̃ts, z s t+1:T |x,y) if z v jk t = z̃ts;
end for end for Finally, we compute the ARSM gradients and update parameters: for t = 1 : T, k = 1 : K (in parallel) do
Let f̄tk = 1V ∑V v=1 f(z v jk t ); for v = 1 : V (in parallel) do Let gtk,v = 1K ( f(zv jkt )− f̄tk ) (1− V πtjk), where gtk,v is the vth component of gtk;
end for end for for t = 1 : T (in parallel) do
θ = θ + ηθ∇θφtgt, where gt = ∑ k gtk with step-size ηθ
end for end while
Algorithm 3: Binary-tree-ARSM policy gradient for fine-tuning a binary-tree contextual categorical sequence generation model. input : MLE pre-trained policy parameter θ, binary code to word mapping ν output : Fine-tuned policy parameter θ; while not converged do
Receive random sample x,y; First, we sample a main trajectory (z1, . . . , zT ): for t = 1 : T do
for d = 1 : D do φt,bt(1:d−1) = Tθ(z1:t−1,x)ν(bt(1:d−1)) Sample πtd ∼ Uniform(0, 1); Let btd = 1[πtd<σ(φt,bt(1:d−1) )]; end for zt = ν(bt(1:D))
end for Second, we compute pseudo actions: for t = 1 : T (in parallel) do
for d = 1 : D do Let b(1)td = btd; Let b(2)td = 1[πtd>σ(−φt,bt(1:d−1) )];
if b(1)td 6= b (2) td then
If d < D, sample b(j)t(d+1:D) ∼ pθ(bt(d+1:D) | z1:t−1, bt,1:d−1, b (j) td ,x), j = 1, 2; Let z(j)td = ν(b (j) t(1:D)), j = 1, 2;
end if end for Let St = unique({z(j)td }d,j) which means St is the set of all unique values in {z (j) td }d,j .
Denote the cardinality of St as |St|, where D ≤ |St| ≤ 2 ∗D. end for Third, we complete sentences and evaluate the rewards for the unique set of pseudo actions: for t = 1 : T (in parallel) do
for z̃ts ∈ St (in parallel) do If t < T , sample zst+1:T ∼ pθ(zt+1:T | z1:t−1, z̃ts,x); end for for d = 1 : D, j = 1 : 2 (in parallel) do
Let f (j)td = r(z1:t−1, z̃ts, z s t+1:T |x,y) if z (j) td = z̃ts;
end for We compute the ARSM gradients and update parameters: for d = 1 : D (in parallel) do
if b(1)td 6= b (2) tl then
Let gt,bt(1:d−1) = 1 2 (f (1) td − f (2) td )(1− 2πtd);
θupdate = ηθ∇θφt,bt(1:d−1)gt,bt(1:d−1) , with step-size ηθ end if θ = θ + θupdate
end for end for
end while

heading: D EXPERIMENTAL SETUP DETAILS
text: D.1 IMAGE CAPTIONING
Image captioning maps an image x to a sentence y = (y1, . . . , yT ) 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.

heading: E QUALITATIVE RESULTS
text: 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.
Main sentence: a man and woman are standing in a of a table. Pseudo sentence 1: a man and woman are standing in an market. Pseudo sentence 2: a man and woman are standing in the street. Pseudo sentence 3: a man and woman are standing in to a tent. Pseudo sentence 4: a man and woman are standing in front of a table.
Main sentence: a city that has a large white building on it. Pseudo sentence 1: a city with a red traffic and a large building. Pseudo sentence 2: a city intersection with a traffic light and a street sign. Pseudo sentence 3: a city bus is driving down the street. Pseudo sentence 4: a city street with a city street with cars parked on it. Pseudo sentence 5: a city road with a traffic light and a street sign.

Review summaries:
The paper presents experimental results on the application of the gradient ARSM estimator of Yin et al. (2019) to challenging structured prediction problems (neural program synthesis and image captioning). The authors also propose two variants, ASR-K which is the ARS estimator computed on a random sample of K (among V) labels, as well as a binary tree version in which the V values are encoded as a path in a binary tree of depth O(log(V)), effectively increasing the length of sequences to be predicted but reducing the action space at each tilmestep.

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.

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.