Sampling

A transformer does not output text directly. It outputs logits: one raw score for every token in the vocabulary. Sampling is the machinery that takes that score vector, applies policy, and turns it into one token ID. That token might be chosen deterministically, or it might be drawn from a filtered distribution. Either way, this is the stage where a huge set of possibilities collapses into one emitted token.

Plain English: what sampling does

For our running example, the prompt is “What is the capital of France?”. After the forward pass, the language model head produces a logit vector z ∈ ℝ^V, where V is the vocabulary size. One entry might correspond to “ Paris”, another to “ Lyon”, another to “ London”, and so on through tens of thousands of candidates. Sampling is not about inventing meaning. It is about translating that score vector into an action.

In a production engine, the sampling pipeline is usually short but deliberate: adjust logits for penalties, divide by temperature, apply any hard masks or grammar rules, optionally filter to a candidate subset, run a numerically stable softmax, then either take the maximum or draw one sample. If the chosen token completes a stop sequence, emit nothing further. If it does not, append it to the context and continue.

A concrete score vector for “What is the capital of France?”

The full vocabulary may be 32K, 50K, or 128K tokens, but the sampler usually reasons about the same vector-shaped object regardless of size. To keep the arithmetic readable, we will show a small visible slice of the distribution. Suppose the model has already projected the final hidden state and produced the following candidate logits for the first generated token:

The numbers are illustrative, not magical. A difference of only a few logit units is enough to make “ Paris” dominate after exponentiation. That is why softmax matters: it converts relative score gaps into a probability distribution that can be interpreted and manipulated consistently.

The full sampling pipeline

logits[V]
  ↓
apply_presence_frequency_or_repetition_penalties()
  ↓
divide_by_temperature()
  ↓
apply_hard_masks(grammar, banned tokens, EOS rules)
  ↓
filter_candidates(top-k / top-p / min-p / typical / beam policy)
  ↓
stable_softmax()
  ↓
argmax_or_random_draw(seed, RNG state)
  ↓
chosen token ID
  ↓
check stop sequences / structured-output completion
  ↓
append to context and continue

This pipeline is operationally cheap compared with the transformer layers that produced the logits, but it is strategically important. Small policy changes here can completely change the style, determinism, safety, and usefulness of the output. A model with identical weights can look terse, repetitive, playful, brittle, or reliable depending on sampler settings.

Maths: logits, temperature, softmax, and numerical stability

Let zᵢ be the logit for token i. With temperature τ, the probability assigned to token i is:

pᵢ = exp(zᵢ / τ) / Σⱼ exp(zⱼ / τ)

Temperature acts before softmax by dividing logits. Lower τ exaggerates differences and makes the distribution sharper. Higher τ compresses differences and makes the distribution flatter. The order of token scores does not change when you divide by a positive scalar, but the gaps between them do.

In real implementations you never compute softmax from raw logits directly. You first subtract the maximum scaled logit. That leaves the probabilities unchanged because the same constant is removed from every numerator and denominator term:

m = maxⱼ (zⱼ / τ)
pᵢ = exp((zᵢ / τ) - m) / Σⱼ exp((zⱼ / τ) - m)

Why does this help? Because exp(80) is enormous and can overflow, while exp(0) is safely 1. By shifting the largest term to zero, every other exponent becomes ≤ 1. The distribution is identical; the implementation becomes safe.

// stable softmax over a visible candidate list
max_val = max(logits)
sum_exp = 0
for i in 0..V-1:
    tmp[i] = exp(logits[i] - max_val)
    sum_exp += tmp[i]
for i in 0..V-1:
    probs[i] = tmp[i] / sum_exp

This is still conceptually the same algorithm on GPU, CPU, or TPU. The differences are in vectorization, memory layout, and whether you fuse surrounding steps to avoid extra reads and writes.

Tensor shapes and sampler state

Notice how small these tensors are compared with the model weights and the KV cache. Sampling is not where most memory is spent. But sampling is still part of the serving hot path, so unnecessary allocations, full sorts, or host-device synchronizations can hurt latency.

Temperature tabs: one scalar, large behavioral effect

One subtle point that matters in practice: temperature should be applied to logits, not probabilities. Scaling probabilities after softmax is not equivalent and usually produces the wrong behavior.

Penalty tabs: discouraging reuse without rewriting the model

Penalty order matters. In most engines, you apply penalties before top-k or top-p filtering, because the whole point is to let those penalties influence which tokens survive the filter. If you filter first and penalize later, some candidates never get a fair chance to replace the repeated token you were trying to suppress.

Decoding strategy tabs: the policy layer

