Adaptive Correlated Monte Carlo for Contextual Categorical Sequence Generation
Authors: Xinjie Fan
Yizhe Zhang
Zhendong Wang
Mingyuan Zhou
Venue:  (2020)
Keywords: binary softmax
discrete variables
policy gradient
pseudo actions
reinforcement learning
variance reduction

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. 

Published as a conference paper at ICLR 2020
ADAPTIVE CORRELATED MONTE CARLO FOR CON- TEXTUAL CATEGORICAL SEQUENCE GENERATION
Xinjie Fan', Yizhe Zhang?, Zhendong Wang', Mingyuan Zhou'
Univeessity Texas at Austin, "Microsoft Research, *Columbia University xfan@utexas edu, yizhe zhang malangmicce crosoft.......mmmmmmmmmmoooooo.. com, zw2533@columbia edu, mingyuan. www.ayyan..huuuueeeee zhou@mccombs. wtw.aas...... edu
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 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.
INTRODUCTION
Contextual categorical sequence generation is modeling component in 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 al. 2015 Xu et 2015).
Typically, an encoder-decoder framework is applied. The encoder maps a contextual input to latent representation, conditioning on which and previously generated tokens the decoder generates categorical tokens in 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. 2013 Ranzato et al. 2016). which refers to the discrepancy between training and generation using the Teacher Forcing (Williams & Zipser 1989) strategy, i.c., 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, is common to refine contextual sequence generation model pre-trained with MLE under the reinforcement learning (RL) framework (Zarembal
Sutskever 2015; Ranzato et 2016 Bahdanauu 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 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 discriminator language model (Yang al. 2018 Yu et al. 2017 Lamb et al. 2016 Caccia et al. 2018: Autume et al.) 2019). The RL training
Code link
https:/gtthuh.ccm/xiiifaaaaCCCC ICLR

---

