summaries: The paper proposes to improve standard variational inference by increasing the flexibility of the variational posterior by introducing a finite set of auxiliary variables. Motivated by the limited expressivity of mean field variational inference the author suggests to iteratively refine a ‘proposal’ posterior by conditioning on a sequence of auxiliary variables with decreasing variance. The key requirement to set the variance of the auxiliary variables such that the integrating over them leaves the original model unchanged. As noted by the authors this is a variant of auxiliary variables introduced by Barber & Agakov. The motivation and theoretical sections seems sound and the experimental results are encouraging, however maybe not completely supporting the claim of new ‘state of the art’ on uncertainty prediction.

summaries: This paper describes a method for training flexible variational posterior distributions, which consists in making iterative locale refinements to an initial mean-field approximation, using auxiliary variables. The focus is on Gaussian latent variables, and the method is applied to Bayesian neural nets to perform variational inference (VI) over the weights. Empirical results show improvements upon the performance of the mean-field approach and some other baselines, on classification and regression tasks.

summaries: The paper proposes a new way to improving the variational posterior. The q(w) is an integral over multiple auxiliary variables. These auxiliary variables can be specified at different scales that can be refined to match different scales of details of the true posterior. They show better performance regression and classification benchmark datasets. They also show that the training time is at a reasonable scale when being parallelized.

abstractText: Variational inference (VI) is a popular approach for approximate Bayesian inference that is particularly promising for highly parameterized models such as deep neural networks. A key challenge of variational inference is to approximate the posterior over model parameters with a distribution that is simpler and tractable yet sufficiently expressive. In this work, we propose a method for training highly flexible variational distributions by starting with a coarse approximation and iteratively refining it. Each refinement step makes cheap, local adjustments and only requires optimization of simple variational families. We demonstrate theoretically that our method always improves a bound on the approximation (the Evidence Lower BOund) and observe this empirically across a variety of benchmark tasks. In experiments, our method consistently outperforms recent variational inference methods for deep learning in terms of log-likelihood and the ELBO. We see that the gains are further amplified on larger scale models, significantly outperforming standard VI and deep ensembles on residual networks on CIFAR10.

authors:
Felix V Agakov
David Barber

title: An auxiliary variational method

venue: In International Conference on Neural Information Processing,

authors:
Dario Amodei
Chris Olah
Jacob Steinhardt
Paul Christiano
John Schulman
Dan Mané

title: Concrete problems in ai safety

venue: arXiv preprint arXiv:1606.06565,

authors:
Charles Blundell
Julien Cornebise
Koray Kavukcuoglu
Daan Wierstra

title: Weight uncertainty in neural network

venue: In International Conference on Machine Learning,

authors:
Chris Cremer
Xuechen Li
David Duvenaud

title: Inference suboptimality in variational autoencoders

venue: arXiv preprint arXiv:1801.03558,

title: Practical variational inference for neural networks

venue: In Advances in neural information processing systems,

authors:
Chuan Guo
Geoff Pleiss
Yu Sun
Kilian Q. Weinberger

title: On calibration of modern neural networks

authors:
Fangjian Guo
Xiangyu Wang
Kai Fan
Tamara Broderick
David B Dunson

title: Boosting variational inference

venue: arXiv preprint arXiv:1611.05559,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
José Miguel Hernández-Lobato
Ryan Adams

title: Probabilistic backpropagation for scalable learning of Bayesian neural networks

authors:
Geoffrey Hinton
Drew Van Camp

title: Keeping neural networks simple by minimizing the description length of the weights

venue: In in Proc. of the 6th Ann. ACM Conf. on Computational Learning Theory. Citeseer,

authors:
Devon Hjelm
Ruslan R Salakhutdinov
Kyunghyun Cho
Nebojsa Jojic
Vince Calhoun
Junyoung Chung

title: Iterative refinement of the approximate posterior for directed belief networks

venue: In Advances in Neural Information Processing Systems,

authors:
Sergey Ioffe
Christian Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: arXiv preprint arXiv:1502.03167,

authors:
Yoon Kim
Sam Wiseman
Andrew C Miller
David Sontag
Alexander M Rush

title: Semi-amortized variational autoencoders

venue: arXiv preprint arXiv:1802.02550,

authors:
Diederik P. Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization, 2014. cite arxiv:1412.6980Comment: Published as a conference paper at the 3rd International Conference for Learning Representations

venue: San Diego,

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114,

authors:
Durk P Kingma
Tim Salimans
Max Welling

title: Variational dropout and the local reparameterization trick

authors:
Balaji Lakshminarayanan
Alexander Pritzel
Charles Blundell

title: Simple and scalable predictive uncertainty estimation using deep ensembles

authors:
Yann LeCun
Yoshua Bengio

title: Convolutional networks for images, speech, and time series

venue: The handbook of brain theory and neural networks,

authors:
Francesco Locatello
Gideon Dresdner
Rajiv Khanna
Isabel Valera
Gunnar Raetsch

title: Boosting black box variational inference

authors:
Christos Louizos
Max Welling

title: Multiplicative normalizing flows for variational Bayesian neural networks

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Joseph Marino
Yisong Yue
Stephan Mandt

title: Iterative amortized inference

venue: arXiv preprint arXiv:1807.09356,

authors:
Andrew C. Miller
Nicholas J. Foti
Ryan P. Adams

title: Variational boosting: Iteratively refining posterior approximations

authors:
Cuong V Nguyen
Yingzhen Li
Thang D Bui
Richard E Turner

title: Variational continual learning

venue: arXiv preprint arXiv:1710.10628,