Typical sampling is harder to explain at first glance, but the intuition is practical: a token can be high probability and still be “too certain” relative to the overall shape of the distribution, just as a token can be low probability and “too surprising.” Typical decoding tries to stay near the expected surprise level.

Entropy, surprise, and why these policies feel different

Two distributions can have the same top token and still feel very different to decode from. In one case the top token may hold 0.95 probability and everything else is dust. In another case the top token may hold 0.30 and ten alternatives are plausible. Top-k does not know that difference unless K changes. Top-p partially adapts because cumulative mass rises faster in low-entropy distributions and slower in high-entropy ones. Typical sampling goes one step further by explicitly using the information content of each token.

surprise(token i) = -log pᵢ
entropy H(p)      = Σᵢ pᵢ (-log pᵢ)

That distinction matters when you want text that is neither robotic nor chaotic. If the distribution is already sharp, typical decoding may reject extremely unsurprising tokens that would make generation collapse into dull loops. If the distribution is already broad, it may reject bizarre tail tokens that are technically allowed by top-p but feel too off-pattern for the current uncertainty level.

There is no universally best policy. Greedy excels when correctness is tightly coupled to the single highest-probability token. Top-k is simple and predictable. Top-p adapts well to variable confidence. Min-p is a relative confidence gate. Typical sampling is often aesthetically appealing because it respects the shape of uncertainty rather than just its rank ordering. The practical lesson is to think in distribution shape, not in brand names.

A worked numerical walkthrough

Let us walk one step all the way through. Assume there are no structural bans yet, no stop condition yet, and we choose τ = 1, top-k = 3, and stochastic selection with a fixed seed. Start from the visible slice above:

Raw logits:
Paris      6.0
Lyon       4.2
London     3.6
Berlin     2.8
Marseille  2.4
French     1.1
The        0.4

Softmax at τ = 1:
Paris      0.749
Lyon       0.124
London     0.068
Berlin     0.031
Marseille  0.020
French     0.006
The        0.003

Top-k = 3 keeps:
Paris      0.796
Lyon       0.132
London     0.072

If the random draw is u = 0.81, cumulative sampling over the renormalized top-3 distribution yields: Paris occupies [0.000, 0.796), Lyon occupies [0.796, 0.928), and London occupies [0.928, 1.000]. A draw of 0.81 therefore picks “ Lyon”, even though “ Paris” was the single most likely token. This is the entire point of stochastic decoding: allow high-probability alternatives to occasionally win.

Now imagine the next step of generation already contains the token “ Paris” twice because the model wrote a clumsy sentence. With a frequency penalty λ_f = 0.4, the logit for “ Paris” loses 0.8. It falls from 6.0 to 5.2. That does not guarantee it loses, but it narrows the gap enough that alternatives can compete. Penalties do not force diversity; they create breathing room for it.

Seeds and determinism

For stochastic decoding, the seed initializes the random number generator. Same prompt + same model weights + same sampler policy + same seed should produce the same random draws and therefore the same sampled tokens. But real systems add caveats: floating-point nondeterminism, fused kernels, vendor library versions, and batch interactions can all perturb results. Greedy decoding bypasses RNG entirely, but it can still differ across implementations if logits are numerically close and hardware behavior changes the tie.

A good implementation treats determinism as a contract with levels. Level 1: same process, same seed, same batch, same code path. Level 2: same hardware and libraries. Level 3: bitwise reproducibility across environments, which is much harder and often not worth the cost for user-facing chat. The practical promise most systems make is “reproducible enough,” not “identical on every machine forever.”

Stop sequences and generation termination

Sampling is not just about choosing a token; it is also where generation stops. There are several common stop conditions:

Implementers often keep a rolling token window or a byte buffer to detect stop sequences without rescanning the full generated output each step. This is a tiny but real form of caching inside the sampler.

What sampler state persists across steps?

Sampling is stateless only in toy examples. A real engine carries small but important state from token to token. The most obvious piece is the random number generator state used for reproducible stochastic draws. Then come repetition counters for presence, frequency, or repetition penalties. Add stop-sequence matching state, grammar-parser state for structured output, and possibly speculative-decoding acceptance bookkeeping when a draft model is involved.

None of that state is large compared with the KV cache, but it still deserves design care. If you stream tokens to a client, pause, resume, or migrate a request between worker threads, you must move the sampler state with the request. If you batch many requests together, each request needs its own independent sampler state even while weights and kernels are shared. This is one of the reasons inference engines usually separate model state from request state: the weights are global, the sampler and cache metadata are per sequence.

Put differently, sampling is not a pure function of current logits once you enable real policies. The same logit vector can produce different legal actions depending on the request’s history. That is exactly what repetition penalties, grammar masks, and stop rules are supposed to do.

