PaperlayerView exact revision on arXiv

Source-linked AI summary

Estimating Gradients for Discrete Random Variables by Sampling without Replacement

Wouter Kool, Herke van Hoof, Max Welling

arXiv:2002.06043v1cs.LGstat.ML

TL;DR

Estimating gradients of expectations over discrete variables is challenging when sampling without replacement makes draws dependent and changes their marginals. The paper derives an unbiased unordered-set estimator through Rao-Blackwellization and combines it with REINFORCE and a built-in baseline. Across the reported entropy regimes, it is consistently among the best estimators.

  • Problem

    Estimating gradients for discrete random variables requires handling variance, and sampling without replacement introduces dependent samples with marginals different from the original distribution.

    Source: §1 · p. 1

  • Method

    The unordered set estimator Rao-Blackwellizes three estimators and combines with REINFORCE using a built-in baseline based on other samples.

    Source: §3 · p. 3, §3.1 · p. 4, §3.2 · p. 5

  • Results

    The unordered set estimator is the only estimator consistently showing the lowest or comparable variance in both high- and low-entropy regimes.

    Source: §4.1 · p. 8

  • Takeaways & Limitations

    The estimator applies to discrete distributions that support sampling without replacement and is parameter-free except for sample size k.

    Source: §3.3 · p. 6, §5 · p. 10

    Computing leave-one-out ratios adds overhead, although they can be computed efficiently even for large k.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

We derive an unbiased estimator for expectations over discrete random variables based on sampling without replacement, which reduces variance as it avoids duplicate samples. We show that our estimator can be derived as the Rao-Blackwellization of three different estimators. Combining our estimator with REINFORCE, we obtain a policy gradient estimator and we reduce its variance using a built-in control variate which is obtained without additional model evaluations. The resulting estimator is closely related to other gradient estimators. Experiments with a toy problem, a categorical Variational Auto-Encoder and a structured prediction problem show that our estimator is the only estimator that is consistently among the best estimators in both high and low entropy settings.

1 INTRODUCTION

The paper introduces an unbiased gradient estimator for discrete variables that samples without replacement to reduce variance, while addressing the dependence and changed marginals this sampling creates. It situates this approach among existing discrete-gradient methods, including REINFORCE, relaxations, control variates, and explicit-summation approaches.

1 INTRODUCTION

  • Discrete-variable gradients arise in reinforcement learning, latent-variable models, structured prediction, and hard attention, where model parameters are learned by gradient descent.The target is the gradient of an expectation over a discrete distribution.Source: §1 · p. 1
  • The unordered set estimator gives an unbiased gradient estimate while reducing variance by avoiding duplicate samples, despite dependent draws and marginals that differ from p(x).Its built-in control variate further reduces variance without requiring additional samples.Source: §1 · p. 1
  • REINFORCE is a widely used estimator; biased continuous Gumbel-Softmax or Concrete relaxations can construct control variates for REINFORCE, as in REBAR (Tucker et al., 2017).Source: §1 · p. 1
  • REINFORCE variance-reduction methods also use control variates and related techniques.Examples include Paisley et al. (2012), Ranganath et al. (2014), Gregor et al. (2014), Mnih & Gregor (2014), Gu et al. (2016), and Mnih & Rezende (2016).Source: §1 · p. 2
  • Other approaches explicitly sum over marginal categories or some categories, approximate gradients by finite differences, or adapt the number of model evaluations to uncertainty.These approaches include Titsias & Lázaro-Gredilla (2015), Liang et al. (2018), Liu et al. (2019), Lorberbom et al. (2018; 2019), and Yin et al. (2019).Source: §1 · p. 2
  • Many structured-prediction methods use supervision or multiple samples but are biased, while Gumbel-Softmax approaches in sequential settings face accumulating bias from mixing errors.The passage notes that many such methods rely on maximum-likelihood pretraining or a gradual transition from supervised learning to reinforcement learning.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define ordered sampling without replacement as sequential draws from a distribution restricted to exclude earlier draws, and define an unordered sample as the set obtained by discarding that order. They also describe Gumbel-Top-k as an alternative way to generate these samples.