authors:
Kazuki Osawa
Siddharth Swaroop
Anirudh Jain
Runa Eschenhagen
Richard E Turner
Rio Yokota
Mohammad Emtiyaz Khan

title: Practical deep learning with Bayesian principles

authors:
Yaniv Ovadia
Emily Fertig
Jie Ren
Zachary Nado
D Sculley
Sebastian Nowozin
Joshua V Dillon
Balaji Lakshminarayanan
Jasper Snoek

title: Can you trust your model’s uncertainty? evaluating predictive uncertainty under dataset shift

title: A mean field theory learning algorithm for neural networks

venue: Complex systems,

authors:
Rajesh Ranganath
Dustin Tran
David Blei

title: Hierarchical variational models

authors:
Carlos Riquelme
George Tucker
Jasper Roland Snoek

title: Deep Bayesian bandits showdown

venue: In International Conference on Representation Learning,

authors:
Francisco Ruiz
Michalis Titsias

title: A contrastive divergence for combining variational inference and mcmc

authors:
Tim Salimans
Diederik Kingma
Max Welling

title: Markov chain monte carlo and variational inference: Bridging the gap

authors:
Yeming Wen
Paul Vicol
Jimmy Ba
Dustin Tran
Roger Grosse

title: Flipout: Efficient pseudoindependent weight perturbations on mini-batches

venue: arXiv preprint arXiv:1803.04386,

authors:
Cheng Zhang
Judith Butepage
Hedvig Kjellstrom
Stephan Mandt

title: Advances in variational inference. IEEE transactions on pattern analysis and machine intelligence, 2018a

authors:
Yichuan Zhang
José Miguel Hernández-Lobato
Zoubin Ghahramani

title: Ergodic measure preserving flows

venue: arXiv preprint arXiv:1805.10377,

heading: 1 INTRODUCTION

text:
Uncertainty plays a crucial role in a multitude of machine learning applications, ranging from weather prediction to drug discovery. Poor predictive uncertainty risks potentially poor outcomes, especially in domains such as medical diagnosis or autonomous vehicles where some forms of high confidence errors may be especially costly (Amodei et al., 2016). Thus, it is becoming increasingly important that the underlying model provides high quality uncertainty estimates along with its predictions. Yet, possibly the most widely used models, deep neural networks (LeCun et al., 2015), are unable to accurately quantify model uncertainty. They are often overconfident in their predictions, even when their predictions are incorrect (Guo et al., 2017; Ovadia et al., 2019).
By marginalizing over a posterior distribution over the parameters given the training data, Bayesian inference provides a principled approach to capturing uncertainty. In contrast, standard training of neural networks employs a point estimate of the parameters, which cannot account for model uncertainty. Unfortunately, exact Bayesian inference is intractable in general for neural networks. To model epistemic uncertainty, variational inference (VI) instead approximates the true posterior with a simpler distribution. The most widely used one for neural networks is the mean-field approximation, where the posterior is represented using an independent Gaussian distribution over all the weights. Variational inference is appealing since it reduces the problem of inference to an optimization problem, minimizing the discrepancy between the true posterior and the variational posterior. The key challenge, however, is the task of training expressive posterior approximations that can capture the true posterior without significantly increasing the computational costs.
This paper describes a novel method for training highly flexible posterior approximations. The idea is to start with a coarse, mean-field approximation q(w) and make iterative, local refinements to it. The regions of the local refinements are determined by sampling the values of additive auxiliary variables. The model parameters w are expressed using a number of auxiliary variables ak (Figure 1 left) for k = 1, . . . ,K that leave the marginal distribution unchanged. In each iteration, we sample the value of an auxiliary variable according to the current variational approximation q(ak) and refine the approximation by conditioning on the newly sampled value q(w) ≈ p(w|x, y, a1:k) (Figure 1 right illustrates the process). Each refinement step makes cheap, local adjustments to the variational posterior in the region of the sampled auxiliary variables. At the end, we draw one sample from
the refined q(w). The refinement iterations have to be repeated for each posterior sample. The algorithm results in samples from a highly complex distribution, starting from a simple mean-field approximation. While the distribution of the samples is difficult to quantify, we show that it is not limited to factorized, uni-modal forms, and that the procedure is guaranteed to improve the resulting ELBO without posing a significant computational overhead.
Summary of contributions:
• A novel algorithm for refining a variational distribution, increasing its flexibility. • We show that the refinement steps are guaranteed to improve the quality of the variational
distribution under mild conditions. • We showcase the effectiveness of the method on Bayesian neural networks using a set
of UCI regression and image classification benchmarks. We set a new state-of-the-art in uncertainty estimation using variational inference at ResNet scale (ResNet-20, (He et al., 2016)) scale on CIFAR10.

heading: 2 METHODS

text: In this section, we first describe standard variational inference (VI), followed by a detailed description of the iterative refinement algorithm. While VI and our proposed methodology are generally applicable to latent variable models, in this work, we consider the application to Bayesian neural networks (Figure 1), where x are inputs, y outputs, and w the weights of network.

heading: 2.1 VARIATIONAL INFERENCE IN BAYESIAN NEURAL NETWORKS