Structured output and grammar-constrained decoding

Structured output is best understood as dynamic masking. At each decode step, an external grammar engine or parser state machine tells the sampler which tokens are legal next moves. Illegal tokens receive -∞ or an equivalent hard ban before softmax. The sampler then runs normally over the remaining legal subset.

This means grammar constraints sit naturally in the sampling pipeline. They do not require retraining the model. They do require careful tokenization-aware engineering: a JSON quote might be one token in one tokenizer and several token fragments in another, so the grammar engine must operate on the actual token vocabulary, not on wishful character abstractions.

Beam search: useful, but less common for chat LLMs

Beam search keeps multiple partial continuations alive at once instead of committing to one token path immediately. At each step, it expands the best B sequences according to cumulative log-probability, then prunes back to the best B. This often improves sequence-level likelihood in tasks like translation or speech recognition, where the best overall sentence may require a locally non-greedy choice early on.

Why is it less common for chat LLMs? Because beam search is more expensive, less diverse than sampling, and often produces overly safe text. User-facing assistants usually benefit more from controlled stochastic decoding than from maximizing sequence likelihood. Still, beam search belongs in the design space because it is the classic reminder that “best next token” is not always “best full sentence.”

Speculative decoding: faster sampling by drafting ahead

Speculative decoding uses a smaller draft model to propose several next tokens, then asks the larger target model to verify them in batches. If the large model agrees, you accept multiple tokens at once and reduce latency. If it disagrees, you fall back to the target model’s choice at the first disagreement point.

This belongs in the sampling chapter because the acceptance test is effectively a decoding policy wrapped around two models. The large model still determines correctness. The small model exists to offer plausible candidates so the expensive model can amortize work over multiple positions.

Cost, hardware, and operation-classification

Compared with attention and FFN matmuls, sampling is small. But small does not mean irrelevant. Decode latency is serial, so even “cheap” steps executed every token deserve engineering attention.

On GPU, the temptation is to leave logits on device and fuse as much as possible so you do not bounce the vocabulary vector back to the CPU just to sample. On CPU, branch-friendly top-k and efficient sparse history structures matter more. On either target, a full O(V log V) sort every token is usually wasteful when a partial selection or threshold-based scan will do.

Optimisations and caching

Although the sampler has no KV cache, it still benefits from small cached state: token occurrence counts for penalties, rolling buffers for stop sequences, parser states for grammar constraints, and RNG state for reproducibility. Engines that implement speculative decoding may also cache draft tokens and acceptance statistics. All of this state is small, but it avoids repeated recomputation on the critical decode path.

Common optimisations include: keep logits resident on the same device as the LM head; fuse temperature, mask, and softmax when possible; use partial selection instead of full sort; maintain sparse repetition maps instead of dense [V] arrays when the vocabulary is huge; and short-circuit directly to argmax when the user requested greedy decoding or temperature zero.

Analogies that help without lying too much

Three analogies are worth keeping. First, roulette wheel: after softmax and filtering, the distribution is a weighted wheel; greedy just always takes the largest slice. Second, thermostat: temperature does not invent new options, it changes how dominant the current favorite feels. Third, bouncer at the door: top-k, top-p, min-p, and grammar constraints decide which candidates even get into the club before the final choice happens.

The analogy to avoid is “the model writes a full answer and then sampling chooses one version.” That is not what happens. The model produces one distribution per position, one token at a time. The future does not exist yet beyond those probabilities.

Common implementation mistakes

If you were implementing this yourself...

The most reliable path is incremental. Implement stable softmax and greedy selection first. Then add seeded categorical sampling. Then add top-k. Then top-p. Then penalties. Then stop sequences. Then grammar masks if you need structured output. Every stage should have tiny deterministic tests using small visible logit vectors like the one in this chapter.

token sample_next(
    float logits[V],
    History history,
    SamplerConfig cfg,
    RNG* rng)
{
    apply_penalties(logits, history, cfg);
    apply_temperature(logits, cfg.temperature);
    apply_masks(logits, cfg.grammar_state, cfg.banned_tokens);

    CandidateSet c = filter_candidates(logits, cfg);   // greedy, top-k, top-p, min-p, typical
    softmax_in_place(c.logits_or_probs);

    token t = cfg.greedy
        ? argmax(c)
        : categorical_sample(c, rng);

    update_stop_state(history, t, cfg.stop_rules);
    update_penalty_state(history, t);
    return t;
}

The engineering mindset is straightforward: the model forward pass produces scores; the sampler is a compact policy engine over those scores. Keep the math honest, keep the data movement small, and keep the semantics explicit. Once you do that, every “creative” or “deterministic” behavior becomes inspectable rather than mystical.