2 PRELIMINARIES

  • B_k denotes an ordered sample without replacement of size k, whereas S_k denotes an unordered sample of that size.Source: §2 · p. 2
  • The distribution restricted to D \ C excludes the already-sampled set C, and subsequent sampling takes place over that reduced domain.The paper uses p_{D\C} to denote this restricted distribution.Source: §2 · p. 2
  • An ordered sample B_k draws items sequentially, restricting each subsequent draw to the domain with previously sampled items removed.This procedure corresponds to a partial ranking under the Plackett-Luce model.Source: §2 · p. 2
  • An unordered sample S_k is formed by generating an ordered sample without replacement and discarding its order; its probability sums over the orderings that could have produced it.Naively computing this probability costs O(k!), while Appendix B gives an efficient computation.Source: §2 · p. 2
  • Gumbel-Top-k generates the ordered sample by ranking perturbed log-probabilities, with each successive choice taken from the remaining domain.The perturbations are Gumbel variables; taking the top k without order yields the unordered set, which underlies the efficient probability computation described in Appendix B.Source: §2 · p. 3

3 METHODOLOGY

The section derives an unbiased unordered-set estimator by Rao–Blackwellizing estimators based on samples without replacement, then combines it with REINFORCE and a built-in baseline. It relates the estimator to other multi-sample methods and compares variance across entropy settings.

3.1 RAO-BLACKWELLIZATION OF THE SINGLE SAMPLE ESTIMATOR

  • The unordered set estimator Rao–Blackwellizes the single-sample estimator by conditioning on the unordered sample set, preserving unbiasedness with variance no greater than the single-sample estimator.The single-sample estimator uses only the first element of the ordered sample.Source: §3.1 · pp. 3–4

3.2 RAO-BLACKWELLIZATION OF OTHER ESTIMATORS

  • Rao–Blackwellizing either the stochastic sum-and-sample estimator or the importance-weighted estimator also yields the unordered set estimator.It has equal or lower variance than the optimal stochastic sum-and-sample estimator without requiring a choice of m.Source: §3.2 · pp. 4–5

3.3 THE UNORDERED SET POLICY GRADIENT ESTIMATOR

  • Combining the unordered set estimator with REINFORCE gives an unbiased policy-gradient estimator, and a built-in baseline based on the other samples preserves unbiasedness.The baseline compares each sample’s value with an estimate of its expectation from the other samples.Source: §3.3 · pp. 5–6

3.4 RELATION TO OTHER MULTI-SAMPLE ESTIMATORS

  • The estimator is related to Murthy’s estimator, while the paper extends its use to expectations and policy gradients; its relation to the biased empirical-risk estimator highlights differences in term weighting.The authors note that additional biases in the empirical-risk implementation, such as length normalization, are incompatible with their estimator.Source: §3.4 · pp. 6–7

4 EXPERIMENTS

Across the Bernoulli toy experiment, VAE, and TSP, the unordered set estimator performs consistently well across entropy settings and compares favorably with several alternatives. In the TSP experiment, it outperforms three estimators and a strong baseline, while matching the biased risk estimator.

4.1 BERNOULLI TOY EXPERIMENT

  • In the Bernoulli toy experiment, the unordered set estimator is the only estimator consistently among the lowest-variance methods in both high- and low-entropy regimes across sample counts.With k = 3 samples, its optimization trajectory was indistinguishable from using the exact gradient.Source: §4.1 · p. 8

4.2 CATEGORICAL VARIATIONAL AUTO-ENCODER

  • The unordered set estimator has the lowest gradient variance in both small- and large-domain VAE settings, on par with the best stochastic sum-and-sample estimator and REINFORCE with replacement.These measurements use k = 4 samples on a trained model.Source: §4.2 · p. 8, Table 1 · p. 8
  • For VAE ELBO optimization, the estimator performs on par with REINFORCE with replacement and outperforms other estimators in at least one setting.The reported experiments use MNIST binarized at a threshold of 0.5; the authors note overfitting and leave regularization outside the paper’s scope.Source: §4.2 · p. 9, Figure 2 · p. 9

4.3 STRUCTURED PREDICTION FOR THE TRAVELLING SALESMAN PROBLEM

  • The estimator outperforms REINFORCE with replacement, stochastic sum-and-sample, and greedy rollout on TSP validation progress, while performing on par with the biased risk estimator.Against single-sample estimators, it also trains with less data and computational cost when results are plotted by number of instances.Source: §4.3 · p. 9, Figure 3 · p. 10

5 DISCUSSION

The unordered set estimator is a low-variance, unbiased alternative to biased Gumbel-Softmax, derived by Rao-Blackwellizing three estimators. The appendix details notation and computation, while the discussion notes that a strong REINFORCE baseline performs similarly in high-entropy settings.

