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: In this paper, an unbiased estimator for expectations over discrete random variables is developed based on a sampling-without-replacement strategy. The proposed estimator is shown to be a Rao-Blackwellization of three existing unbiased estimators with guaranteed reduction in estimation variance. The connections of the method to other gradient estimators are discussed. Experimental results on several toy and real-data DL/RL problems are reported to demonstrate the applicability of the proposed estimators in the practice of machine learning.