text:
Exact Bayesian inference in Bayesian neural networks (BNN) (Figure 1) is often intractable and is NP-hard in the general case. Variational inference attempts to approximate the true posterior p(w|x, y) with an approximate posterior qφ(w), typically from a simple family of distributions, for example independent Gaussians over the weights, i.e. the mean-field approximation. To ensure that the approximate posterior is close to the true posterior, the parameters of qφ(w), φ are optimized to minimize their Kullback-Leibler divergence: DKL(qφ(w) ‖ p(w|x, y)). At the limit of DKL(qφ(w) ‖ p(w|x, y)) = 0, the approximate posterior exactly captures the true posterior, although this might not be achievable if p(w|x, y) is outside of the distribution family of qφ(w). In order to minimize the KL-divergence, variational inference optimizes the Evidence Lower Bound (ELBO) w.r.t. φ (denoted as L(φ)), which is a lower bound to the log marginal likelihood log p(y|x). Since the the marginal likelihood can be expressed as the sum of the KL-divergence and the ELBO, maximizing the ELBO is equivalent to minimizing the KL divergence:
log p(y|x) = DKL ( qφ(w) ∥∥ p(w|x, y))︸ ︷︷ ︸ ≥0 +L(φ) ≥ L(φ) = Eqφ [ log p(y|x,w) ] −DKL ( qφ(w) ∥∥ p(w)) , due to non-negativity of the KL-divergence.
1: procedure REFINEANDSAMPLE(φ) 2: for m = 1, . . . ,M do 3: φ0 ← φ 4: for k = 1, . . . ,K do 5: ak ∼ qφk−1(ak) . Sample ak 6: qφk(w)← qφk−1(w|ak) . Initialize φk 7: φk ← argmaxφk L|a1:k(φk) . Optimize φk 8: end for 9: wm ∼ qφK (w) . Sample the refined posterior
10: end for 11: return w1:M 12: end procedure
Algorithm 1: Pseudocode for the iterative refinement of the variational posterior
Following the optimization of φ, the model can be used to make predictions on unseen data. For a new input x′, the predictive distribution p(y′|x′, y, x) can be approximated by stochastically drawing a small number (around M ≤ 10) of sample model parameters and averaging their prediction in an ensemble model:
w1:M ∼ qφ(w), p(y′|x′, y, x) ≈ 1
M M∑ i=1 p(y′|x′, wi) .

heading: 2.2 REFINING THE POSTERIOR

text:
The main issue with variational inference is the inflexibility of the posterior approximation. The most widely used variant of variational inference, mean-field variational inference, approximates the posterior with independent Gaussians across all dimensions. This approximation is too simplistic to capture the complexities of the posterior for complicated models such as BNNs. Our idea is to refine the posterior approximation through iterative optimization. Using the refinement procedure, it is feasible to train a detailed posterior in the regions of the posterior samples used for prediction while relying on a coarse-grained approximation further away from these samples.
The graphical model is augmented with a finite number of auxiliary variables a1:K as shown in Figure 1. The constraints are that (x, y) must be conditionally independent of the auxiliary variables given w and that they must not affect the prior distribution p(w). This is important in justifying the use of the initial variational approximation. A significant way in which we distinguish ourselves from hierarchical variational models (Ranganath et al., 2016) is that the model is unaffected by the presence of the auxiliary variables. Their purpose is solely to aid the inference procedure. Given a Gaussian prior N (0, σ2w) over w, we express w as a sum of independent auxiliary variables2
w = K∑ k=1 ak, with p(ak) = N (0, σ2ak) for k = 1, . . . ,K ,
while ensuring that ∑K k=1 σ 2 ak
= σ2w so that the prior p(w) = N (0, σ2w) is unchanged. Locally refining the approximate posterior refers to iteratively sampling the values of the auxiliary variables a1:K and then approximating the posterior conditional on the sampled values, i.e. qφk(w) approximates p(w|x, y, a1:k) for iterations k = 1, . . . ,K (Algorithm 1). Starting from the initial mean field approximation qφ(w), we sample the value of a1 from qφ(a1) = ∫ p(a1|w)qφ(w) dw,
1The sudden drops after sampling are optimizer artefacts due to having to reset the parameters of Adam. 2While we are focusing on one specific definition of the auxiliary variables, additive auxiliary variables, note that all of our results straight-forwardly generalize to arbitrary joint distributions p(w, a1:K) that meet the constraints.
then optimize the approximation to the conditional posterior: qφ1(w) ≈ p(w|x, y, a1). This procedure is then iteratively repeated for a2:K . In iteration k,
1) ak ∼ ∫ p(ak|a1:k−1, w)qφk−1(w) dw 2) φk = argminDKL ( qφk(w) ∥∥ p(w|x, y, a1:k)) . Analogously to variational inference, the KL divergence is minimized through the optimization of the conditional ELBO in each iteration: L|a1:k(φk) = Eqφk [ log p(y|x,w) ] − DKL(qφk(w) ‖ p(w|a1:k)). In order to get independent samples from the variational posterior, we have to repeat the iterative refinement for each ensemble member w1:M .
Toy example We use a toy example to demonstrate the procedure. In this toy example, we have a single weight w with prior p(w) = N (0, 1) and a complicated posterior with four modes (synthetically generated data). We express w as the sum of two auxiliary variables w = a1 + a2 with p(a1) = N (0, 0.8) and p(a2) = N (0, 0.2) (which recovers p(w) as per the constraint). As Figure 3b shows, a Gaussian approximation to the posterior fails to capture the multimodal nature of the true posterior. The first step of the refinement is to sample using qφ: a1 ∼ qφ(a1) =∫ p(a1|w)qφ(w) dw. Next, we condition on the value of a1 and initialize the parameters of the vari-
ational posterior accordingly: qφ1(w) = N ( w|µφ1 , σ2φ1 ) ← qφ(w|a1). Both qφ(a1) and qφ(w|a1) can be computed analytically for Gaussian distributions. After optimizing φ1, the approximate posterior qφ1(w) is able to capture a good, local approximation to the posterior p(w|a1, x, y) (Figure 3c). In Figure 3d, we can see the histogram of the refined posterior, that is, the distribution we are generating samples from (for each sample from w, we drew a sample from qφ(a1) and optimized qφ1(w)). Clearly, it is a much better approximation to the true posterior than the Gaussian approximation we started with, although it is important to note that the true posterior is not recovered exactly.