5 DISCUSSION

  • The unordered set estimator is low-variance, unbiased, parameter free except for sample size k, and competitive in both high- and low-entropy regimes.It is the result of Rao-Blackwellizing three existing estimators, which guarantees equal or lower variance.Source: §5 · p. 10
  • REINFORCE with replacement, multiple samples, and a VIMCO-inspired baseline performs similarly to the proposed estimator in the high-entropy setting.The authors identify applying their estimator locally as future work, expecting more duplicate samples in the smaller domain.Source: §5 · p. 10

A NOTATION

C.1 UNBIASEDNESS OF THE SUM-AND-SAMPLE ESTIMATOR

The section establishes unbiasedness for the sum-and-sample estimator and shows that Rao-Blackwellizing its stochastic form yields the unordered set estimator.

C.1 UNBIASEDNESS OF THE SUM-AND-SAMPLE ESTIMATOR

  • The sum-and-sample estimator is unbiased for any set C ⊂ D.Its proof rewrites the conditional expectation as a sum over elements of the sampled set.Source: §C.1 · p. 16

C.2 RAO-BLACKWELLIZATION OF THE STOCHASTIC SUM-AND-SAMPLE ESTIMATOR

  • Rao-Blackwellizing the stochastic sum-and-sample estimator results in the unordered set estimator.The proof conditions on the sampled set and uses the conditional distribution of the last sampled element.Source: §C.2 · p. 17
  • The proof completes the Rao-Blackwellization by adding and subtracting the same sampled term, then moving terms independent of the sampled element outside the expectation.It applies Lemma 3 to obtain the theorem's estimator form.Source: §C.2 · p. 19

C.3 THE STOCHASTIC SUM-AND-SAMPLE ESTIMATOR WITH MULTIPLE SAMPLES

The appendix generalizes sum-and-sample estimation to multiple samples and proves that Rao-Blackwellization recovers the unordered set estimator, including for an importance-weighted estimator.

C.3 THE STOCHASTIC SUM-AND-SAMPLE ESTIMATOR WITH MULTIPLE SAMPLES

  • The generalized estimator sums k − m terms and samples m > 1 terms without replacement, applying the unordered set estimator on the domain restricted by the summed samples.The authors describe this as a trade-off between the number of summed and sampled terms.Source: §C.3 · p. 19
  • Rao-Blackwellizing the multi-sample stochastic sum-and-sample estimator results in the unordered set estimator.The proof reduces it to a single-sample version by conditioning on the initial ordered samples.Source: §C.3 · pp. 19–20

D.1 RAO-BLACKWELLIZATION OF THE IMPORTANCE-WEIGHTED ESTIMATOR

  • Rao-Blackwellizing the importance-weighted estimator also results in the unordered set estimator.The appendix uses the Gumbel threshold corresponding to the largest perturbed log-probability outside the sampled set.Source: §D.1 · pp. 20–21

D.2 THE IMPORTANCE-WEIGHTED POLICY GRADIENT ESTIMATOR WITH BUILT-IN BASELINE

The section combines REINFORCE with importance weighting to form an unbiased policy-gradient estimator. It also describes normalized variants and a built-in baseline.

D.2 THE IMPORTANCE-WEIGHTED POLICY GRADIENT ESTIMATOR WITH BUILT-IN BASELINE

  • Combining REINFORCE with the importance-weighted estimator yields an unbiased estimate of the policy gradient ∇θEpθ(x)[fθ(x)].The section presents the construction as an importance-weighted policy-gradient estimator.Source: §D.2 · p. 21
  • Normalizing the importance weights produces a lower-variance but biased variant, with a correction that weights individual terms.The correction uses the sample-dependent factor 1 − pθ(s) + pθ(s)qθ,κ(s).Source: §D.2 · pp. 21–22
  • The estimator includes a built-in baseline, and its normalized implementation uses a different normalization for each sample term.The resulting form is similar to the earlier estimator with a different baseline for each sample.Source: §D.2 · p. 22

E.1 PROOF OF UNBIASEDNESS OF THE UNORDERED SET POLICY GRADIENT ESTIMATOR WITH BASELINE

The proof establishes that the control variate has zero expectation, then identifies it as a Rao-Blackwellized estimator. This preserves unbiasedness while supporting the unordered set policy gradient estimator.

