---
title: Gumbel-top-k Sampling
url: https://www.emergentmind.com/topics/gumbel-top-k-trick-9e40f0db-57cf-41af-941e-4da5b209f849
type: topic
---

# Gumbel-top-k Sampling

The Gumbel-top-$k$ trick is a stochastic sampling method that enables efficient, exact, and unbiased sampling of $k$ distinct items without replacement from a categorical or Plackett–Luce distribution, with probability proportional to each item's assigned nonnegative weight or exponential score. This approach generalizes the Gumbel-max trick (which samples a single item) and is foundational for learning-to-rank, diverse text generation, scalable hashing, and low-variance inference in structured probabilistic models. Algorithmic advances, such as FastGM and quasi-Monte Carlo variants, further optimize its computational efficiency and variance properties for large-scale applications.

## 1. Mathematical Foundations and Core Algorithm

Given a collection of $n$ discrete items with associated positive weights $\{w_i \}_{i=1}^n$ (or equivalently, real-valued scores $s_i$, with $w_i = \exp(s_i)$), the Gumbel-top-$k$ trick enables sampling a $k$-tuple $(r_1, \dots, r_k)$ of distinct indices such that

\[
\Pr(r_1, \dots, r_k) = \prod_{j=1}^k \frac{w_{r_j}}{\sum_{i \notin \{ r_1,\dots,r_{j-1} \}} w_i},
\]

which is the Plackett–Luce distribution on ordered $k$-tuples.

The sampling procedure consists of drawing i.i.d. Gumbel$(0,1)$ random variables $g_i$ for each item, computing perturbed keys $K_i = \log w_i + g_i$, and selecting the $k$ indices of the largest $K_i$ in descending order. This process exactly simulates $k$ sequential draws without replacement from the normalized weights, but at the computational cost of a single vector perturbation and sort. The method's unbiasedness and joint distributional correctness follow from the max-stability and memoryless properties of Gumbel and exponential distributions [2110.01515, 2110.15072, 1903.06059].

**Canonical pseudocode** for sampling one $k$-length tuple:

```python
# Inputs: array of scores s[1..n], desired list-length k
for i in 1..n:
    u_i = Uniform(0,1)
    g_i = -log(-log(u_i))
    key_i = s_i + g_i
r = indices of the k largest key_i (descending order)
return r[1..k]
```

## 2. Theoretical Properties and Extensions

The Gumbel-top-$k$ trick inherits the unbiasedness, joint law, and diversity guarantees of the base Gumbel-max approach. For any ordered $k$-tuple $r = (r_1, ..., r_k)$,

\[
\Pr(\text{Top-$k$}(K) = r) = \prod_{j=1}^k \frac{w_{r_j}}{\sum_{i \notin \{ r_1, ..., r_{j-1} \}} w_i},
\]

ensuring consistency with exact sampling without replacement under the Plackett–Luce law [2110.01515, 1903.06059, 2205.06024].

**Marginals** satisfy the expected "no-replacement" probabilities, contrasting with sampling with replacement. The method generalizes naturally to combinatorial and structured sampling tasks by recursively applying the same probabilistic invariance (so-called "stochastic invariant")—including structured domains such as permutations, matchings, trees, and arborescences [2110.15072].

Variants encompassing continuous relaxations (e.g., Gumbel-softmax, Gumbel-Sinkhorn, relaxed top-$k$ masking) allow for differentiable sampling and are used extensively in neural optimization contexts where discrete gradients are otherwise problematic [2502.11116].

## 3. Computational Complexity and Fast Algorithms

The naive Gumbel-top-$k$ implementation samples $n$ Gumbels and finds the top-$k$ via sorting, resulting in $O(n \log k)$ time with a min-heap, or $O(n \log n)$ with full sort [2110.01515].

**FastGM algorithm** [2002.00413, 2302.05176] improves scaling for large $k$ and/or $n$ by recasting the problem as finding the first $k$ arrivals in $k$ servers, each fed by $n$ independent exponential queues (rates $v_i$). Using priority-based pruning and event-driven simulation, it achieves expected $O(k \log k + n)$ time, reducing superfluous Gumbel evaluations for low-weighted items. FastGM is orders of magnitude faster than the naive baseline in high-dimensional and high-sketch-length regimes, preserving both unbiasedness and the exact output distribution.

| Approach        | Time Complexity  | Output            |
|-----------------|-----------------|-------------------|
| Naive           | $O(k n)$        | Exact, unbiased   |
| FastGM          | $O(k \ln k + n)$| Exact, unbiased   |

**Practical remarks:** FastGM requires only $O(1)$ per arrival and supports both subset and structured sketching; its correctness rests on the order-statistics of Poisson/exponential processes matching those of the Gumbel-perturbed sorting [2002.00413, 2302.05176].