heading: 2.3 THEORETICAL JUSTIFICATION

text:
Our theoretical claims are twofold. Firstly, that through this procedure, we are optimizing a lower bound to the ELBO and secondly, that the refinement cannot result in a worse posterior approximation than the initial mean-field approximation that we start with (in the ELBO sense). That is,
ELBOref ≥ ELBOaux ≥ ELBOinit ,
where ELBOref denotes the ELBO of the refined posterior qref, ELBOaux refers to the objective that the refinement process is optimizing, and ELBOinit is the ELBO of the initial mean-field approximation.
Lower bound to the ELBO Consider the case with two auxiliary variables a1 and a2. The initial training optimizes the ELBOinit = Eqφ [ log p(y|x,w) ] −DKL ( qφ(w) ∥∥ p(w)) and the refinement step optimizes the conditional ELBO, L|a1(φ1) = Eqφ1 [ log p(y|x,w) ] −DKL ( qφ1(w) ∥∥ p(w|a1)).
The key observation is that we can define ELBOaux that is a lower bound to ELBOref and is increased both by the initial training and the refinement steps:
ELBOaux = Eqφ [ Eqφ1 [ log p(y|x,w)− log qφ1(w)
p(w|a1)
] − log qφ(a1)
p(a1) ] = Eqref [ log p(y|x,w) ] −DKL ( qref(w, a1)
∥∥ p(w, a1)) ≤ Eqref [ log p(y|x,w) ] −DKL ( qref(w)
∥∥ p(w)) = ELBOref , since the KL divergence of the joint distribution is greater than or equal to that of the marginals.
Guarantee of improvement Improvement in the ELBO (ELBOref ≥ ELBOinit) is guaranteed under two assumptions. First, that the conditional variational posterior, qφ(w|a1), is within the variational family of qφ1 . Second, that the process that optimizes φ1 does not make it worse than the value it was initialized with. The first assumption holds for Gaussian families: qφ(w|a1) is Gaussian and can be computed in closed form. The second assumption is reasonable to assume for most optimizers and, in addition, it can be ensured by comparing the initial value to the final value and choosing the one with the more desirable objective value.
The argument goes as follows. By initializing qφ1 such that it coincides with qφ(w|a1), we can ensure that ELBOaux ≥ ELBOinit, since they are equal at initialization time and the optimization process does not decrease ELBOaux. From this combined with our previous result, it follows that ELBOref ≥ ELBOinit and therefore ensuring that the ELBO improves through the refining steps. Note that this also implies that it is not necessary to optimize until convergence: any amount of optimization increases the ELBO.
Figure 2 shows that the ELBO improvement occurs on real world datasets. With the sampling of each auxiliary variable, the ELBO improves forming a staircase pattern. Table 1 and 2 serve as further empirical evidence that the ELBO improves as a result of the refinement steps.
Extending to multiple auxiliary variables For simplicity, we stated the arguments with two auxiliary variables, but they straight-forwardly extend to any finite number of auxiliary variables. In this scenario, there are K ELBOaux-s, one for each auxiliary variable, upper bounded by ELBOref and lower bounded by ELBOinit.

heading: 3 RELATED WORKS

text:
While in theory, the Bayesian approach can accurately capture uncertainty, in practice we find that exact inference is computationally infeasible in most scenarios and thus we have to resort to approximate inference methods. There is a wealth of research on approximate inference methods, here we focus on works closely related to this paper.
Variational inference (Hinton & Van Camp, 1993) tries to approximate the true posterior distribution over parameters with a variational posterior from a simple family of distributions. Mean-field VI, which for neural networks traces back to Peterson (1987), uses independent Gaussian distributions over the parameters to try to capture the posterior. The advantage of the mean-field approximation is that the network can be efficiently trained using the reparameterization trick (Kingma & Welling, 2013) and the variational posterior has a proper density over the parameter space which then can be used across tasks like continual learning (Osawa et al., 2019; Nguyen et al., 2017) and contextual bandits (Riquelme et al., 2018). Recently, Louizos & Welling (2017) showed that normalizing flows can be employed to further increase the flexibility of the variational posterior. Zhang et al. (2018a) provide a detailed survey of recent advances in VI.
Our method is a novel variant of the auxiliary variable approaches to VI (Agakov & Barber, 2004; Ranganath et al., 2016) that increase the flexibility of the variational posterior through the use of auxiliary variables. The key distinction, however, is that instead of trying to train a complex variational approximation over the joint distribution, we iteratively train simple, mean-field approximations at the sampled values of the auxiliary variables. While this poses anO(MK) (K is the number of auxiliary variables andM is the number of posterior samples) overhead over mean-field VI, the training itself is kept straightforward and efficient. The introduction of every new auxiliary variable increases
the flexibility of the posterior approximation. In contrast to MCMC methods, it is unclear whether the algorithm approaches the true posterior in the limit of infinitely many auxiliary variables.
There are also numerous methods that start with an initial variational approximation and refine it through a few MCMC steps (Salimans et al., 2015; Zhang et al., 2018b; Ruiz & Titsias, 2019). The distinction from our algorithm is that we refine the posterior starting at large scale and iteratively move towards smaller scale refinements whereas these methods only refine the posterior at the scale of the MCMC steps. Guo et al. (2016); Miller et al. (2017) and Locatello et al. (2018) used boosting to refine the variational posterior, iteratively adding parameters such as mixture components to minimize the residual of the ELBO. Our method does not add parameters at training time but instead iteratively refines the samples through the introduction of auxiliary variables. We did not include these models among our baselines because they have yet to be applied to Bayesian multi-layer neural networks.
Further related works include methods that iteratively refine the posterior in latent variable models (Hjelm et al., 2016; Cremer et al., 2018; Kim et al., 2018; Marino et al., 2018). These methods, however, focus on reducing the amortization gap and they do not increase the flexibility of the variational approximation.
Lastly, there are non-Bayesian strategies for estimating epistemic uncertainty in deep learning. Bootstrapping (Breiman, 1996) and deep ensembles (Lakshminarayanan et al., 2017) are perhaps the most promising. Deep ensembles, in particular, have been demonstrated to achieve strong performance on benchmark regression and classification problems and uncertainty benchmarks including out-of-distribution detection (Lakshminarayanan et al., 2017) and prediction under distribution shift (Ovadia et al., 2019). Both methods rely on constructing a set of independently trained models to estimate the uncertainty. Intuitively, the amount of disagreement across models reflects the uncertainty in the ensemble prediction. To induce diversity among the ensemble members, bootstrapping subsamples the training set for each member while deep ensembles use the randomness in weight initialization and mini-batch sampling.

