PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement

Wouter Kool, Herke van Hoof, Max Welling

arXiv:1903.06059v2cs.LGstat.ML

TL;DR

Sampling with replacement can duplicate sequence outputs, whereas beam search lacks probabilistic interpretation for statistical estimation. The paper introduces Stochastic Beam Search, an implicit Gumbel-Top-k procedure that samples unique sequences efficiently. Experiments report favorable translation diversity–quality trade-offs and lower-variance sentence-level BLEU estimation.

  • Problem

    Sequence models need representative sequence sets, but standard sampling can duplicate outputs and beam search produces deterministic, low-variability outputs unsuitable for statistical estimators.

  • Method

    Stochastic Beam Search implicitly applies Gumbel-Top-k to factorized sequence distributions, sampling without replacement without enumerating all sequences.

  • Results

    Stochastic Beam Search achieves favorable diversity–BLEU trade-offs in translation and, for T < 0.5, its normalized estimator significantly reduces variance versus Monte Carlo without observable bias.

  • Takeaways & Limitations

    The method provides a probabilistic bridge between sampling and beam search for diverse sequence generation and statistical estimation.

  • Takeaways & Limitations

    The empirical estimator has high variance, making normalized importance weights preferable in practice.

Abstract

from arXiv · show

The well-known Gumbel-Max trick for sampling from a categorical distribution can be extended to sample $k$ elements without replacement. We show how to implicitly apply this 'Gumbel-Top-$k$' trick on a factorized distribution over sequences, allowing to draw exact samples without replacement using a Stochastic Beam Search. Even for exponentially large domains, the number of model evaluations grows only linear in $k$ and the maximum sampled sequence length. The algorithm creates a theoretical connection between sampling and (deterministic) beam search and can be used as a principled intermediate alternative. In a translation task, the proposed method compares favourably against alternatives to obtain diverse yet good quality translations. We show that sequences sampled without replacement can be used to construct low-variance estimators for expected sentence-level BLEU score and model entropy.

1. Introduction

The paper frames sampling without replacement as a way to obtain representative, unique sequences while retaining a probabilistic interpretation. Stochastic Beam Search implicitly applies Gumbel-Top-k to sequence models, combining sampling’s statistical uses with beam search’s high-probability outputs.

  • Motivation: Standard sampling can produce duplicate sequences, whereas beam search finds unique high-probability sequences but lacks variability and cannot construct statistical estimators.This motivates an intermediate method for representative sequence sets.
  • Contribution: Gumbel-Top-k enables sampling k sequence-model outputs without replacement without enumerating the exponentially large sequence domain.The procedure uses top-down sampling on a tree of partial and complete sequences.
  • Contribution: Stochastic Beam Search expands stochastically perturbed log-probabilities and requires model evaluations growing only linearly with k and maximum sampled sequence length.It is derived as an implicit application of Gumbel-Top-k.
  • Interpretation: The method avoids duplicate samples while retaining a probabilistic interpretation, allowing uses such as importance sampling.This connects stochastic sampling with deterministic beam search.
  • Applications: Experiments use Stochastic Beam Search for diverse translation outputs and low-variance estimators of sentence-level BLEU score and model entropy.The paper presents it as a principled alternative to ordinary sampling or beam search.

2. Preliminaries

The preliminaries define categorical sampling and the Gumbel perturbation tricks, then formulate sequence models as trees where beam search can identify high-probability completions. Gumbel-Top-k selects ordered samples without replacement by ranking perturbed probabilities, while tree structure enables efficient top-down search.

  • The categorical distribution: A categorical variable assigns probability p_i to each domain element, represented by log-probabilities φ_i with exp φ_i proportional to p_i.The probabilities are normalized across the categorical domain.
  • The Gumbel-Max trick: Adding independent Gumbel noise to category log-probabilities and taking the maximum samples exactly from the categorical distribution.The perturbed value is written G^φ_i = G_i + φ_i.
  • The Gumbel-Top-k trick: The Gumbel-Top-k trick returns the indices of the k largest perturbed log-probabilities in decreasing order, producing an ordered categorical sample without replacement.Sampling without replacement sequentially renormalizes the remaining probabilities.
  • Sequence models: A sequence model factorizes probabilities into conditional next-token distributions, typically using softmax-normalized log-probabilities with optional temperature T.The model defines probabilities for partial and complete sequences through the chain rule.
  • Beam search: Beam search expands at most k highest-probability partial sequences per step and terminates with k complete sequences.It is a limited-width breadth-first search used to approximate high-probability sequence retrieval.
  • Gumbel-Top-k on a tree: In the tree representation, leaf bars are perturbed sequence log-probabilities, while internal-node bars are maxima over perturbed descendant leaves.Expanding the top k internal nodes at each level is sufficient to construct the top k leaves.

