I used to think of categorical sampling as bookkeeping.

Normalize weights. Draw a uniform random number. Walk along the cumulative sum until the mass crosses it.

That implementation is fine. It is also the dullest mental model. There is another one:

give every candidate a private clock;
the first clock to ring is the sample

The clock view is not a metaphor. If item \(i\) has positive weight \(w_i\) and its clock time is

\[T_i = E_i / w_i, \quad E_i \sim \operatorname{Exp}(1),\]

then

\[P(T_i = \min_j T_j) = \frac{w_i}{\sum_j w_j}.\]

A heavier item simply has a faster clock. The first arrival is a categorical draw.

Now take logs:

\[-\log T_i = \log w_i - \log E_i.\]

Because \(-\log E_i\) has a standard Gumbel distribution, the same race can be written as

\[\arg\max_i \{ \log w_i + G_i \}, \quad G_i \sim \operatorname{Gumbel}(0,1).\]

That is the Gumbel-max trick. It converts sampling into optimization. Draw one independent noise value per item, add it to the log weight, and take an argmax. The maximum is random, but the operation is deterministic after the perturbation.

This tiny identity is one of those ideas that keeps reappearing under different names: random utility models, Luce choice, perturb-and-MAP, weighted reservoir sampling, stochastic beam search. The interesting part is not that it saves a few lines of sampling code. The interesting part is that it lets “sample” and “search” share the same object: a noisy score table.

Softmax as a Race

Suppose a language model gives unnormalized logits \(\phi_i\) for the next token. After temperature, the categorical law is

\[p_i = \frac{\exp(\phi_i)}{\sum_j \exp(\phi_j)}.\]

Gumbel-max says we may sample from this distribution by drawing independent Gumbel noises and returning

\[\arg\max_i \{\phi_i + G_i\}.\]

This is more than a trick for avoiding numerical underflow. It says that a softmax can be represented as a competition among noisy utilities. That connection goes back to choice theory. Luce’s choice axiom gives the ratio form of choice probabilities.1 Yellott showed the relationship between Luce’s axiom, Thurstone-style random utility, and the double-exponential distribution, which is another name tied to the Gumbel law.2

The exponential clock proof is short enough to keep in your head. Independent exponential clocks have two useful properties:

  1. The minimum of exponentials is exponential with rate equal to the sum of the rates.
  2. The chance that clock \(i\) is the first one equals its rate divided by the sum of rates.

With rates \(w_i\), this is exactly categorical sampling. With \(w_i=\exp(\phi_i)\), it is exactly softmax sampling.

The Gumbel form is just the log-space version of the same race. Let \(E_i \sim \operatorname{Exp}(1)\). Then

\[P(-\log E_i \le x) = P(E_i \ge e^{-x}) = \exp(-e^{-x}),\]

which is the standard Gumbel CDF. So sorting \(\phi_i - \log E_i\) descending is the same as sorting \(E_i / \exp(\phi_i)\) ascending.

One uniform random number per item is enough to build both views:

E_i = -log(U_i)
G_i = -log(-log(U_i))
clock_i = E_i / exp(phi_i)
key_i = phi_i + G_i

The key with the largest value and the clock with the smallest value name the same item.

Do Not Stop at the First Bell

The next move is the part that makes the identity feel larger than a categorical sampler.

Do not stop at the maximum.

Sort all the noisy keys:

\[\phi_i + G_i.\]

The first item is a categorical sample. The second item is what you would have sampled from the remaining items after removing the first. The first \(k\) items are an exact weighted sample without replacement. Their ordered probability is

\[P(i_1,\ldots,i_k) = \prod_{t=1}^{k} \frac{\exp(\phi_{i_t})}{ \sum_{j \notin \{i_1,\ldots,i_{t-1}\}} \exp(\phi_j) }.\]

That is the Plackett-Luce ranking law. In machine learning papers this sorted perturbation view is often called Gumbel-top-k. Kool, van Hoof, and Welling used it to construct stochastic beam search: a way to sample unique sequences without replacement while still using beam-search-like structure.3 The broader Gumbel process view in Maddison, Tarlow, and Minka’s A* sampling paper makes the same philosophical move for continuous spaces: sampling can be phrased as optimizing a random perturbation.4

This is not the same thing as drawing \(k\) independent samples from the categorical distribution. Independent draws collide. If the head item has probability 0.47 and we draw four times with replacement, a lot of the budget is spent rediscovering the same item.

The lab below makes that failure visible. It creates a small next-token-like categorical law, adds Gumbel noise, and checks the empirical first-pick and top-k inclusion frequencies against the exact formulas.

Exact law Gumbel simulation Exponential race With-replacement inclusion Selected top-k order