heading: 4 EXPERIMENTS

text:
The goal of the experiments is twofold. First, we empirically confirm the improvement in the ELBO as claimed in Section 2.3. Second, we quantify the improvement in the uncertainty estimates due to the refinement. We conducted experiments on a selection of regression and classification benchmarks using Bayesian neural networks as the underlying model. We looked at the marginal loglikelihood of the predictions as well as accuracy in classification tasks.
Refinement (Refined VI) After training the initial mean-field approximation, we refine M = 10 ensemble members, each with K = 5 auxiliary variables. The means on their prior distributions are fixed at 0., and their variances form a geometric series (each auxiliary variable has variance 0.3 times the previous one, which roughly halves the standard deviation of the prior each time): σ2a1 = 0.7σ 2 w, σ2a2 = 0.21σ 2 w, σ 2 a3 = 0.063σ 2 w, σ 2 a4 = 0.0189σ 2 w, and σ 2 a5 = 0.0081σ 2 w. In each refinement iteration, we optimized the posterior with Adam (Kingma & Ba, 2014) for 200 iterations. To keep the training stable, we scaled the learning rate according to the standard deviation of the conditional posterior: in iteration k, lr=0.3 k 2 0.001.

heading: 4.1 BASELINES

text:
We used three baselines. First, mean-field variational inference in order to quantify the improvement provided by refining. Second, multiplicative normalizing flows (MNF) to compare against a more flexible posterior approximation and third, deep ensemble models to compare against a state-of-theart non-Bayesian approach. For all methods we used a batch size of 256 and the Adam optimizer with the default learning rate of 0.0013.
3The hyperparameters of each baseline were tuned using a Bayesian optimization package, however we found batch size and learning rate to be consistent across methods.
Variational inference (VI) (Graves, 2011; Blundell et al., 2015) Naturally, we investigate the improvement over variational inference with a mean-field Gaussian posterior approximation. We do inference over all weights and biases with a Gaussian prior centered at 0. The variance of the prior is tuned through empirical Bayes. This model is trained for 30000 iterations.
Multiplicative Normalizing Flows (MNF), (Louizos & Welling, 2017) To measure the performance against a more flexible class of posterior approximations, we look at Multiplicative Normalizing Flows. In this work, the posterior means are augmented with a multiplier from a flexible distribution parameterized by the masked RealNVP. This model is trained with the default flow parameters for 30000 iterations.
Deep ensemble models, (Lakshminarayanan et al., 2017) Deep ensemble models are shown to be surprisingly effective at quantifying uncertainty. While they are less principled than Bayesian methods, they are certainly a competitive baseline. For the regression datasets, we used adversarial training ( = 0.01) whereas in classification we did not use adversarial training (adversarial training did not give an improvement on the classification benchmarks). For each dataset, we trained 10 ensemble members for 5000 iterations each.

heading: 4.2 REGRESSION BENCHMARKS

text:
Following Hernández-Lobato & Adams (2015), we evaluate the methods on a set of UCI regression benchmarks (Table 1). The datasets used a random 80-20 split for training and testing. The model used for these datasets is a feed forward neural network with a single hidden layer containing 50 units with a ReLU activation function. As it is common practice, we utilize the local reparameterization trick (Kingma et al., 2015).
On these benchmarks, we are able to confirm that the refinement step consistently improves both the ELBO and the uncertainty estimates over VI. On 7 out of the 9 datasets, Refined VI is one of the best performing approaches.

heading: 4.3 CLASSIFICATION BENCHMARKS

text:
We examine the performance on commonly used image classification benchmarks (Table 2). The architecture used for this experiment is the LeNet5 (LeCun et al., 1995) architecture containing three convolutional layers with 6, 16 and 120 channels respectively followed by a feed-forward layer with 84 units and an output layer with ReLu activations throughout the network. We use the local reparameterization trick (Kingma et al., 2015) for the dense layers and Flipout (Wen et al., 2018) for the convolutional layers to reduce the gradient noise.
On the classification benchmarks, we again are able to confirm that the refinement step consistently improves both the ELBO and the uncertainty estimates over VI. While Refined VI is unable to outperform Deep Ensembles in classification accuracy, it does outperform them in MLL on the largest dataset, CIFAR10.