3. Stochastic Beam Search

Stochastic Beam Search implicitly applies Gumbel-Top-k to sequence models through top-down sampling, avoiding complete-tree instantiation while producing exact samples without replacement. It retains only the most promising perturbed partial sequences, connecting stochastic sampling with beam search.

  • Derivation: The method derives Stochastic Beam Search by replacing explicit complete-tree Gumbel-Top-k sampling with an equivalent top-down process.The explicit approach instantiates all sequences, whereas the equivalent process instantiates only sequences with the k largest perturbed log-probabilities.
  • Node representation: Each tree node represents a partial or completed sequence, with perturbed log-probability equal to the maximum perturbation among leaves in its subtree.The node score has a Gumbel distribution whose location is the node's sequence log-probability.
  • Top-down sampling: Top-down sampling recursively samples child perturbations conditioned on parent maxima, while preserving independent Gumbel perturbations at the leaves.The procedure can be viewed as ancestral sampling in a tree-structured graphical model, with an alternative conditional construction based on independently sampled child Gumbels.
  • Beam pruning: At every level, expanding only the k nodes with highest perturbed log-probabilities is sufficient to recover the top k leaves.Ancestors of the top k leaves are among these nodes; node scores upper-bound the leaves in their subtrees, enabling pruning.
  • Guarantee and interpretation: By Gumbel-Top-k, the resulting leaves are an ordered sample without replacement, and the procedure requires computation linear in k and maximum sequence length.The algorithm therefore provides a probabilistic beam-search procedure rather than ordinary deterministic beam search.
  • Interpretation: Stochastic Beam Search consistently propagates perturbations down a subtree, avoiding the repeated low-probability reselection problem of independently sampling partial sequences.This soft commitment lets a partial sequence remain on the beam while its total perturbed score stays among the top k.

4. Experiments

The experiments evaluate Stochastic Beam Search for diverse translation generation and for estimating expected sentence-level BLEU and model entropy. It provides favorable diversity–quality trade-offs and lower-variance estimates under the tested temperatures.

  • 4.1. Diverse Beam Search: The translation experiment compares Beam Search, Sampling, Stochastic Beam Search, and Diverse Beam Search across temperatures T = 0.1 to 0.8 and sample sizes k = 5, 10, 20.Diversity is measured using unique n-grams, while mean and maximum BLEU indicate sample quality.
  • 4.1. Diverse Beam Search: For the same diversity, Stochastic Beam Search achieves higher mean and maximum BLEU than Diverse Beam Search.At low temperatures below 0.5, its maximum BLEU is comparable to deterministic Beam Search.
  • 4.1. Diverse Beam Search: Stochastic Beam Search increases diversity without sacrificing the best translation at low temperatures below 0.5.Sampling attains higher mean BLEU at lower diversity, while Sampling and Diverse Beam Search have lower maximum BLEU than Beam Search and Stochastic Beam Search.
  • 4.2. BLEU score estimation: Sampling without replacement addresses inefficient repeated samples in low-entropy distributions but requires importance weights for unbiased estimation.The fixed-size estimator is unbiased, while normalizing the weights reduces variance but introduces bias that vanishes as k approaches the full domain.
  • 4.2. BLEU score estimation: The normalized Stochastic Beam Search estimator has significantly lower variance than the unnormalized estimator and, for T < 0.5, than Monte Carlo without observable bias.The comparison uses three sentences, temperatures T = 0.05, 0.1, 0.2, 0.5, and sample sizes k = 1 to 250.
  • 4.3. Conditional Entropy Estimation: Stochastic Beam Search also yields lower-variance entropy estimates than Monte Carlo for T < 0.5.Entropy estimation is relevant when optimization includes an entropy loss to encourage diversity.

5. Related Work

The paper extends prior Gumbel-based and optimization-based sampling ideas to multiple samples without replacement from factorized sequence distributions. It relates the resulting procedure to weighted sampling, Plackett–Luce rankings, and beam-search research.

  • 5. Related Work: The work extends top-down sampling to draw multiple samples without replacement from factorized distributions with potentially exponentially large domains.It does so by implicitly applying the Gumbel-Top-k trick and presents this as a new practical sampling method.
  • 5. Related Work: The Gumbel-Top-k trick is related to Weighted Reservoir Sampling and, over the complete domain, to the Thurstonian interpretation of Plackett–Luce rankings.These connections situate sampling without replacement within established weighted-sampling and ranking frameworks.
  • 5. Related Work: Earlier work used the Gumbel-Max trick for categorical relaxations and reparameterizable low-variance but biased gradient estimators.The paper positions its sequence-model application as a step toward improving related methods.
  • 5. Related Work: Beam search is widely used for approximate inference and sequence-level training across machine translation, image captioning, speech recognition, and structured prediction.The related literature includes many beam-search variants and methods that incorporate beam search into training.