Published as a conference paper at ICLR 2020
enables direct improvement of the user-defined learned reward. Moreover, in cases where only weak-superrision is available, e.g., in neural program synthesis, RL may considerably improve the model performance (Bunel et al. 2018 Zhong et al. 2017p. However, the gradients of the expected reward in RI. 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 reparameterizattin trick (Kingma & Welling 2013 Rezende et al., 2014). which, however, is applicable to neither discrete variables nor non- con-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) [2033] Williams 1992, Paisley et al. 2012 Ranganath 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 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 token from large vocabulary via the softmax output layer is often computationally heavy. This prevents a categorical sequence generation model from being deployed 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, is so far unclear how 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 Segmen--EINFFORCE-waaaa (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 categorical softmax output layer, with the generation of 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,
adapt the augment- mugment-REINFORCE REINFORCE--errge (ARM) estimator (Yin & Zhou 2019), which provides unbiased, low-variance gradients for binary variables, to backpropagate the gradients through the sequence binary sequences. Under this binary-tree construction, the cost of generating categorical token reduces from O(V) to Ollogz()) 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).
PRELIMINARIES ON CONTEXTUAL SEQUENCE GENERATION
For dataset of context-output pairs D = {x; yr J our goal is to learn the conditional distri- bution of output y; given its context æ;, expressed as po xz)). Below we drop the data index subscript for brevity. We focus on the case that an output is a sequence of T categorical vari- able as y {ur.- .Vr), where Yt E fI....,, V). A common way to model ( x)) is to decompose it as (y æ) II Po a)), where the t-th term in the product, which models the distribution of token Ye conditioning on the context æ and previously generated to- kens V1-t-1, is commonly parameterized by recurrent neural network (Sutskever et al.) 2014). MLE is a common way to train the model: θ argmaxg E [(r.y))~ummxx..) [log Po x)]. Viewing (e x)) as 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
0R.. = argmax, E ((..)))-pa..t..) Ez-pe (: (=) [r(2 x. y)], where r(x x. 3) denotes the reward of the generated (hypothesis) sequence z given the context and the reference target sequence y. For
2

---

Published as a conference paper at ICLR 2020
example, for image captioning, the reward could be the CIDEr score that measures the similarity between the generated caption & and the reference y (Rennie et al. 2017).
Denote (-) the softmax function and (- deterministic function defined by deep neural network with parameter 0. We model po 2 æ) II== (0 21.--1 where
(2 E, 11: C))( o, To (z, 21:--)).
(1)
For context-target pair fæ, y), we can expand the expected reward ER under policy po 2x)
ER = E- (r m. y)] = E P 2 (ato a y)],
(2)
where the partial-sentence reward is defined as r(= 22, y) = E Exp (1a w)
Using the chain rule and REINFORCE (Williams 1992) estimator, have
V VØ = E V V Φ ER = E... 1-pe(- an) Vo, Ez Cattrt v)),
V (o (r x, y) E Eu,-Ca(o(
x, y) V In (
The main challenge here is control the variance in estimating V VoER. A variety methods have been proposed. For example, drawing sentence 2 (o x)) and using its reward r(z x. y) to approximate all partial-sentence rewards {r(#1: æ, y) h,r, Ranzato al. (2016) introduce MIXER with scheduled training iterating between MLE and RL, estimating the RL gradient as
VØER= = Ei=(((z) y) b(2 In p(2;; a (o,))
where b(21;-1) is baseline function. To improve MIXER, Rennie al. (2017) introduces the self-critic (SC) sequence training algorithm that sets 6(21=-1) r(= x. y) for all 4, where & is a greedy sequence rollout under po æ)). As using sentence-level reward r(= ac, y) to guide the learning is found to be sensitive to the algorithm parameters such as the learning rate, [Lu et al. (2017) follow Yu et al.) (2017) to use token-level rewards that approximate each partial-sentence reward with (k)
K independent MC rollouts K) as f(zize ac, y) = K E r(21,; (1-1):: x. y). (((+1):T ~
p- 2, 21.)). With these token-level rewards f(21: x. y) Liu et al. 2017 estimate the gradient
VØER E (f(za x y) 6( V s, Inp(a; (0)))
(3)
Following SC, for token t, one may choose baseline b( ( r(#10-1, žcr x. y). where ă.r is greedy rollout following partial sentence -1:--1. In addition to being more robust, using token-level rewards f(21=t) to guide the learning is often necessary when the reward signal is sparse, e.g., in neural program synthesis, it is common that full MC roll receives zero reward with high probability.
In addition to these methods mentioned above, the actor-critic method (Sutton 1988 used to reduce gradient variance at the expense of introducing bias (Bahdanau et al. 2016 Zhang 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).
POLICY GRADIENT WITH ADAPTIVE CORRELATED MC ROLLOUTS
Correlated MC samples, if well designed, can 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 Dugment-REINFORCCE
(ARS) and ARS-merge (ARSM) estimators of Yin 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(a as a univariate categorical variable such that P(z v Φ) (Φ)). e //V/ c Denote π = (T....., .Ty ~ Dir(ly) as a random probability vector drawn from the Dirichlet distribution whose V parameters are all ones. Denoting j as the operation of
3

---

Published as a conference paper at ICLR 2020
swapping the mth and jth elements of a vector, we have am am Jz. Am and x;
T4, Vi & {m,j). With = = argmin .Y Tee Tze=", E(Φ) = (c-CCa(((()) [r(c)] can be reexpressed
as E {(Φ) = B [+()], whose gradient under the ARS estimator can expressed as
Vo. {()) Di( (TA , J)]] DARS( f). (r(--)))
722== ][( Vny), (4)
where is reference category randomly selected from (1....,, V} and 2mm
argminie( (I.... e e ARSM further improves ARS by adding a merge step as
Vo. E(Φ) = Di (T))sM(T)))) (ARSM((T), :- + E (T)
(5)
We refer to 2 as the true action and zm~ 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 unique pairwise swaps, after MLE pre-train with V ~ 10*, 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 ( to the expected reward shown in ( we have
V
ER E Ex
| E8 1 Dir(1v) [AARSM (T)))) JARSM (Tx) - - JARS (T++;
V
1x, 3) (1 VTte), x,,)) where 27 amm = argmin 1.... v/ Tii m 0es, Te ~ Dir(Iv), and r(a1=-1,, = y) za na -0. - [r( x, y)]. We can therefore estimate each expectation using reg-
gARS TT. J). [r(21:11, 2 x. + " ((112-1,27
ular MC in the augmented space. Detailed formulations are deferred to Appendix B. Note if given Te, all pseudo actions 2," are equal to true action 2,, then 9ARSM ( ). = 9ARS (T: = O.
We note while the notation appears cumbersome, its implementation is not difficult, as described in Algorithm[2 2 Appendix C The intuitive explanation ARSM is that given π I.T, it first generates true action sequence 21. (main trajectory) with 2₁ argmin, The it then performs embarrassingly parallel MC rollouts for all unique pseudo actions that differ from their corresponding true actions: at step f, given the true actions 21:--1, it generates pseudo actions 2, 2m= and for each unique value of them that differs from z, 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 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 the policy parameter is far from convergence. We note while in theory ARSM enjoys embarrass- ingly 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..... YK- randomly sampled from {I...., V} without replacement, to perform the swapping operations for pseudo action generation, and averaging over their corresponding ARS estimators as
gARSk((T) (T) x Ex DARS(TT 75)0-
(6)
We refer this gradient estimator as the ARS-K gradient estimator. Whether this remedy could be successful depends on how large K needs to be 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 considcrably reduce the generation cost during testing, though at the expense of reduced performance. We describe such solution below.
4

---

Published as a conference paper at ICLR 2020
3.2 TREAA--REE--ARM 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 softmax output layer. This softmax output layer often becomes the computation bottleneck when V is large, making it difficult to be applied to shource-constrained environments. such as mobile devices. To mitigate this issue, related to the hierarchical softmax idea (Morin & Bengio [005) Grave et al. 2017; Goodman (00T), we first construct a binary tree to allocate each word 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, as the word embedding vector of word v. In this paper, we use agglomerative clustering (Sibson 1973) on eI....: Ey to recursively merge two closest clusters at time until there is only one cluster. The root linked V leaf nodes via overlapping root-to-leaf paths, each of which can be represented by unique binary code b. of length D. where D (logg 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 eet 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 sequence of binary decisions b (bi, (bi,...., bp). If making 1 < D binary decisions (bI,..., b)) has already led to a leaf node, then the sequence is terminated and br+1,--...hp hn all become zeros. There a one-to-one mapping between the V root-to-leaf paths and V vocabulary words. We denote v(b) E {l...., } as the word that path b is mapped to, and B(v) E 10,11 1) as the path that word is mapped to. Note for binary tree with V leaves, there will be V 1 non-leaf nodes, each of which needs logit 6 for its Bernoulli probability. Thus in total we need V 1 logits $1,-. , -1- The computational saving in generating categorical sequences comes from the fact that to generate word token we need D ø's at most rather than all V V-1 1 Ø's. Therefore. with the binary tree, the computation for the softmax output layer to generate a token decreases from O(V) to Óllog, 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 wugment-REINFFRREE... Zhou 2019), which expresses the gradient of (d) = E (e(o((())))) 1/(1+ x),
(ARM) Yin &
as
VoE((0) = -Un [_ARM (4)], 9ARM (T) == (r(bunue - rlbsmdo)](1122 x),
where Diue := 1 [x)ca(()] and bundo == 1 [=)))(--)] are referred to as the true and pseudo actions, respectively. We note if we represent a V-way categorical variable as sequence of D = Ologg 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 24 is changed from to
pe xx, 11: P=: Bernouli (oubaaaa (o....... 1)) == To (xx, -1:111,, (8)
where Ubi.,...,,,,,... ) == B(2x), (1,b411:-1) is the parameter of the non-leaf node at the end of the path defined by b((1:t1)) and Dz, is the number of non-leaf nodes in the root-to-leaf path that
leads to 2t. Similar to the derivation Section 3.1 we apply the ARSM gradient estimation to the decomposed binary sequences (BT-ARSM).
We provide the detailed formulations Appendix B.2 and pseudo code Algorithm 3 Appendix C Intuitively, it first samples the true sequence of binary sequences {(br)..... b.D., ) it then per- forms embarrassingly parallel MC rollouts for all pseudo actions that differ from their corresponding true actions: at step f, and depth I, given the previous true tokens 21:--1, and true binary code be....., b,.1-11 generates pseudo binary code ((5ud0) and if it differs from bus 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 ((sud)) is the same as by, then the corresponding ARM gradient is zero ftl
4 EXPERIMENTS
We evaluate our models with both neural program synthesis (NPS) and image captioning.
5

---

Published as a conference paper at ICLR 2020
Table 1: Comparison of various algorithms in terms the Generalization score on the Karel dataset.
Methods
MLE
SC
MC-0
MC-2
RL.beam
ARSM
Generalization (validation)
13.6
12.51
12.64
13.56
14.76
17.07
Generalization (test)
12.76
12.12
12.56
12.76
14.92
16.28
4.1 NEURAL PROGRAM SYNTHESIS
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 { O Im=.,,,, the goal is to learn synthesizer parameterized by 0 to generate a program As, which will produce a sequence of categorical actions to map input state I," to output state O (i.e., A = O!Th) for all m E {1..... M, J. The evaluation metric is Generalization (Bunel et al. 2018). defined as the proportion of the test instances {IT) O. } Im I.M, that satisfy Ay (I 0 for all m E (1,.... M. We evaluate on the Karel dataset (Devlin et al. 20)7a), consisting of 10, 000 training reference Karel programs with 2, 500 validation and 2, 500 test samples. Each program
consists of sequence of actions to move agent inside a grid-world from one starting grid (input)
to an end grid (output). The size of the action space is 53 and average program length is around 20.
Baselines We incorporate five baseline algorithms in our evaluation. (i) MC-2 (Eq 3), using token- level rewards and greedy baselines. (tt) MC-0, using sentence-level reward and token-level gokendeevelgrredy baselines, which corresponds to the TD-SCST in Rennie et al. (017)) ).. (iii) REINFORCE. using
sentence-level reward and with mini-batch mean as the baseline. (iv) Celf-Critic (SC) as in Rennie
et al. (017). (v) Remm 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 Re.bbaam 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) 8) 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. la and Tb 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 pscudo actions) and hence takes relative longer time per iteration, but soon becomes more and more confident (reflected as fewer and fewer pseudo actions per iteration) and turns faster. We note that the gradient variance at given iteration is related to both the property of the gradient estimator and the parameter value 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. Ic and Id plot the Generalization scores against training time on the training and validation sets. Due to large gradient variance, all methods except ARSM and RI. 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
'The original dataset contains I million training instances. Bunel et al. 2018) proposed to reduce the dataset to 10, 000 examples and observed significant improvement of RT. upon MLE when the reference program data limited. Our experiments are hased on the same reduced dataset.

---

Published as a conference paper at ICLR 2020
Nunio of idluus
RESM
teet
UC0
been
Teakeno
M24
08
(a) Log variance
(b) Number of millouts
(c) Training generalization
(d) Testing generalizatin
Figure 1: From left to right are the comparisons of various methods terms of gradient variance,
number of sequence rollouts, training Generalization score, and validation Generalization score.
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 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 Both MLE and RI. beam (Bunel et al. 2018) perform reasonably well, but are outperformed by ARSM with 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 adaptive number of correlated MC rollouts over a fixed number of independent MC rollouts.
4.2 IMAGE CAPTIONING
Image captioning, mapping an image æ to a summary sentence = (y/1-... "aa....,,rrr.. ) has become a standard task to compare different policy gradient based RL methods. We conduct experiments 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
ARS-K for computation-suuffiienn deployments We first investigate the effectiveness of the proposed method when it is computationally feasible to use the categorical softmax out- put layer at the test time. We con- sider MLE, REINFORCE, SC, and ARS-K with the same vocabulary size of V = 9788. For ARS-K, we ex- periment with several different K val- ues. We report CIDEr score, and other commonly used metries for the test set in 2 [ ARS-1 underperforms SC, ARS-K quickly improves K increases: ARS-5 becomes comparable SC performance; ARS- 10 and ARS-20 outperform SC by large margin with statistical significance (standard error is about 0.2). The superior performance of ARS-K with large is also evidenced by Fig. 3 The gradient variance of ARS-20 is significantly lower than other algorithms (Fig. 3b Sbjuperr). In Fig. (upper), we compare the average number of correlated MC rollouts of ARS-K K for different K. While in theory the number of unique pseudo actions in ARS-K could be many V at each step, 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 evident decreasing trend during training. Moreover, it increases slowly as increases (clearly below a linear increasing rate).
Mis artl
whire
Phrado modks with and bally luge
mun riding motoncycle with dag. nossgyylee pack
Picato anlls has driiing drvc the snict Paaao:i well mai with vallss and atret sign.
Przndo
We observe that while
Figure Main and pseudo trajectories for image captioning.
Fig. shows two pictures (see more plots in Appendix E.I with their main sentences. which 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
7

---

Published as a conference paper at ICLR 2020
PaarsonR=-0 2; p=t 36-15
AR5-20
Eerotion
Bonaton
06
Fkamber of Roblosh
ST-RRMM
PaarsonR=-0 31: p=? 1p-35
51-ANSM
6T-80
01-40
BT-6-000
BT 82
42
Tasttn
Ralkaza
(a) CIDEr
(b) Log variance
(c) Rollouts
Rolloats CIDE
Figure 3: Comparison of different gradient estimators for image captioning task. "RF" denotes REINFORCE
and "SC" denotes Self-Critic. Upper (lower) row: models using a regular softmax (binary- softmax).
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 pscudo 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- -iimted deployments We evaluate the binary- tree ARSM (BT-ARSM) described in Section 3.2 which decomposes the action space to sequence of binary actions.
(a) Binary free constructions and MLE pretraining. We explore three different ways to construct binary trees over the action space: (i) tree_We WV : We apply agglomerative clustering to off-the-shelf pre-trained Word-to-Vent Vector (WV) embeddings (Mikolov et al. 2013) to get a binary tree with a depth of 25; (ii) tree tree_DIS: We follow tree_WY Wee_ except that we use the word embeddings pre-trained for image captioning with standard MLE objective and full vocabulary to distill (DIS) the soll-softmax model's knowledge to the binary tree; (iii) tree_RD: We randomly permute the leaves of tree DIS to produce tree with no meaningful structure, referred as tree RD. We pre-train all these three models with MILE, and report the CIDEr score in T 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 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 binary- REe-REINFORCE (BT-RF), Lmary-tree-Self-CCree (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 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. and Fig. 4 (in Appendix A). our proposed methods can adaptively choose the number of correlated MC rollouts four aspects: (i)dat 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); (æi) (adapt across iterations) as shown in Fig. 3c] during the training, the number of correlated MC rollouts decreases as the model improves and converges; (adapt across sentence positions) as shown in Figs. 4a 4b 4d and 4e more MC
8

---

Published as a conference paper at ICLR 2020
Table 2: Performance comparison on test of COCO- caption dataset.
Method
CIDEr
BLEU-
BLUE-3
BLEU-2
BLEU-1
ROUGE
METEOR
Soft Attention Xu et al. 2015
24.3
34.4
49.2
70.7
23,9
Hard Attention XXu of al. 2015
25.0
35.7
50.4
71.8
23.0
Show Tell [Vinyals
85.5
27.7
23.7
ATT-FCN YYu
2016
30.4
40.2
53.7
70.9
24.3
SCN-LSTM Gan al. 2017
101.2
33.0
43.3
56.6
72.8
25.7
Vanilla Softmax with MI.E
93.3
30.4
40.2
53.6
70.7
52.2
24.7
REINFORCE
103.6
31.6
42.9
57.8
74.8
54.0
25.1
Self-Critic [Rennie al. 2017
106.5
32.3
43.8
25.6
ARS-I
103.6
30.8
42.4
57.8
75.4
54.0
25,2
ARS-5
106.1
31.6
43.3
58.6
76.1
54.6
25.5
ARS-10
107.7
31.8
43,5
58.8
76.0
54.7
25.7
ARS-20
108.4
32.1
43.8
$9.0
76.4
54.8
25.8
tree RD with MLE
77.5
23.1
33.7
48.6
67.1
49.2
22.4
tree WV with MLE
80.2
23.3
34.1
49.3
68.0
49.9
22.9
tree DIS with MLE tree. DIS with BT-RF
84.2
24.9
35.1
49.7
67.6
50.5
23.7
93.6
28.8
39.7
55.0
72.5
52.2
23.7
tree_DIS with BT-SC
96.5
29.4
55.4
72.7
52.8
24.1
Iree_DI with BT-ARSM
99.2
29.9
41,2
56.2
73.3
52.7
24.3
rollouts appear at the timesteps in the middle range of the generated sequence, as they are associated with higher uncertainty; (ad)(t across depths) as shown in Figs. [gg./4 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 predict. More details are provided in Appendix
CONCLUSION
In this paper, we demonstrate the adaptation of ARSM policy gradient estimator, utilizing token- level rewards of correlated Monte Carlo (MC) rollouts, 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 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 cold-start setting where we start from complete random policy, it is still challenging to make our methods work efficiently as the number of pseudo actions may be too large is large. We consider 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 (0017) and d' Autume et al. (2019).
for generating
categorical tokens and hence is suited for computation-limins
ACKNOWLEDGMENTS
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 NVIDIA Corporation with the donation of the Titan Xp GPU used for this research, and the computational support of Texas Advanced Computing Center.
REFERENCES
Peter Anderson, Xiaodong He. Chris Buchler, Damien Teney, Mark Johnson, Stephen Gould, and Lei Zhang. Bottom-up and top-down attention for image captioning and visual question answering. In Proceedings of the IEFE Conference Computer Vision and Pattern Recognition, pp. 6077-6086. 2018.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. arXiv preprint arXiv: A09i::449.0075: 2014.

---

Published as a conference paper at ICLR 2020
Dzmitry Bahdanau, Philemon Brakel, Kelvin Xu, Anirudh Goyal, Ryan Lowe, Joelle Pineau, Aaron Courville, and Yoshua Bengio. An actor-critic algorithm for sequence prediction. arXiv preprint arXiv:1607.07086. 2016.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In Advances in Neural Information Processing Systems, pp. 1171-1179, 2015.
Rudy Bunel, Matthew Hausknecht, Jacob Devlin, Rishabh Singh, and Pushmeet Kohli, Lever- aging grammar and reinforcement learning for neural program synthesis. arXiv preprint 1.X5::8805.04276, 2018.
Massimo Caccia, Lucas Caccia, William Fedus, Hugo Larochelle, Joelle Pineau, and Laurent Charlin. Language GANs falling short. arXiv preprint #xXiv:1811.02545 2018.
Patrick H Chen, Si Si, Sanjiv Kumar, Yang Li, and Cho-Jui Hsich. Learning to screen for fast softmax inference on large vocabulary neural networks. arXiv preprint arXiv: /8/0. /2406, 2018.
Xinyun Chen, Chang Liu, and Dawn Song. Execution-gunee guided neural program synthesis. In Inter- national Conference on Learning Representations, 2019. /tt/htttps///opeeeeeeeee
Eorum Perrumid--im -imeenddd--iooo -.g.-igggiAqiim
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1406. 1078, 2014.
Sumit Chopra, Michael Auli, and Alexander M Rush. Abstractive sentence summarization with attentive recurrent neural networks. In NAACL, 2016.
Cyprien de Masson d' Autume, Mihaela Rosca, Jack Rae, and Shakir Mohamed, Training language GANs from scratch. arXiv preprint arXiv:/905.099 09922,, 2019.
Jacob Devlin, Rudy R Bunel, Rishabh Singh, Matthew Hausknecht, and Pushmeet Kohli. Neural program meta-induction. In Advances in Neural Information Processing Systems, pp. 2080-2088 2017a.
Jacob Devlin, Jonathan Uesato, Surya Bhupatiraju, Rishabh Singh, Abdel-rahman Mohamed, and Pushmeet Kohli. Robustfill: Neural program learning under noisy i/o. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 990-998. JMLR. org, 2017b.
Nan Ding and Radu Soricut. Cold-start reinforcement learning with softmax policy gradient. In Advances in Neural Informaiion Processing Systems, pp. 2817-2826, 2017.
Zhe Gan, Chuang Gan, Xiaodong He, Yunchen Pu, Kenneth Tran, Jianfeng Gao, Lawrence Carin, and Li Deng. Semantic compositional networks for visual captioning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5630-5639, 2017.
Joshua Goodman. Classes for fast maximum entropy training. arXiv preprint cx/0108006, 2001.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoff Roeder, and David Duvenaud. Backpropagation through the Void: Optimizing control variates for black-box gradient estimation. In ICLR, 2018.
Edouard Grave, Armand Joulin, Moustapha Cissé, Hervé Jégou, et al. Efficient softmax approximation for GPUs. In Proceedings of the 34th International Conference Machine Learning-Volume Volume 70, pp. 1302-1310. JMLR. org, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-softmax. In ICI.R 2017.
10

---

Published as a conference paper at ICLR 2020
Andrej Karpathy and Li Fei-Fei. Deep visual-semante aligmments for generating image descriptions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2015.
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational Bayes. arXiv preprint
Alex M Lamb, Anirudh Goyal Alias Parth Goyal, Ying Zhang, Saizheng Zhang, Aaron C Courville, and Yoshua Bengio. Professor forcing: A new algorithm for training recurrent networks, In Advances In Neural Information Processing Systems, pp. 4601-4609, 2016.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and Lawrence Zitnick. Microsoft COCO: Common objects in context. In European conference on computer vision, pp. 740-755. Springer, 2014.
Siqi Liu, Zhenhai Zhu, Ning Ye. Sergio Guadarrama, and Kevin Murphy. Improved image captioning via policy gradient optimization of SPIDEr. In Proceedings of the IEEE international conference on computer vision, pp. 373-881,2017,
Ruotian Luo, Brian Price, Scott Cohen, and Gregory Shakhnarovich. Discriminability objective for training descriptive captions. arXiv preprint /803.v8803.000 04376, 2018.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: continuous relaxation of discrete random variables. In ICLR, 2017.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representa- tions in vector space. arXiv preprint arXiv:1301.370 2013.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In ICML, pp. 1791-1799, 2014.
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte Carlo gradient estimation in machine learning. arXiv preprint arXiv: 10652, 2019.
Frederic Morin and Yoshua Bengio. Hierarchical probabilistic neural network language model. In AISTATS, volume 5, pp. 246-25222 252.. Citeseer, 2005.
Mohammad Norouzi, Samy Bengio, Navdeep Jaitly, Mike Schuster, Yonghui Wu, Dale Schuurmans, et al. Reward augmented maximum likelihood for neural structured prediction. In Advances In Neural Information Processing Systems, pp. 1723-1731, 2016.
Art B. Owen. Monte Carlo Theory, Methods and Examples, chapter Variance Reduction. 2013.
John Paisley David M Blei, and Michael Jordan. Variational Bayesian inference with stochastic search. In ICML, pp. 1363-13700 2012.
Romain Paulus, Caiming Xiong, and Richard Socher. A deep reinforced model for abstractive summarization. arXiv preprint arXiv:/70s. 04304. 2017.
Jeffrey Pennington, Richard Socher, and Christopher Manning. Glove: Global vectors for word
representation. In Proceedings of the 2014 conference empirical methods in natural language processing (EMNLP), pp. 1532-1543, 2014.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. AISTATS, pp. 814-822, 2014.
Marc' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. In ICLR, 2016.
Steven Rennie, Etienne Marcheret, Youssef Mrouch, Jerret Ross, and Vaibhava Goel. Self-critical sequence training for image captioning. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 7008-7024, 2017.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In ICML., pp. 1278-1286, 2014.
11

---

Published as a conference paper at ICLR 2020
Alexander M Rush, Sumit Chopra, and Jason Weston. A neural attention model for abstractive sentence summarization. arXiv preprint arXiv:1509.0000 2015.
Kyuhong Shim, Minjae Lee, Iksoo Choi, Yoonho Boo, and Wonyong Sung. SVD-softmax: Fast softmax approximation on large vocabulary neural networks. In Advances in Neural Information Processing Systems, pp. 1463-5473, 2017.
Xujie Si, Hanjun Dai, Mukund Raghothaman, Mayur Naik, and Le Song. Learning loop invariants for program verification. In Advances in Neural Information Processing Systems, pp. 7751-7762, 2018.
Robin Sibson. Slink: an optimally efficient algorithm for the single-link cluster method. The computer jourmal, 16(1):30-34 1973.
Ilya Sutskever, Oriol Vinyals, and Quoe V Le. Sequence to sequence learning with neural networks. In Advances neural information processing systems, pp. 110411122,, 2014.
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, (1):9 44, 1988.
George Tucker, Andriy Mnih, Chris Maddison, John Lawson, and Jascha Sohl-Dickstein. REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models. In NIPS, pp. 2624-2633, 2017.
Ramakrishna Vedantam, C Lawrence Zitnick, and Devi Parikh. CIDEr: Consensus- based image description evaluation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4566 2015.
Oriol Vinyals, Alexander Toshev, Samy Bengio, and Dumitru Erhan. Show and tell: A neural image caption generator. In CVPR, 2015.
Jason D Williams and Steve Young. Partially observable Markov decision processes for spoken dialog systems. Computer Speech Language, 21(2):393 2007.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-22 1992.
Ronald Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280 1989.
Sam Wiseman and Alexander M Rush. Sequence-to-sequence learning as beam-search optimization. arXiv preprint
wew.:666602960 2016.
Lijun Wu, Fei Tian, Tao Qin, Jianhuang Lai, and Tie- Yan Liu. study of reinforcement learning for neural machine translation. arXiv preprint arXiv: 1rxiv:1808.08866. 2018.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhutdinov, Richard Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. arXiv preprint 100.:::1502.004::: 2015.
Zichao Yang, Zhiting Hu, Chris Dyer, Eric P Xing, and Taylor Kerg-Kirkpatrick. Unsupervised text style transfer using language models as discriminators. In Advances in Neural Information Processing Systems, pp. 7287-7298, 2018.
Mingzhang Yin and Mingyuan Zhou ARM: MngmenttREINFORCE-MT gradient for stochastic binary networks. In International Conference on Learning Representations, 2019. URL https
net. (/geeneeview.ntt/fooum.2.eee.e /ttp:/oummmidsllg011aa
Mingzhang Yin. Yuguang Yue, and Mingyuan Zhou ARSM: CaimeentREINNORCEE swap-megge estimator for gradient backpropagation through categorical variables. In International Conference Machine Learning, 2019.
Quanzeng You, Hailin Jin. Zhaowen Wang, Chen Fang, and Jiebo Luo. Image captioning with semantic attention. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4651-4659, 2016.
12

---

Published as a conference paper at ICLR 2020
Lantao Yu, Weinan Zhang. Jun Wang, and Yong Yu. Seqgan: Sequence generative adversarial nets with policy gradient. In Thirty-First AAAI Conference on Artificial Intelligence, 2017.
Wojciech Zaremba and Ilya Sutskever. Reinforcement learning neural turing machines-revised, arXiv preprint arXiv:1505.005: 2015.
Li Zhang, Flood Sung. Feng Liu, Tao Xiang, Shaogang Gong, Yongxin Yang, and Timothy M Hospedales. Actor-critic sequence training for image captioning. arXiv preprint arXiv: Ttt1::7006.90000 2017.
Minjia Zhang, Wenhan Wang, Xiaodong Liu, Jianfeng Gao, and Yuxiong He. Navigating with graph representations for fast and scalable decoding of neural language models. In Advances in Neural Information Processing Systems, pp. 1308-6319 2018.
Victor Zhong, Caiming Xiong, and Richard Socher. Seq2SQL: Generating structured queries from natural language using reinforcement learning. arXiv preprint arXiv:1709.00003, 2017.
13

---

Published as a conference paper at ICLR 2020
A ADAPTIVENESS OF ARS- K/BT-ARSM
7.5
4.5
Depti
Time
Iteration
Beration
Beration
steps 20)
(b) Rollouts vs. sleps (BT-ARSM)
Rollouts depth (BT-ARSM)
MC Rolouk
Rallora
MC
(d) Rollouts time sleps (ARS 20)
(c) Rollouts vs. time steps (T--ARSM)
(f) Rollouts vs. dopth (BT-ARSM)
Figure 4: Adaptiveness of MC rollouts with BT-ARSM/ARS-K
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 statistically significant negative correlation (p 0.05) between the number of rollouts and CIDEr scores, indicating that our algorithms CIDEr scores) and less rollouts for easier samples (higher CIDEr scores).
adaptively generate more rollouts for harder samples (lower
2. Adaptiveness across time steps: As shown in Figs. will be more MC rollouts at the middle time step of the generated sequence, since the initial words and end words more easily be learned due their cardinality, and ARSM model will be more confident about the choices for these words, leading to fewer MC rollouts.
4d and 4c is reasonable that there
Adaptiveness across depths: For binary softmax model, the terms in the vocabulary are represented as leaf nodes in the tree. As shown in Figs. and 4f ASRM successfully captures the adaptiveness of MC rollouts across depths in the tree. The top layers (close the root) of the tree are associated with more MC rollouts, since they are more uncertain about what to predict.
Adaptiveness across iterations: In Figs. lb and 3c observe that, during training, the number of correlated MC rollouts decreases as the model improves and converges.
B DETAILED FORMULATIONS
B. ARS/M GRADIENT ESTIMATORS
Applying the ARSM estimator in (5) to the expected reward shown in ② we have
ER E Ezur-rmpo( (-1 E m [ØARSM (Te)), JARSM (Tt), v Ey JARS (T. J).
1
BARS( T JT. (1(21111,,2 2 x. 4) Em= r x. 3) VTTe),,
where 27t mm = argmin mm e de,, Thus we can approximate the gradient VoER as
Σÿ= ÒARS J) JARS( (ARS(T::))) = [f 2 x.y) EN f(21=-1,, x. v) (1 Vmg)) (9)
VÆER ΣT ΣV @ARSM Vo Vodtos (ARSM(TT)),
sid Dir(Iv); 2 a. is an approximation of la m= [r(2 v)), which can be estimated with r(212-1, m Z a. y). where ž p- 2c, 21: mm))) ) is an MC rollout.
where T1...., NT
((11,111 mm xx, v)
14

---

Published as a conference paper at ICLR 2020
B.2 BT-ARSM GRADIENT ESTIMATORS
(11-11111) ER E Pint-rppol By Di 0, ( Uar..
(TH) = [r( (21--111 11 11 b(true) xz v) -r (211:11, 1, ((add) x. (r.y)](1/
(true) = ( 1 1) bu
( 1 Touzal-0rkul 2,)l
(10)
where r(21:-11, b1(1:1) 20, (0,
anc s1 x. y)]. Thus we have VOER = Ei=1 Eei Voweaan ER V Voor which can be approximated with VoER = Et= Eii (ARM( 2TH), where (ARM approximates (TR) shown in (TO) via MC
integration; note given Tet--tiif( b((ad))) = ((au)) then 9ARM(TH)
C ALGORITHM
Efficient Pseudo-Action Computation: To implement ARS-K, we need to compute the correspond- ing pscudo actions for each reference category j in reference set J. In other words, given π, Ø, need compute m= argmin implementation would involve taking the minimum over V-dimensionnl 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 -diimnssional vectors |J| times. We first explain the notations and the basic ideas, and then present the algorithm in Algorithm !
In x; - Φi. for all m E "i....,V,,.
J. A naive
Let Oú, = In ग; Φj. Denote mi, as the indexes for the top 2 smallest in for, respectively (this is the only place where we need take the minimum over whole V-dimensional vectors). Let ID(o denote the function returning the second index of Oigi which means
ID j-
the following, we show that is not necessary to take minimum over whole V-dimensional vectors for each pseudo actions. We can compute == efficiently by observing:
If m4 & (m,J), the smallest value in In m= can only be among the updated two values (0)m j) and the original smallest value 0m1,m1 Hence, the index of the smallest value is
m=j ([imin(min(()....) )m.),, 0m.)m)))
2. If m1 {m,j). 0m, will be updated to 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
m
(oj.m, jm,)); 0mggn
Proof. Assume m. E [m. m2 & {m, J. since Omimm has been changed, Omg.nng will be the smallest value in the vector excluding the two updated values. Hence, the smallest value will be among the three values {or.m. jm mmz,,,), among which we can find the index the smallest item. If m2 E {m,j}. j.. then {m. becomes {m1, m2) Hence, Om Om,,my and 0m2, will be updated. For any m & {m1, m2). min(O mm)) S Ommz.mz S O Therefore, the smallest value will be among {0mag. Orn 10000mmn The index of the smallest value will be ID(min(Om my mmly ))
which is equivalent to IDmin monin(min(()). (min(o).m. m, Om, im. mm),
15

---

Published as a conference paper at ICLR 2020
Algorithm 1: Compute Pseudo-Action Matrix for Reference Category Set in Parallel
input Batched ग and Φ. Ref-Cat Set J, ID Function
output: Pseudo-Action Matrix
Compute 0m1 M2, Omg m2 Top2(In Φ);
for j E m E {1...., V} (in parallel) do
Compute jj,m = in JT pm
Compute Om.j = nm 6
end for
Initialize P with size (IJ). V)
for j E J, m E {1...., V} (in parallel by using index matrix) do
if mm E {j.m] m] then
Plj, ml (Dmin(mmin( im..); Omn
else
Plj, m] ID(min(nin(()... Om.j) (mgg.
end if
end for
16

---

Published as a conference paper at ICLR 2020
Algorithm 2: (RS-K/ARSM(K = V) policy gradient for fine-tuning contextual categorical sequence generation model with discrete-action space of actions.
input MLE pre-trained policy parameter 0, 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 æ; y; First, sample main trajectory (21,...;,1)) 2T):
mt = greedily sample' then
for 1::7 T. let ān gr (- where Φ, (z
else
for = T. let 2 argmin (1...., Te Φu), where Dirichlet(1v)) (or let TH In(nni ((0, 1))), and o, 11
end if
Second, we compute pseudo actions:
for f = : 7 do
Let πτς ~ Dirichlet(11) 1v) (or let Tis = ==(n(Unif((,1)) for i 1,..., and then normalize them to have unit norm);
Let ji, be reference categories randomly sampled from {I,..., V} without replacement;
for k = ==.1.... K, I...., V (in parallel) do
Let 2, -= == arg. (I...... (In Tt v=
oii) as the (v, k)th pseudo action;
end for
Let S, unique {2,"² Denote the cardinality of St as [S,], where 1 < |s S V end for Third, we complete sentences and evaluate the reward for the unique of pseudo actions: for f = 1: (in parallel) do
"=""", Jus) which means S, is the of all unique values in {="">) Jugi
if ||| then
continue
end if
for 2ns E Sy (in parallel) do if T then
if pt 'greedily sample' then
greedily sample 2++1: po(2141:11 æ1:-11 Ets x)
else
randomly sample 2141 ~ po(2141:T| #1.t-1, 21,,
end if end if end for
for e = 1: 1iv K (in parallel) do
Let f(z = r(z11-11111 2t.. Zit1- x. b) if 2t == Z105
end for
end for
Finally, we compute the ARSM gradients and update parameters: for 1 = 1::: T,, k 1: K (in parallel) do
Let fa & EV f(2
for v = 1: V (in parallel) do
Let gik.v R (f(z) (k)= - Fex (1 ViTeje), where Iek.v is the eth component of 9
end for
end for for
: T (in parallel) do
0 f= 0 + no Ve where g, = EE Itk
with step-size
end for
end while
17

---

Published as a conference paper at ICLR 2020
Algorithm 3: Binary-tree-ARSM policy gradient for fine-tuning a binary-tree contextual categorical sequence generation model.
input MLE pre-trained policy parameter 0, binary code to word mapping output: Fine-tuned policy parameter 0:
while not converged do Receive random sample y/:
First, we sample a main trajectory (21,...;2r))
for i 1: : do
for 1 D do (t,ba(1e--) To( () n)
Sample td Uniform(0, 1); Let bed = 1 muso 152
end for =bb(())))
end for Second, we compute pseudo actions: for t 1: (in parallel) do
for d = D do
Let ( bid
Let bed I (e
id # Ued then
(1)
If d < D. sample (6
((d+1:D) (po(baati:))) Z1 b,1-d-1, (0) xx) = 1,2;
( v((())
Let 2td (i 1,2;
end if end for
Let St unique( Denote the cardinality of Sy as [S,), where D s |Sp < 2 D. end for Third, we complete sentences and evaluate the rewards for the unique set of pseudo actions: for = 1: (in parallel) do
{ (z Jas)) which means S, is the of all unique values in {=(d) Jd.-
for 2ιo E S, (in parallel) do If < T, sample 25 ~ (= 31a-1, Eess w);
end for
for d = 1: D.j 1: (in parallel) do
Let (() = 12 #441.T x y) 2nd
(7)
Zzsi
end for We compute the ARSM gradients and update parameters:
for d = 1:: (in parallel) do
if bbbl)) * #b then
Let Itbbuad- 1 ( 1 )1 2Tee);
0 Qupdate- 170 V Vothe 9t.
with step-size n
end if
0 = 0 + 0 Puplate
end for
end for end while
18

---

Published as a conference paper at ICLR 2020
D EXPERIMENTAL SETUP DETAILS
D.1 IMAGE CAPTIONING
Image captioning maps an image æ to sentence y (......,, .rr) that summarizes the image information. It has become a standard task to compare different RL methods using policy gradient. 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 al.) 2018 Xu et al. 2015). We conduct that consists of 123,287 images. Each image has at least five captions. We use standard data split from Karpathy & Fei-Fei 2015), with 113,287 training, 5000 validation, and 5000 testing images. The vocabulary size is 9488 and the max caption length T is 16. For the model architecture, we employ Fully-Connected (FC) model without attention (Rennie et al. (017). Image features extracted from a pre-trained ResNet (He et al.) 2016). Our implementation (2018). We pre-train model with MLE until convergence and it for initialization. The CIDEr scores between the generated captions and references are used as rewards.
experiments on the MS COCO dataset (Lin et al. 2014)
based on Luo et al.
E
QUALITATIVE RESULTS
E. PSEUDO SENTENCES PRODUCED BY ARS-K ALGORITHM.
Main sentence: a man riding motorcycle with dog.
Main sentence: group of people flying kites Pseudo sentence I: a man riding motorcycle in a field. with hay on it.
Pseudo sentence 1: group of teenagers stand- Pseudo sentence 2: a man riding motorcycle ing in field flying kites. with pack of sheep.
19

---

Published as a conference paper at ICLR 2020
MISSOURI ERRRDWERS
Main sentence: city that has large white building on it.
Pseudo sentence I: city with red traffic and large building.
Pseudo sentence 2: city intersection with traffic light and a street sign.
Main sentence: man and woman are standing in of a table.
Pseudo sentence 3: city bus is driving down the street.
Pseudo sentence I: man and woman are standing an market.
Pseudo sentence a city street with city street with cars parked on it.
Pseudo sentence 2: man and woman are standing the street.
Pseudo sentence 5: city road with a traffic light and a street sign.
Pseudo sentence 3: man and woman are standing in to tent.
Pseudo sentence 4: man and woman are standing in front of table.
20


Reviews:
[{"review_id": "HklBH2ICFr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe 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.\n\nSeveral questions on the technical side:\n1. 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)?\n2. 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?\n3. I understand that the ARSM estimator should be unbiased for V = 2. Does the estimator remain unbiased when V > 2?\n4. 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?\n5. How does the runtime of the proposed algorithm compare to the competitors?\n\nExperiments:\n1. 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?\n2. 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?\n\nMinor:\n1. p. 4, “expected award” >> “expected reward”\n2. g_{ARSM} defined twice in (5) and in the beginning of p. 4\n3. “Fig. 1 (left two) plots” and “Fig. 1 (right two) plots” not good\n\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "HJgX78WIcr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe paper is self-contained and clear. The main value of the paper is to present good experimental results on challenging tasks; the ARS-K variant, although fairly straightforward, seems to be a reasonable implementation of the ARS(M) estimator.\n\nMy main criticism on the paper is that the exact nature of the contribution is not properly stated. As far as I understand, the main value of the paper is to demonstrate the effectiveness of ASR-K/M on challenging tasks. In a first read however, it seems that the authors claim an algorithmic/theoretical contribution compared to the state-of-the-art. Comparing with the paper by Yin et al. (2019), it seems to me that the technical contribution is rather incremental (the binary tree version is a variant of the hierarchical softmax, and ASR-K seems very straightforward), up to the point that the first set of experiments is actually only about vanilla ARSM.\n\nother comments:\n- what is j in Eq 4?\n- RL_beam vs ASRM on neural program synthesis: the authors say that \"RL_beam overfits [...] because of biased gradients\", whereas \"ASRM converges to a local minimum that generalizes better\". I do not see why biased gradients would help fitting the data (compared to unbiased gradients). And as far as I understood, ASRM is about getting a better gradient (hence better optimization, and hence better fitting of the data), so I really do not understand this argument.\n\n- RL_beam vs ASRM on NPS: I do not see why ASRM cannot fit the data as well as RL_beam. Is there some regularization involved? \n\nminor:\n- \"expected award\" (first line section 3.1)\n\n------ after author rebuttal\n\nThe authors answered my main concerns, I raises my score to weak accept.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "B1xMsRpaFH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe proposed algorithm is novel, and the results are promising. Implementation of the idea seems non-trivial, but the authors provide open source code. The proposed algorithm could be impactful. The paper is clearly written.\n\nQuestions for the authors:\n- Can you say anything about the optimality of scaling the number of rollouts with the policy uncertainty? Does the algorithm make optimal use of the number of rollouts? i.e. is the variance minimal for the number of rollouts, or is there scope for improvement?\n- The number of rollouts being random possibly complicates efficient parallel evaluation of the rollouts (batch sizes are effectively varying). This is presumably not a problem for the chosen applications, but could you discuss the limitations in a broader setting?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a novel reinforcement learning-based algorithm for contextual sequence generation. Specifically, 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 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. Numerical experiments are presented with promising performance. 

Reviewers were in agreement that this is a non-trivial extension of previous work with broad potential application. Some concerns about better framing of contributions were mostly resolved during the author rebuttal phase. Therefore, the AC recommends publication. 
Decision: Accept (Poster)