heading: 4.4 LARGE SCALE MODELS

text:
To demonstrate the performance on larger scale models, we apply the refining algorithm to Residual Networks (He et al., 2016) with 20 layers (based on Keras’s ResNet implementation). We look at two models. Firstly, a model where we do inference over all of the residual blocks and secondly, following Ovadia et al. (2019), a hybrid model (ResNet Hybrid) where inference is only done over the final layer of each residual block and every other layer is treated as a regular layer. For this model, we used a batch-size of 256 and we decayed the learning rate starting from 0.001 over 200 epochs. We used 10 auxiliary variables each reducing the prior variance by a factor of 0.5. Furthermore we investigate the effect of Batch Normalization (Ioffe & Szegedy, 2015). While it is difficult to incorporate batch normalization into the Bayesian framework, its positive effects are undeniable.
Regarding batch normalization, we can confirm the findings of Osawa et al. (2019), that it provides a substantial improvement for VI, although interestingly, this improvement disappears in the case of the hybrid model. To our knowledge, the refined hybrid model is state-of-the-art in terms of MLL. It outperforms Osawa et al. (2019) in both MLL and accuracy.

heading: 4.5 COMPUTATIONAL COSTS

text: When introducing a novel algorithm for variational inference, we have to discuss the computational costs. Table 4 shows the wall-time required to train each model including the possibility of parallel training. Deep ensembles require fewer training iterations and parallelize very well, although it is important to note that the system we tested on is heavily optimized towards training these models both in hardware and software. For Refined VI, training the initial mean-field approximation cannot be parallelized, but the following refinement iterations can be straight-forwardly split into M threads.

heading: 4.6 CONCLUSIONS

text: In this paper, we describe a novel algorithm for refining a coarse variational approximation to the Bayesian posterior. We show, both theoretically and empirically, that the refined posterior is a better approximation to the posterior than the initial variational distribution. Our method outperforms the baseline variational approximations in both uncertainty estimation as well as computational requirements. It sets a new state-of-the-art in uncertainty estimation using variational inference at ResNet scale (ResNet-20) on CIFAR10.

heading: A CLOSED FORMS OF THE SAMPLING DISTRIBUTIONS

text:
In Section 2.2, we claim that qφ(a1) (used to sample a1) and qφ(w|a1) (used to initialize qφ1 ) have analytic solutions for additive Gaussians. For p(w) = N (0, σ2w), q(w) = N (µφ, σ2φ), p(a1) = N (0, σ2a1) and p(a2) = N (0, σ 2 a2),
qφ(a1) = ∫ p(a1|w)qφ(w) dw = N ( µφ σ2a1 σ2w , σ2φσ 4 a1 σ4w + σ2a1σ 2 a2 σ2w ) qφ1(w) = N ( µφ1 , σ 2 φ1 ) ← qφ(w|a1) = N (a1σ2φσ2w + µφσ2a2σ2w σ2φσ 2 a1 + σ 2 wσ 2 a2 , σ2φσ 2 wσ 2 a2 σ2a1σ 2 φ + σ 2 wσ 2 a2 ) . (1)
These are derived by applying Bayes rule to the Gaussian probability density functions.
B NOTE ON ELBOAUX ≥ ELBOINIT
We formally show that ELBOaux ≥ ELBOinit under the assumption that the conditional variational posterior, qφ(w|a1), is within the variational family of qφ1 for all values of a1. For a given a1, let φ10|a1 be a set of variational parameters such that ∀w, qφ10|a1 (w) = qφ(w|a1). Such φ10|a1 must exist as a result of the initial assumption. The formula for Gaussian distributions is given in Appendix A.
Set φ1|a1 4 such that
φ1|a1 = φ10|a1 if Eqφ10|a1 [ log p(y|x,w)− log qφ10|a1 (w) p(w|a1) ] ≥ Eqφ1opt|a1 [ log p(y|x,w)− log qφ1opt|a1 (w) p(w|a1) ] φ1opt|a1 otherwise (2) where φ1opt|a1 is the result of the stochastic optimizer that attempts to maximize
Eqφ1|a1
[ log p(y|x,w)− log qφ1|a1 (w)
p(w|a1)
] .
By construction, for all a1,
Eqφ1|a1
[ log p(y|x,w)− log qφ1|a1 (w)
p(w|a1)
] ≥
Eqφ10|a1
[ log p(y|x,w)− log qφ10|a1 (w)
p(w|a1)
] =
Eqφ(w|a1) [ log p(y|x,w)− log qφ(w|a1)
p(w|a1)
] .
(3)
Substituting this into ELBOaux gives
ELBOaux = Eqφ [ Eqφ1|a1 [ log p(y|x,w)− log qφ1|a1 (w)
p(w|a1)
] − log qφ(a1)
p(a1)
] ≥
Eqφ [ Eqφ(w|a1) [ log p(y|x,w)− log qφ(w|a1)
p(w|a1)
] − log qφ(a1)
p(a1)
] =
Eqφ [ log p(y|x,w)− log qφ(w)
p(w)
] =
ELBOinit
(4)
concluding the proof.
4We use the notation φ1|a1 instead of φ1 to emphasize the dependence on the value of a1.


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.

summaries: 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.

summaries: 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.

abstractText: 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.