6. Discussion

The discussion presents Stochastic Beam Search as an easy-to-implement bridge between sampling and beam search. Experiments support its use for diverse translation outputs and lower-variance statistical estimators.

  • 6. Discussion: Stochastic Beam Search samples sequences without replacement while retaining a probabilistic interpretation unlike ordinary beam search.This combination supports use as a principled alternative where sampling or beam search is applied.
  • 6. Discussion: Experiments show that Stochastic Beam Search produces high-diversity neural machine translation samples and lower-variance estimators.The authors describe it as a drop-in replacement and propose future statistical learning methods based on beam search's probabilistic interpretation.

A. Proof of the Gumbel-Top-k trick

The proof establishes the ordered Gumbel-Top-k sample without replacement by induction on k. The base case is Gumbel-Max, and the induction step uses independence of the maximum and argmax.

  • A. Proof of the Gumbel-Top-k trick: The result concerns an ordered sample without replacement from the categorical distribution.For each successive draw, previously selected indices are excluded from the domain.
  • A. Proof of the Gumbel-Top-k trick: The proof proceeds by induction on k, with the case k = 1 given by the Gumbel-Max trick.The induction step applies the result for k − 1 to establish the result for k.
  • A. Proof of the Gumbel-Top-k trick: The transition from the maximum to its argmax uses their independence, followed by the Gumbel-Max trick.These identities supply the key steps in the induction argument.

B. Sampling set of Gumbels with maximum T

This section defines the truncated Gumbel distribution through its location, maximum, and cumulative distribution function, then gives its inverse CDF.

  • A truncated Gumbel variable is characterized by location φ, maximum T, and CDF F_φ,T(g).
  • The notation G′ ∼ TruncatedGumbel(φ, T) denotes a Gumbel variable truncated at maximum T.
  • The inverse CDF provides the transformation used to sample from this truncated distribution.

B.2. Sampling set of Gumbels with maximum T

The procedure samples a set of Gumbels conditioned on their maximum by first selecting the arg max, fixing its value, and truncating the remaining variables below that maximum.

  • The method first samples the arg max i* when conditioning Gumbel variables on maximum T.
  • The arg max i* is independent of the maximum T, so the procedure need not condition on T when selecting it.
  • The selected variable is set to T, while every other variable is sampled from TruncatedGumbel(φ_i, T).
  • An equivalent construction starts with ordinary Gumbels, computes their maximum Z, and applies a monotonic transformation that preserves the arg max.
  • This construction recovers the same three steps: Gumbel-Max selects the arg max, the selected value becomes T, and the others are truncated Gumbels.

B.3. Numeric stability of truncated Gumbel computation

The section addresses numerical instability in truncated-Gumbel and importance-weight calculations by using stable logarithmic identities and specialized computations for small differences.

  • Direct computation of the truncated-Gumbel transformation can be unstable because large terms must be exponentiated.
  • The implementation uses log1pexp(a) = log(1 + exp(a)) and related log1p and expm1 identities for accurate evaluation.
  • The stable formulation rewrites the transformation as T − log(1 + exp(T − G_φi + log(1 − exp(G_φi − Z)))).
  • Importance weights require care because p_θ(y_i|x) / q_θ(y_i|x) can become very small, and direct tail-probability computation can suffer catastrophic cancellation.
  • For φ_i − κ < −10, the method computes the log importance weight directly using z = exp(φ_i − κ) and normalized φ_i = log p_θ(y_i|x).In this regime, 0 < z < 10^-6, so the alternative computation does not lose significant digits.

D. Proof of unbiasedness of priority sampling estimator

The proof establishes unbiasedness for priority sampling by associating random keys with domain elements, defining inclusion probabilities through a threshold, and showing that sampled contributions recover the target expectation.

  • The proof considers a finite domain with normalized probabilities p_i and an expectation E[f(i)] to estimate.
  • Each element receives a random key h_i, and q_i(a) = P(h_i > a) gives its inclusion probability at threshold a.
  • Using Gumbel-perturbed log-probabilities as keys yields the PPSWOR sampling scheme, while key distributions do not affect estimator unbiasedness.
  • For sample size k, κ is the (k+1)-th largest key, while κ′_i is the k-th largest key among all elements except i.
  • The proof shows that every sampled element contributes an unbiased estimate of E[f(i)], so summing those contributions gives an unbiased total estimator.
  • An element is sampled exactly when its key exceeds the relevant threshold, with the reverse implication also established through the order-statistic characterization.
Loading 1903.06059v2…