Estimating Gradients for Discrete Random Variables by Sampling without Replacement

1 INTRODUCTION . Put replacement in your basement ! We derive the unordered set estimator1 : an unbiased ( gradient ) estimator for expectations over discrete random variables based on ( unordered sets of ) samples without replacement . In particular , we consider the problem of estimating ( the gradient of ) the expectation of f ( x ) where x has a discrete distribution p over the domain D , i.e . Ex∼p ( x ) [ f ( x ) ] = ∑ x∈D p ( x ) f ( x ) . ( 1 ) This expectation comes up in reinforcement learning , discrete latent variable modelling ( e.g . for compression ) , structured prediction ( e.g . for translation ) , hard attention and many other tasks that use models with discrete operations in their computational graphs ( see e.g . Jang et al . ( 2016 ) ) . In general , x has structure ( such as a sequence ) , but we can treat it as a ‘ flat ’ distribution , omitting the bold notation , so x has a categorical distribution over D given by p ( x ) , x ∈ D. Typically , the distribution has parameters θ , which are learnt through gradient descent . This requires estimating the gradient ∇θEx∼pθ ( x ) [ f ( x ) ] , using a set of samples S. A gradient estimate e ( S ) is unbiased if ES [ e ( S ) ] = ∇θEx∼pθ ( x ) [ f ( x ) ] . ( 2 ) The samples S can be sampled independently or using alternatives such as stratified sampling which reduce variance to increase the speed of learning . In this paper , we derive an unbiased gradient estimator that reduces variance by avoiding duplicate samples , i.e . by sampling S without replacement . This is challenging as samples without replacement are dependent and have marginal distributions that are different from p ( x ) . We further reduce the variance by deriving a built-in control variate , which maintains the unbiasedness and does not require additional samples . Related work . Many algorithms for estimating gradients for discrete distributions have been proposed . A general and widely used estimator is REINFORCE ( Williams , 1992 ) . Biased gradients based on a continuous relaxations of the discrete distribution ( known as Gumbel-Softmax or Concrete ) were jointly introduced by Jang et al . ( 2016 ) and Maddison et al . ( 2016 ) . These can be combined with the straight through estimator ( Bengio et al. , 2013 ) if the model requires discrete samples or be used to construct control variates for REINFORCE , as in REBAR ( Tucker et al. , 2017 ) or 1Code available at https : //github.com/wouterkool/estimating-gradients-without-replacement . RELAX ( Grathwohl et al. , 2018 ) . Many other methods use control variates and other techniques to reduce the variance of REINFORCE ( Paisley et al. , 2012 ; Ranganath et al. , 2014 ; Gregor et al. , 2014 ; Mnih & Gregor , 2014 ; Gu et al. , 2016 ; Mnih & Rezende , 2016 ) . Some works rely on explicit summation of the expectation , either for the marginal distribution ( Titsias & Lázaro-Gredilla , 2015 ) or globally summing some categories while sampling from the remainder ( Liang et al. , 2018 ; Liu et al. , 2019 ) . Other approaches use a finite difference approximation to the gradient ( Lorberbom et al. , 2018 ; 2019 ) . Yin et al . ( 2019 ) introduced ARSM , which uses multiple model evaluations where the number adapts automatically to the uncertainty . In the structured prediction setting , there are many algorithms for optimizing a quantity under a sequence of discrete decisions , using ( weak ) supervision , multiple samples ( or deterministic model evaluations ) , or a combination both ( Ranzato et al. , 2016 ; Shen et al. , 2016 ; He et al. , 2016 ; Norouzi et al. , 2016 ; Bahdanau et al. , 2017 ; Edunov et al. , 2018 ; Leblond et al. , 2018 ; Negrinho et al. , 2018 ) . Most of these algorithms are biased and rely on pretraining using maximum likelihood or gradually transitioning from supervised to reinforcement learning . Using Gumbel-Softmax based approaches in a sequential setting is difficult as the bias accumulates because of mixing errors ( Gu et al. , 2018 ) . 2 PRELIMINARIES . Throughout this paper , we will denote with Bk an ordered sample without replacement of size k and with Sk an unordered sample ( of size k ) from the categorical distribution p. Restricted distribution . When sampling without replacement , we remove the set C ⊂ D already sampled from the domain and we denote with pD\C the distribution restricted to the domain D \C : pD\C ( x ) = p ( x ) 1− ∑ c∈C p ( c ) , x ∈ D \ C. ( 3 ) Ordered sample without replacement Bk . Let Bk = ( b1 , ... , bk ) , bi ∈ D be an ordered sample without replacement , which is generated from the distribution p as follows : first , sample b1 ∼ p , then sample b2 ∼ pD\ { b1 } , b3 ∼ pD\ { b1 , b2 } , etc . i.e . elements are sampled one by one without replacement . Using this procedure , Bk can be seen as a ( partial ) ranking according to the PlackettLuce model ( Plackett , 1975 ; Luce , 1959 ) and the probability of obtaining the vector Bk is p ( Bk ) = k∏ i=1 pD\B i−1 ( bi ) = k∏ i=1 p ( bi ) 1− ∑ j < i p ( bj ) . ( 4 ) We can also restrict Bk to the domain D \ C , which means that bi 6∈ C for i = 1 , ... , k : pD\C ( Bk ) = k∏ i=1 pD\C ( bi ) 1− ∑ j < i pD\C ( bj ) = k∏ i=1 p ( bi ) 1− ∑ c∈C p ( c ) − ∑ j < i p ( bj ) . ( 5 ) Unordered sample without replacement . Let Sk ⊆ D be an unordered sample without replacement from the distribution p , which can be generated simply by generating an ordered sample and discarding the order . We denote elements in the sample with s ∈ Sk ( so without index ) and we write B ( Sk ) as the set of all k ! permutations ( orderings ) Bk that correspond to ( could have generated ) Sk . It follows that the probability for sampling Sk is given by : p ( Sk ) = ∑ Bk∈B ( Sk ) p ( Bk ) = ∑ Bk∈B ( Sk ) k∏ i=1 p ( bi ) 1− ∑ j < i p ( bj ) = ∏ s∈Sk p ( s ) · ∑ Bk∈B ( Sk ) k∏ i=1 1 1− ∑ j < i p ( bj ) . ( 6 ) The last step follows since Bk ∈ B ( Sk ) is an ordering of Sk , such that ∏k i=1 p ( bi ) = ∏ s∈S p ( s ) . Naive computation of p ( Sk ) is O ( k ! ) , but in Appendix B we show how to compute it efficiently . When sampling from the distribution restricted to D \ C , we sample Sk ⊆ D \ C with probability : pD\C ( Sk ) = ∏ s∈Sk p ( s ) · ∑ Bk∈B ( Sk ) k∏ i=1 1 1− ∑ c∈C p ( c ) − ∑ j < i p ( bj ) . ( 7 ) The Gumbel-Top-k trick . As an alternative to sequential sampling , we can also sample Bk and Sk by taking the top k of Gumbel variables ( Yellott , 1977 ; Vieira , 2014 ; Kim et al. , 2016 ) . Following notation from Kool et al . ( 2019c ) , we define the perturbed log-probability gφi = φi + gi , where φi = log p ( i ) and gi ∼ Gumbel ( 0 ) . Then let b1 = arg maxi∈D gφi , b2 = arg maxi∈D\ { b1 } gφi , etc. , soBk is the top k of the perturbed log-probabilities in decreasing order . The probability of obtaining Bk using this procedure is given by equation 4 , so this provides an alternative sampling method which is effectively a ( non-differentiable ) reparameterization of sampling without replacement . For a differentiable reparameterization , see Grover et al . ( 2019 ) . It follows that taking the top k perturbed log-probabilities without order , we obtain the unordered sample set Sk . This way of sampling underlies the efficient computation of p ( Sk ) in Appendix B . 3 METHODOLOGY . In this section , we derive the unordered set policy gradient estimator : a low-variance , unbiased estimator of ∇θEpθ ( x ) [ f ( x ) ] based on an unordered sample without replacement Sk . First , we derive the generic ( non-gradient ) estimator for E [ f ( x ) ] as the Rao-Blackwellized version of a single sample Monte Carlo estimator ( and two other estimators ! ) . Then we combine this estimator with REINFORCE ( Williams , 1992 ) and we show how to reduce its variance using a built-in baseline . 3.1 RAO-BLACKWELLIZATION OF THE SINGLE SAMPLE ESTIMATOR . A very crude but simple estimator for E [ f ( x ) ] based on the ordered sample Bk is to only use the first element b1 , which by definition is a sample from the distribution p. We define this estimator as the single sample estimator , which is unbiased , since EBk∼p ( Bk ) [ f ( b1 ) ] = Eb1∼p ( b1 ) [ f ( b1 ) ] = Ex∼p ( x ) [ f ( x ) ] . ( 8 ) Discarding all but one sample , the single sample estimator is inefficient , but we can use RaoBlackwellization ( Casella & Robert , 1996 ) to signficantly improve it . To this end , we consider the distribution Bk|Sk , which is , knowing the unordered sample Sk , the conditional distribution over ordered samples Bk ∈ B ( Sk ) that could have generated Sk.2 Using Bk|Sk , we rewrite E [ f ( b1 ) ] as EBk∼p ( Bk ) [ f ( b1 ) ] = ESk∼p ( Sk ) [ EBk∼p ( Bk|Sk ) [ f ( b1 ) ] ] = ESk∼p ( Sk ) [ Eb1∼p ( b1|Sk ) [ f ( b1 ) ] ] . The Rao-Blackwellized version of the single sample estimator computes the inner conditional expectation exactly . Since Bk is an ordering of Sk , we have b1 ∈ Sk and we can compute this as Eb1∼p ( b1|Sk ) [ f ( b1 ) ] = ∑ s∈Sk P ( b1 = s|Sk ) f ( s ) ( 9 ) where , in a slight abuse of notation , P ( b1 = s|Sk ) is the probability that the first sampled element b1 takes the value s , given that the complete set of k samples is Sk . Using Bayes ’ Theorem we find P ( b1 = s|Sk ) = p ( Sk|b1 = s ) P ( b1 = s ) p ( Sk ) = pD\ { s } ( Sk \ { s } ) p ( s ) p ( Sk ) . ( 10 ) The step p ( Sk|b1 = s ) = pD\ { s } ( Sk \ { s } ) comes from analyzing sequential sampling without replacement : given that the first element sampled is s , the remaining elements have a distribution restricted toD\ { s } , so sampling Sk ( including s ) given the first element s is equivalent to sampling the remainder Sk \ { s } from the restricted distribution , which has probability pD\ { s } ( Sk \ { s } ) ( see equation 7 ) . 2Note that Bk|Sk is not a Plackett-Luce distribution restricted to Sk ! The unordered set estimator . For notational convenience , we introduce the leave-one-out ratio . Definition 1 . The leave-one-out ratio of s w.r.t . the set S is given by R ( Sk , s ) = p D\ { s } ( Sk\ { s } ) p ( Sk ) . Rewriting equation 10 as P ( b1 = s|Sk ) = p ( s ) R ( Sk , s ) shows that the probability of sampling s first , given Sk , is simply the unconditional probability multiplied by the leave-one-out ratio . We now define the unordered set estimator as the Rao-Blackwellized version of the single-sample estimator . Theorem 1 . The unordered set estimator , given by eUS ( Sk ) = ∑ s∈Sk p ( s ) R ( Sk , s ) f ( s ) ( 11 ) is the Rao-Blackwellized version of the ( unbiased ! ) single sample estimator . Proof . Using P ( b1 = s|Sk ) = p ( s ) R ( Sk , s ) in equation 9 we have Eb1∼p ( b1|Sk ) [ f ( b1 ) ] = ∑ s∈Sk P ( b1 = s|Sk ) f ( s ) = ∑ s∈Sk p ( s ) R ( Sk , s ) f ( s ) . ( 12 ) The implication of this theorem is that the unordered set estimator , in explicit form given by equation 11 , is an unbiased estimator of E [ f ( x ) ] since it is the Rao-Blackwellized version of the unbiased single sample estimator . Also , as expected by taking multiple samples , it has variance equal or lower than the single sample estimator by the Rao-Blackwell Theorem ( Lehmann & Scheffé , 1950 ) .

Summary: This paper introduces an gradient estimator for loss functions that are expectations over discrete random variables. The basic idea is that an estimator over a discrete distribution can be Rao-Blackwellized by conditioning on the event that the discrete realization was produced by being the first sample drawn from an unordered set of samples drawn with replacement. Much of the paper is spent showing how this Rao-Blackwellized estimator can be computed in practice and how it compares to other known estimators.