authors:
Peter Anderson
Xiaodong He
Chris Buehler
Damien Teney
Mark Johnson
Stephen Gould
Lei Zhang

title: Bottom-up and top-down attention for image captioning and visual question answering

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Dzmitry Bahdanau
Kyunghyun Cho
Yoshua Bengio

title: Neural machine translation by jointly learning to align and translate

venue: arXiv preprint arXiv:1409.0473,

authors:
Dzmitry Bahdanau
Philemon Brakel
Kelvin Xu
Anirudh Goyal
Ryan Lowe
Joelle Pineau
Aaron Courville
Yoshua Bengio

title: An actor-critic algorithm for sequence prediction

venue: arXiv preprint arXiv:1607.07086,

authors:
Samy Bengio
Oriol Vinyals
Navdeep Jaitly
Noam Shazeer

title: Scheduled sampling for sequence prediction with recurrent neural networks

authors:
Rudy Bunel
Matthew Hausknecht
Jacob Devlin
Rishabh Singh
Pushmeet Kohli

title: Leveraging grammar and reinforcement learning for neural program synthesis

venue: arXiv preprint arXiv:1805.04276,

authors:
Massimo Caccia
Lucas Caccia
William Fedus
Hugo Larochelle
Joelle Pineau
Laurent Charlin

title: Language GANs falling short

venue: arXiv preprint arXiv:1811.02549,

authors:
Patrick H Chen
Si Si
Sanjiv Kumar
Yang Li
Cho-Jui Hsieh

title: Learning to screen for fast softmax inference on large vocabulary neural networks

venue: arXiv preprint arXiv:1810.12406,

authors:
Xinyun Chen
Chang Liu
Dawn Song

title: Execution-guided neural program synthesis

venue: In International Conference on Learning Representations,

authors:
Kyunghyun Cho
Bart Van Merriënboer
Caglar Gulcehre
Dzmitry Bahdanau
Fethi Bougares
Holger Schwenk
Yoshua Bengio

title: Learning phrase representations using RNN encoder-decoder for statistical machine translation

venue: arXiv preprint arXiv:1406.1078,

authors:
Sumit Chopra
Michael Auli
Alexander M Rush

title: Abstractive sentence summarization with attentive recurrent neural networks

venue: In NAACL,

authors:
Cyprien de Masson d’Autume
Mihaela Rosca
Jack Rae
Shakir Mohamed

title: Training language GANs from scratch

venue: arXiv preprint arXiv:1905.09922,

authors:
Jacob Devlin
Rudy R Bunel
Rishabh Singh
Matthew Hausknecht
Pushmeet Kohli

title: Neural program meta-induction

authors:
Jacob Devlin
Jonathan Uesato
Surya Bhupatiraju
Rishabh Singh
Abdel-rahman Mohamed
Pushmeet Kohli

title: Robustfill: Neural program learning under noisy i/o

authors:
Nan Ding
Radu Soricut

title: Cold-start reinforcement learning with softmax policy gradient

authors:
Zhe Gan
Chuang Gan
Xiaodong He
Yunchen Pu
Kenneth Tran
Jianfeng Gao
Lawrence Carin
Li Deng

title: Semantic compositional networks for visual captioning

title: Classes for fast maximum entropy training

venue: arXiv preprint cs/0108006,

authors:
Will Grathwohl
Dami Choi
Yuhuai Wu
Geoff Roeder
David Duvenaud

title: Backpropagation through the Void: Optimizing control variates for black-box gradient estimation

venue: In ICLR,

authors:
Edouard Grave
Armand Joulin
Moustapha Cissé
Hervé Jégou

title: Efficient softmax approximation for GPUs

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

authors:
Eric Jang
Shixiang Gu
Ben Poole

title: Categorical reparameterization with Gumbel-softmax

authors:
Andrej Karpathy
Li Fei-Fei

title: Deep visual-semantic alignments for generating image descriptions

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational Bayes

authors:
Alex M Lamb
Anirudh Goyal Alias Parth Goyal
Ying Zhang
Saizheng Zhang
Aaron C Courville
Yoshua Bengio

title: Professor forcing: A new algorithm for training recurrent networks

venue: In Advances In Neural Information Processing Systems,

authors:
Tsung-Yi Lin
Michael Maire
Serge Belongie
James Hays
Pietro Perona
Deva Ramanan
Piotr Dollár
C Lawrence Zitnick

title: Microsoft COCO: Common objects in context

venue: In European conference on computer vision,

authors:
Siqi Liu
Zhenhai Zhu
Ning Ye
Sergio Guadarrama
Kevin Murphy

title: Improved image captioning via policy gradient optimization of SPIDEr

venue: In Proceedings of the IEEE international conference on computer vision,

authors:
Ruotian Luo
Brian Price
Scott Cohen
Gregory Shakhnarovich

title: Discriminability objective for training descriptive captions

venue: arXiv preprint arXiv:1803.04376,

authors:
Chris J Maddison
Andriy Mnih
Yee Whye Teh

title: The concrete distribution: A continuous relaxation of discrete random variables

authors:
Tomas Mikolov
Kai Chen
Greg Corrado
Jeffrey Dean

title: Efficient estimation of word representations in vector space

venue: arXiv preprint arXiv:1301.3781,

authors:
Andriy Mnih
Karol Gregor

title: Neural variational inference and learning in belief networks

venue: In ICML, pp

authors:
Shakir Mohamed
Mihaela Rosca
Michael Figurnov
Andriy Mnih

title: Monte Carlo gradient estimation in machine learning

venue: arXiv preprint arXiv:1906.10652,

