Top-K Sampling Math

The rule: a rank cut, then a renormalization

Given logits z: [V] and their softmax p, top-k keeps the survivor set S_k of the k highest entries, zeroes everything else, and rescales so the result is a distribution again:

S_k  = { i : rank(z_i) ≤ k }        rank descending
M_k  = Σ_{i ∈ S_k} p_i           retained mass
p’_i = p_i / M_k  if i ∈ S_k,   else 0

Two things deserve more attention than they usually get. First, S_k can be found on the raw logits: softmax is strictly increasing, so ranking by p and ranking by z give the same set. You never need to exponentiate the whole vocabulary. Second, M_k alone describes the damage — the divergence from truncating collapses to log(1/M_k), as the companion sampling article derives. Everything interesting about top-k is a statement about how M_k behaves when k is fixed and the context varies.

Rank is the one thing temperature cannot move

Dividing logits by T > 0 is strictly increasing, so it cannot reorder anything. Top-k therefore selects the identical k tokens at every temperature — but the mass those tokens carry is not identical at all. Take a 32k-vocabulary model with top logits [18.2, 14.6, 13.9, 13.1, 12.4, …] over a bulk near zero, held at k = 40:

Same 40 tokens, same k, and the distortion moves five orders of magnitude. At T = 0.7 the setting is a no-op; at T = 2.0 it deletes 78% of the model’s own belief. The parameter did not change, and the thing it controls changed completely.

The order-statistics view of z_(k)

Write z_(1) ≥ z_(2) ≥ … ≥ z_(V) for the sorted logits. Then z_(k) is an order statistic, and top-k is really a threshold at z_(k) whose numeric value you never specify. If the vocabulary bulk is roughly N(μ, σ) — a decent first approximation away from the head — then the standard quantile estimate gives:

z_(k) ≈ μ + σ · Φ^-1(1 − k/V)

It is accurate: the step above has a N(0, 2) bulk and z_(40) = 6.03, against a predicted 0 + 3.02 · 2 = 6.04. For V = 32000: k = 10 → 3.42σ, 40 → 3.02σ, 100 → 2.73σ, 400 → 2.24σ. A 40× increase in k moves the cut by only 1.2σ, because for a Gaussian tail the threshold grows roughly like sqrt(2 ln(V/k)) — logarithmically. The same collapse shows in the spacings: in the step above, z_(1) − z_(2) = 3.6 logits but z_(40) − z_(100) = 0.50. Moving k from 1 to 2 is a violent change; moving it from 40 to 100 barely moves the boundary at all.

Failure one: the confident step, where k does nothing

Keep that step at T = 1. Its entropy is H = 0.31 nats, so the effective support exp(H) = 1.4 tokens; the leader holds 0.946 and four tokens already cover 99% of the mass. Now apply k = 40. The retained mass is M_40 = 0.9976, and the thirty tokens ranked 11 through 40 hold, after renormalization, a combined 2.7 × 10^-4 — roughly one draw in 3,700.

This is not a harmful setting, it is an inert one. Top-k protected you from nothing, because the model had already crushed its own tail; you paid a selection pass over 32,000 logits to move the sampling distribution by a quarter of a percent. The failure is not bad output but false confidence: clean text on confident steps gets credited to k = 40 when the softmax did all of it.

Failure two: the open step, where k cuts the model off

Now a genuinely open step — the same model mid-narrative, where about a hundred continuations are live. Entropy H = 4.46 nats, effective support exp(H) = 86, leader at p_(1) = 0.074, and it takes 80 tokens to accumulate 0.90 of the mass.

Apply the identical k = 40: M_40 = 0.690. You have deleted 31% of the model’s belief, log(1/0.690) = 0.371 nats of distortion — 155× the distortion the same setting imposed on the confident step. The first dropped token, which the model rated at p_(41) = 0.0072, is now strictly impossible, and over a long generation those unreachable-but-plausible continuations compound into a measurably narrower voice — at precisely the steps where diversity was the point. Fixed k is loosest where the model is certain and tightest where it is not: exactly backwards.