## 4. Monte Carlo, Variance Reduction, and QMC

A primary application is estimating expectations under the Plackett–Luce model using Monte Carlo (MC) sampling, where each sample is generated by the Gumbel-top-$k$ trick. However, standard MC suffers from $O(1/N)$ mean-squared error (MSE) convergence in the number of samples.

**Quasi-Monte Carlo (QMC) augmentation** replaces the independent uniform draws by points from a low-discrepancy sequence (e.g., scrambled Sobol), which are then transformed to Gumbels via the inverse CDF. Under mild smoothness assumptions, randomized QMC estimators retain unbiasedness but attain improved MSE decay, up to $O(1/N^2)$ [2205.06024]. This is highly beneficial for learning-to-rank, counterfactual evaluation, and propensity estimation, where empirical studies show consistent (and sometimes dramatic) MSE reduction compared to MC—especially for moderate sample sizes.

| Sampler               | MSE decay         | Empirical effect             |
|-----------------------|-------------------|------------------------------|
| Monte Carlo (MC)      | $O(1/N)$          | Higher variance              |
| Randomized QMC (Sobol)| $O(1/N^2)$        | Lower variance, same compute |

Detailed pseudocode for QMC-augmented Gumbel-top-$k$ is presented in [2205.06024]. Caveats include diminishing gains as $N$ grows, potential high-dimensional curse if $m$ is very large, and necessity of randomized QMC for unbiasedness.

## 5. Structured, Recursive, and Differentiable Variants

The recursive application of Gumbel-top-$k$ underpins combinatorial sampling schemes (e.g., for permutations, matchings, and spanning trees) via the preservation of the stochastic invariant at each recursion or split [2110.15072]. This supports efficient, exact "perturb-and-MAP" sampling for structured domains.

**Differentiable top-$k$ relaxations** are used to enable end-to-end gradient-based optimization in discrete settings, exemplified by document reranking with differentiable top-$k$ masks using Gumbel perturbations and softmax smoothing [2502.11116]. The relaxation involves sampling multiple Gumbel-perturbed softmax vectors and aggregating via elementwise maxima to approximate a smooth top-$k$ mask enabling backpropagation.

The versatility of the Gumbel-top-$k$ family thus spans discrete stochastic optimization, structured inference, efficient hashing, and modern neural end-to-end learning.

## 6. Main Applications and Practical Guidance

The Gumbel-top-$k$ trick is widely adopted across machine learning and information retrieval. Key domains include:

- **Learning-to-rank and recommendation:** Fast unbiased sampling of top-$k$ ranked lists from Plackett–Luce models; low-variance QMC estimators for loss and metric evaluation [2205.06024].
- **Sequence and set generation in NLP:** Sampling diverse sequences or token subsets in neural text generation; used within stochastic beam search for diversity–quality tradeoffs [1903.06059].
- **Sketching and similarity estimation:** Efficient Gumbel-top-$k$ sketches for scalable Jaccard similarity (e.g., ℙ-MinHash), graph embeddings, and active learning [2002.00413, 2302.05176].
- **Structured probabilistic inference:** Sampling from complex structured domains (e.g., matchings, trees) via recursive Gumbel invariants [2110.15072].
- **Differentiable subset selection:** Relaxed Gumbel-top-$k$ for end-to-end training of discrete selectors, such as differentiable document reranking in retrieval-augmented generation (RAG) [2502.11116].

**Practical recommendations** include favoring QMC sequences for low-variance MC when sample budgets are small or moderate, utilizing FastGM for large-scale sketching, and applying differentiable relaxations when gradient-based optimization is required. Scrambled Sobol QMC, $N=2^k$ sample sizes, and analytic smoothing for very small/large scores are advisable for robustness and computational efficiency.

## 7. Limitations, Open Problems, and Future Directions

While the Gumbel-top-$k$ trick achieves exact, scalable, and versatile sampling without replacement, certain settings remain challenging:

- **Very high dimensionality:** QMC benefits diminish as the number of items grows due to curse-of-dimensionality effects [2205.06024].
- **Complex structured relaxations:** For intricate combinatorial objects, the construction and efficient evaluation of the recursive or relaxational sampling map may require sophisticated algorithm design [2110.15072].
- **Bias–variance–efficiency trade-offs:** In differentiable relaxations, exploration/exploitation and temperature scaling must be tuned to balance gradient quality and estimator bias [2502.11116].

Ongoing research explores further structural generalizations, integration with optimal transport relaxations, and principled variance reduction for deeper neural architectures and reinforcement learning settings. The method remains a central tool for scalable, unbiased sampling in modern probabilistic modeling.

Source: https://www.emergentmind.com/topics/gumbel-top-k-trick-9e40f0db-57cf-41af-941e-4da5b209f849