authors:
Frederic Morin
Yoshua Bengio

title: Hierarchical probabilistic neural network language model

venue: In AISTATS,

authors:
Mohammad Norouzi
Samy Bengio
Navdeep Jaitly
Mike Schuster
Yonghui Wu
Dale Schuurmans

title: Reward augmented maximum likelihood for neural structured prediction

title: Monte Carlo Theory, Methods and Examples, chapter 8 Variance Reduction

authors:
John Paisley
David M Blei
Michael I Jordan

title: Variational Bayesian inference with stochastic search

authors:
Romain Paulus
Caiming Xiong
Richard Socher

title: A deep reinforced model for abstractive summarization

venue: arXiv preprint arXiv:1705.04304,

authors:
Jeffrey Pennington
Richard Socher
Christopher Manning

title: Glove: Global vectors for word representation

venue: In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP),

authors:
Rajesh Ranganath
Sean Gerrish
David Blei

title: Black box variational inference

venue: In AISTATS, pp

authors:
Marc’Aurelio Ranzato
Sumit Chopra
Michael Auli
Wojciech Zaremba

title: Sequence level training with recurrent neural networks

authors:
Steven J Rennie
Etienne Marcheret
Youssef Mroueh
Jerret Ross
Vaibhava Goel

title: Self-critical sequence training for image captioning

authors:
Danilo Jimenez Rezende
Shakir Mohamed
Daan Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

venue: In ICML,

authors:
Alexander M Rush
Sumit Chopra
Jason Weston

title: A neural attention model for abstractive sentence summarization

venue: arXiv preprint arXiv:1509.00685,

authors:
Kyuhong Shim
Minjae Lee
Iksoo Choi
Yoonho Boo
Wonyong Sung

title: SVD-softmax: Fast softmax approximation on large vocabulary neural networks

authors:
Xujie Si
Hanjun Dai
Mukund Raghothaman
Mayur Naik
Le Song

title: Learning loop invariants for program verification

title: Slink: an optimally efficient algorithm for the single-link cluster method

venue: The computer journal,

authors:
Ilya Sutskever
Oriol Vinyals
Quoc V Le

title: Sequence to sequence learning with neural networks. In Advances in neural information processing

title: Learning to predict by the methods of temporal differences

venue: Machine learning,

authors:
George Tucker
Andriy Mnih
Chris J Maddison
John Lawson
Jascha Sohl-Dickstein

title: REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models

venue: In NIPS,

authors:
Ramakrishna Vedantam
C Lawrence Zitnick
Devi Parikh

title: CIDEr: Consensus-based image description evaluation

authors:
Oriol Vinyals
Alexander Toshev
Samy Bengio
Dumitru Erhan

title: Show and tell: A neural image caption generator

venue: In CVPR,

authors:
Jason D Williams
Steve Young

title: Partially observable Markov decision processes for spoken dialog systems

venue: Computer Speech & Language,

title: Simple statistical gradient-following algorithms for connectionist reinforcement learning

authors:
Ronald J Williams
David Zipser

title: A learning algorithm for continually running fully recurrent neural networks

venue: Neural computation,

authors:
Sam Wiseman
Alexander M Rush

title: Sequence-to-sequence learning as beam-search optimization

venue: arXiv preprint arXiv:1606.02960,

authors:
Lijun Wu
Fei Tian
Tao Qin
Jianhuang Lai
Tie-Yan Liu

title: A study of reinforcement learning for neural machine translation

venue: arXiv preprint arXiv:1808.08866,

authors:
Kelvin Xu
Jimmy Ba
Ryan Kiros
Kyunghyun Cho
Aaron Courville
Ruslan Salakhutdinov
Richard Zemel
Yoshua Bengio

title: Show, attend and tell: Neural image caption generation with visual attention

venue: arXiv preprint arXiv:1502.03044,

authors:
Zichao Yang
Zhiting Hu
Chris Dyer
Eric P Xing
Taylor Berg-Kirkpatrick

title: Unsupervised text style transfer using language models as discriminators

authors:
Mingzhang Yin
Mingyuan Zhou

title: ARM: Augment-REINFORCE-merge gradient for stochastic binary networks

authors:
Mingzhang Yin
Yuguang Yue
Mingyuan Zhou

title: ARSM: Augment-REINFORCE-swap-merge estimator for gradient backpropagation through categorical variables

venue: In International Conference on Machine Learning,

authors:
Quanzeng You
Hailin Jin
Zhaowen Wang
Chen Fang
Jiebo Luo

title: Image captioning with semantic attention

authors:
Lantao Yu
Weinan Zhang
Jun Wang
Yong Yu

title: Seqgan: Sequence generative adversarial nets with policy gradient

venue: In Thirty-First AAAI Conference on Artificial Intelligence,

authors:
Wojciech Zaremba
Ilya Sutskever

title: Reinforcement learning neural turing machines-revised

venue: arXiv preprint arXiv:1505.00521,

authors:
Li Zhang
Flood Sung
Feng Liu
Tao Xiang
Shaogang Gong
Yongxin Yang
Timothy M Hospedales

title: Actor-critic sequence training for image captioning

venue: arXiv preprint arXiv:1706.09601,

authors:
Minjia Zhang
Wenhan Wang
Xiaodong Liu
Jianfeng Gao
Yuxiong He

title: Navigating with graph representations for fast and scalable decoding of neural language models

authors:
Victor Zhong
Caiming Xiong
Richard Socher

title: Seq2SQL: Generating structured queries from natural language using reinforcement learning

venue: arXiv preprint arXiv:1709.00103,

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.