[{"content":"How a language model turns last-step logits into the next token. The model itself only produces a score per vocabulary id. Sampling decides which id is actually emitted.\nInput: logits with shape [batch, vocab_size].\nOutput: next token ids with shape [batch, 1].\nTypical knobs: temperature, top_k, top_p. Common defaults for greedy decoding: temperature = 0, top_k = 0 (off), top_p = 1.0 (off).\nPipeline Each step below is optional except the first branch and the final sample. Greedy (temperature == 0) skips the rest.\nlogits → temperature == 0? argmax → done → logits / temperature → top-k mask (optional) → top-p / nucleus mask (optional) → softmax → multinomial sample 1 token 1. Greedy (temperature \u0026lt;= 0) Pick the highest logit. Ranking is deterministic; no softmax, no randomness. z_i is the logit for token i, and V is the vocabulary.\n$$ \\hat{y} = \\arg\\max_{i \\in V} z_i $$return torch.argmax(logits, dim=-1, keepdim=True) 2. Temperature (logits = logits / temperature) Write scaled logits as z_i / T with T \u0026gt; 0. Then apply softmax. Textbook softmax is the T=1 case:\n$$ \\mathrm{softmax}(z)_i = \\frac{\\exp(z_i)}{\\sum_{j \\in V} \\exp(z_j)} $$Temperature is the same function on z/T:\n$$ p_i = \\mathrm{softmax}(z / T)_i = \\frac{\\exp(z_i / T)}{\\sum_{j \\in V} \\exp(z_j / T)} $$Grey bars = softmax(z) (T=1, before). Colored bars = softmax(z/T) (after). Four logits z=(10,8,6,3). Ranking stays A \u0026gt; B \u0026gt; C \u0026gt; D; only the heights change. T\u0026gt;1 (right) flattens the distribution: tokens that were low-probability become more likely to be sampled. T\u0026lt;1 (left) does the opposite — mass piles on the top token.\nGrey dashed line = input logits z (same for every T). Colored curves = softmax(z/T). T only changes how peaked the distribution is: smaller T → sharper peak on the top token; larger T → flatter, so other tokens get more probability. Argmax stays index 10.\nWhen T \u0026gt; 0 the sampler does not take argmax; it divides logits by T and later samples from softmax. Smaller T makes the distribution peakier (almost always the top token). Larger T makes it flatter (lower-ranked tokens get sampled more often). That is why different T can yield different text.\nlogits = logits / temperature # softmax runs after optional top-k / top-p 3. Top-k (top_k \u0026gt; 0) Temperature already ran. Top-k cuts by rank: keep the k largest logits, send the rest to -inf. z_(k) is the k-th largest value (a threshold, not the top-1). Anything strictly below it is dropped. Ties at the cutoff can keep more than k tokens.\n$$ \\tilde{z}_i = \\begin{cases} z_i / T \u0026 \\text{if } z_i \\ge z_{(k)} \\\\ -\\infty \u0026 \\text{otherwise} \\end{cases} $$Then exp(z̃_i) = 0 for dropped tokens, so they get probability 0 after softmax. Survivors are renormalized to sum to 1. top_k=0 disables this step.\nkth = torch.topk(logits, k, dim=-1).values[..., -1:] logits = logits.masked_fill(logits \u0026lt; kth, -inf) 4. Top-p (0 \u0026lt; top_p \u0026lt; 1) Top-p cuts by probability mass, not by count. Sort remaining tokens by probability descending: p_(1) ≥ p_(2) ≥ ⋯. Walk down the list; keep the prefix whose running sum is still ≤ p; drop the rest. How many tokens survive is dynamic: a peaky softmax fills p with few ids; a flat softmax needs many.\n$$ V_p = \\bigl\\{ (1),\\ldots,(m) \\bigm| m = \\max\\{ m' : \\textstyle\\sum_{i=1}^{m'} p_{(i)} \\le p \\} \\bigr\\} $$The implementation always keeps the top-1 token (mask[..., 0] = False), then sets z̃_(i) = -∞ for i \\notin V_p. top_p=1.0 disables this step (V_p = V). Top-k and top-p can both be on: k first, then p.\nsorted_logits, sorted_idx = torch.sort(logits, descending=True) probs = softmax(sorted_logits) cum = cumsum(probs) mask = cum \u0026gt; top_p # drop the tail after cumulative mass exceeds p mask[..., 0] = False # always keep the top-1 token # scatter masked scores back to vocab order 5. Sample After masking, renormalize with softmax and draw one token. This is not argmax. Greedy is only the T ≤ 0 branch.\n$$ p_i = \\frac{\\exp(\\tilde{z}_i)}{\\sum_{j \\in V} \\exp(\\tilde{z}_j)}, \\qquad \\hat{y} \\sim \\mathrm{Categorical}(p) $$probs = softmax(logits) return torch.multinomial(probs, num_samples=1) Full sample funciton A complete PyTorch sampler matching the pipeline above. Any causal LM can call this once per generated token.\ndef sample( logits: torch.Tensor, *, temperature: float, top_k: int, top_p: float, ) -\u0026gt; torch.Tensor: if temperature \u0026lt;= 0: return torch.argmax(logits, dim=-1, keepdim=True) logits = logits / temperature if top_k \u0026gt; 0: kth = torch.topk(logits, min(top_k, logits.size(-1)), dim=-1).values[..., -1:] logits = logits.masked_fill(logits \u0026lt; kth, torch.finfo(logits.dtype).min) if 0.0 \u0026lt; top_p \u0026lt; 1.0: sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) probs = torch.softmax(sorted_logits.float(), dim=-1) cum = torch.cumsum(probs, dim=-1) mask = cum \u0026gt; top_p mask[..., 0] = False sorted_logits = sorted_logits.masked_fill(mask, torch.finfo(logits.dtype).min) logits = torch.zeros_like(logits).scatter_(-1, sorted_idx, sorted_logits) probs = torch.softmax(logits.float(), dim=-1) return torch.multinomial(probs, num_samples=1) ","permalink":"https://notes.leo-draw.xyz/posts/token-sampling/","summary":"\u003cp\u003eHow a language model turns last-step \u003cstrong\u003elogits\u003c/strong\u003e into the next token. The model itself only produces a score per vocabulary id. Sampling decides which id is actually emitted.\u003c/p\u003e\n\u003cp\u003eInput: \u003ccode\u003elogits\u003c/code\u003e with shape \u003ccode\u003e[batch, vocab_size]\u003c/code\u003e.\u003cbr\u003e\nOutput: next token ids with shape \u003ccode\u003e[batch, 1]\u003c/code\u003e.\u003c/p\u003e\n\u003cp\u003eTypical knobs: \u003ccode\u003etemperature\u003c/code\u003e, \u003ccode\u003etop_k\u003c/code\u003e, \u003ccode\u003etop_p\u003c/code\u003e. Common defaults for greedy decoding: \u003ccode\u003etemperature = 0\u003c/code\u003e, \u003ccode\u003etop_k = 0\u003c/code\u003e (off), \u003ccode\u003etop_p = 1.0\u003c/code\u003e (off).\u003c/p\u003e\n\u003ch2 id=\"pipeline\"\u003ePipeline\u003c/h2\u003e\n\u003cp\u003eEach step below is optional except the first branch and the final sample. Greedy (\u003ccode\u003etemperature == 0\u003c/code\u003e) skips the rest.\u003c/p\u003e","title":"Token Sampling"},{"content":"","permalink":"https://notes.leo-draw.xyz/visitors/","summary":"Where readers have been","title":"Visitors"}]