Retained mass is a random variable

Put the two steps side by side and the structural problem is plain: with k pinned at 40, M_k took the values 0.9976 and 0.690 within a single generation from a single model. M_k is not a constant you configured — it is a random variable over contexts, and k only fixes its index, never its value.

That reframes tuning k. You are not choosing a truncation strength; you are choosing a rank and letting the next token’s entropy choose the strength for you. Per-token entropy in real text swings by an order of magnitude — near-deterministic inside a word or a known proper noun, wide open after a sentence boundary — so the induced spread of M_k is enormous. Mass-based and peak-relative rules exist to shrink that spread by construction; each has its own article in this series.

Selection: you do not need a sort

The naive implementation sorts the whole vocabulary, O(V log V), then slices. That is more work than the problem requires, because renormalization does not care about order — you need the survivor set, not a ranking of it.

full sort        O(V log V)   sorts 32,000 to use 40
quickselect      O(V) expected, in place, unordered output
bounded heap     O(V log k)   one pass, k-sized working set

A single-threaded NumPy measurement on a 32,000-element float32 vector: full sort 156 µs, argpartition for k = 40 73 µs — a bit over 2×. The larger win is upstream: because S_k is decidable on raw logits, a top-k-only pipeline never exponentiates the vocabulary. Select in O(V) comparisons, then exp over k values — 40 calls, not 32,000.

GPU, CPU, and vocabulary size

Whatever the algorithm, every top-k must read all V logits, which is why the operation is bandwidth-bound rather than compute-bound. A 128k vocabulary in fp32 is 500 KB per sequence per step; at batch 64 that is 33 MB streamed every step, and CUDA topk answers with a radix select that makes a small number of passes over exactly that data. On CPU a 32k float32 vector is 125 KB — it lives in L2, and at 30 tokens/second a 0.07–0.16 ms selection is well under 1% of a ~30 ms step.

Vocabulary size also silently changes what k means. k = 40 is the top 0.125% of a 32k vocabulary and the top 0.031% of a 128k one — four times stricter in rank, though in logit space the cut deepens only from 3.02σ to 3.42σ. Modest, but not zero: a k tuned against one tokenizer does not transfer unexamined to another.

Worked end to end

The two steps, the same k, all the arithmetic in one place:

              STEP A (confident)        STEP B (open)
H             0.31 nats                 4.46 nats
exp(H)        1.4 tokens                86 tokens
p_(1)         0.9461                    0.0736
cover 0.90    1 token                   80 tokens

k = 40
M_40          0.99764                   0.6901
scale 1/M     1.0024×                  1.4491×
p’_(1)        0.9461/0.99764 = 0.9483   0.0736/0.6901 = 0.1067
last kept     p_(40) = 4.9e-6           p_(40) = 7.6e-3
first dropped p_(41) = 4.8e-6           p_(41) = 7.2e-3
log(1/M)      0.0024 nats               0.371 nats

Renormalization is the tell. On step A it scales survivors by 1.0024 — invisible. On step B it scales them by 1.449, inflating the leader’s share by 45% — confidence the model never claimed. Truncation does not merely remove options; it redistributes the removed mass onto survivors in proportion to what they already had, which systematically favours the head.

Choosing k, and why it is rarely the first knob

k = 1 is greedy decoding written the long way. The conventional band is 20–100, and the order-statistics result above explains why nobody can tell 40 from 50: those cut points differ by 0.07σ.

The honest role for top-k is a rail, not a control: set k generously (50–100) as a hard ceiling on how deep into the tail one unlucky draw can reach, and let a shape-aware rule do the actual shaping. Reach for temperature or a mass-based cut first, because k is the only knob of the three that ignores the numbers entirely.

Three traps. Boundary ties are broken by index in most kernels — deterministic but arbitrary — and in bf16, with 8 mantissa bits, deep-tail ties are routine, so the exact 40th token is dtype- and device-dependent. k = 0 means ‘disabled’ in some stacks and ‘keep everything’ in others. And mask with -inf, never a large finite negative — but never mask all V, because a softmax over all -inf is NaN.