E.1 PROOF OF UNBIASEDNESS OF THE UNORDERED SET POLICY GRADIENT ESTIMATOR WITH BASELINE

  • The proof reduces unbiasedness of the baseline estimator to showing that its control variate has expectation zero.The section introduces this as the condition needed to prove unbiasedness.Source: §E.1 · p. 22
  • Conditioning on b1 and applying Bayes’ theorem handles the cases s′ ≠ s and s′ = s.For s′ = s, the restricted-distribution ratio is 1 by definition.Source: §E.1 · p. 22
  • The control variate is a Rao-Blackwellization: its expression depends on b1 and b2 and matches the k = 2 stochastic sum-and-sample estimator used as the baseline.The proof then uses the zero expectation of the score ∇θ log pθ(b1) to establish the control variate’s zero expectation.Source: §E.1 · p. 23

F.1 PROOF OF BUILT-IN BASELINE

The proof derives a built-in baseline for the RISK estimator by rewriting the gradient of a ratio and rearranging the resulting sum.

F.1 PROOF OF BUILT-IN BASELINE

  • The RISK estimator has a built-in baseline when gradients are taken through its normalization factor.The proof applies the log-derivative trick to express the gradient of the ratio as the ratio times a log-gradient, then swaps summation variables.Source: §F.1 · p. 24

G.1 EXPERIMENTAL DETAILS

The categorical VAE experiments reproduce a prior setup on MNIST, optimizing an ELBO with several gradient estimators. The section specifies the model, gradient expressions, and training settings.

G.1 EXPERIMENTAL DETAILS

  • The categorical VAE uses 20 latent dimensions with 10 categories each, a uniform prior, and fully connected encoder and decoder networks.The encoder has hidden layers of 512 and 256 units; the decoder has 256 and 512, both with LeakyReLU activations.Source: §G.1 · p. 24
  • For encoder parameters, the analytic-KL gradient combines a score-function term weighted by the reconstruction log likelihood with the KL gradient.The alternative sample-KL expression includes the prior and posterior log probabilities in the score-function weight.Source: §G.1 · pp. 24–25
  • The implementation omits zero-expectation score terms and warns that REINFORCE-style and direct-objective gradient estimates require different automatic-differentiation handling.With a uniform prior, the prior-weighted score term also has zero expectation; the built-in baseline cancels it automatically.Source: §G.1 · p. 25
  • Training uses analytic-KL ELBO optimization for 1,000 epochs with Adam; Gumbel-Softmax and RELAX use a 10^-4 learning rate after diverging at 10^-3.Other estimators use 10^-3, while ARSM uses sample KL and 3 × 10^-4; reported ELBO values use analytic KL.Source: §G.1 · p. 25

G.2 ADDITIONAL RESULTS

Additional VAE experiments compare estimator behavior during training and on standard binarized MNIST, while the TSP experiment applies the estimators to an Attention Model. The reported curves show overfitting in validation performance and broadly similar relative estimator behavior across binarizations.

G.2 ADDITIONAL RESULTS

  • During training, gradient-variance results resemble the trained-model results in Table 1, except at the start, and estimator rankings are mostly unchanged.Figure 4 tracks log gradient variance with k = 4 samples across training.Source: §G.2 · p. 25, Figure 4 · p. 26
  • Validation -ELBO rises after its minimum, likely from overfitting; before then, REINFORCE without replacement and the unordered set estimator perform similarly to the other estimators.The authors note that early stopping is usable in this setting.Source: §G.2 · p. 25, Figure 5 · p. 26
  • On standard binarized MNIST, -ELBO scores are higher and overfitting remains; relative estimator performance is generally similar to threshold-at-0.5 MNIST.One REINFORCE-with-replacement run diverged in this less stable setting.Source: §G.2 · p. 25, Figure 6 · p. 27

H TRAVELLING SALESMAN PROBLEM

  • The TSP experiment treats tour construction as sequence modeling and evaluates estimators on minimizing predicted tour length.The task is to order locations given as coordinates so that the tour length is minimized.Source: H TRAVELLING SALESMAN PROBLEM · p. 28
  • The experiment uses Kool et al.’s Attention Model, which selects tour nodes autoregressively and was trained with REINFORCE using a greedy-rollout baseline.The model encodes locations as a fully connected graph and decodes a tour one node at a time.Source: H TRAVELLING SALESMAN PROBLEM · p. 28
  • TSP training reuses the Attention Model code and training details for 100 epochs, with Adam at 10^-4 and no learning-rate decay.For k = 4 estimators, the batch size is reduced to 128 to compensate for the additional samples.Source: H TRAVELLING SALESMAN PROBLEM · p. 28

Related papers

Loading 2002.06043v1…