Deterministic synthetic experiment. The labels are toy candidate tokens. The first panel checks the categorical law; the second shows that Gumbel keys and exponential clocks sort identically; the third compares exact without-replacement inclusion probabilities with a Monte Carlo estimate.

With the default settings, the distribution has only about 4.63 effective choices even though there are eight candidates. The largest candidate has probability 47.2%.

For k=4, independent with-replacement sampling would use only about 2.66 distinct items in expectation. In other words, about 33.5% of the sample slots are wasted on repeats. Sorting one Gumbel key per item spends all four slots on distinct candidates, while still favoring high-probability candidates in the right Plackett-Luce proportions.

The numerical audit is deliberately boring:

default first-pick total variation:     0.0171
default top-k inclusion mean abs error: 0.0051
clock order matches Gumbel order:       yes

The point is not that Monte Carlo error is zero. It is that the errors shrink around the exact categorical and Plackett-Luce laws, and the clock/key ordering is algebraically identical for every trial.

Top-k, Unfortunately, Means Two Things

There is an annoying naming trap here.

In LLM decoding, “top-k sampling” often means:

keep the k highest-probability tokens;
renormalize;
draw one token from that truncated distribution

Gumbel-top-k means something different:

add one random Gumbel key to every candidate;
take the k largest noisy keys

The first procedure returns one token after a deterministic truncation. The second returns a random ordered set of \(k\) distinct candidates from the full distribution. They answer different questions.

The decoding procedure asks, “How should I restrict the support before drawing one continuation?” The Gumbel-top-k procedure asks, “How can I draw several unique alternatives according to the same original weights?”

That distinction matters for sequence generation. If you need one continuation, ordinary categorical sampling is enough. If you need a diverse list of candidate continuations without duplicates, repeatedly sampling with replacement is a noisy way to buy diversity. Gumbel-top-k is a cleaner primitive.

Why Search Code Likes This

Beam search is deterministic: keep the best partial sequences under the model. Sampling is stochastic: roll the model forward according to its probabilities.

Gumbel perturbations make a bridge between the two. If every complete sequence has a log probability and receives an independent Gumbel perturbation, then the top perturbed sequence is a sample from the model. The top \(k\) perturbed sequences are samples without replacement.

Of course, a language model has exponentially many complete sequences. You cannot materialize every sentence, add noise, and sort. Stochastic beam search uses structure in the sequence model to avoid doing that explicitly.3 A* sampling uses bounds and a Gumbel process to search continuous or enormous spaces for the maximum of a random perturbation.4

The common shape is:

randomness first, optimization second

Once the perturbations are drawn, the rest can look like search. That is a useful engineering pattern because search code can exploit structure, pruning, bounds, and batching, while the perturbation gives the result a sampling interpretation.

Trick Still Has a Boundary

Gumbel-max is exact when the candidates you perturb are exactly the domain you want to sample from. If you only perturb local factors in a structured model, you generally get an approximation unless the construction preserves the right global perturbation law. That is the warning behind perturb-and-MAP methods: turning sampling into optimization is powerful, but the perturbation must match the object being optimized.

It also does not make low-temperature distributions diverse. If the top item has nearly all the mass, sampling without replacement will still include it almost every time. The lab shows this if you lower temperature: the top candidate dominates the first-pick law, and it remains nearly guaranteed to appear in the top-k set. No sampling trick can create probability mass that the model did not assign.

And it does not remove the usual numerical obligations. In real code, use stable logits, avoid taking logs of exact zero, and remember that a Gumbel perturbation has a heavy right tail. That tail is the whole point. Occasionally a low-weight candidate receives a very large positive shock and wins.

The Picture I Keep

I like the race picture because it makes several facts feel like one fact.

Temperature changes the clock rates. Lower temperature makes the fastest clock much faster. Higher temperature makes the race more even.

Categorical sampling is the first clock.

Sampling without replacement is the order in which the clocks ring.

Gumbel noise is the log-space record of those same clocks.

That is the whole trick. A sample is not only a draw from a distribution. It can also be the winner of a noisy optimization problem. Once you see that, a lot of algorithms that look like separate species start to share a skeleton.

  1. R. Duncan Luce, Individual Choice Behavior: A Theoretical Analysis, Wiley, 1959. ↩

  2. John I. Yellott, “The Relationship Between Luce’s Choice Axiom, Thurstone’s Theory of Comparative Judgment, and the Double Exponential Distribution”, Journal of Mathematical Psychology, 15(2):109-144, 1977. ↩

  3. Wouter Kool, Herke van Hoof, and Max Welling, “Stochastic Beams and Where To Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement”, ICML 2019. ↩ ↩2

  4. Chris J. Maddison, Daniel Tarlow, and Tom Minka, “A* Sampling”, NeurIPS 2014. ↩ ↩2