Sampling Methods in Practice: Temperature Scaling with Top-p/Top-k

Multinomial sampling is the core alternative to greedy decoding: instead of always taking the single highest-probability token, the next token is sampled from the model's probability distribution over the vocabulary — a token with 30% probability genuinely has roughly a 30% chance of being picked, not a guarantee either way. This is what unlocks non-deterministic, varied output — running the exact same prompt twice can produce two different (both perfectly valid) continuations.

In practice, raw multinomial sampling straight from the model's unmodified distribution is rarely used on its own, because the unmodified distribution usually has a long tail of very low-probability, often nonsensical tokens that can still occasionally get sampled purely by chance, producing incoherent output. This is exactly why temperature, top-p, and top-k exist as companions to sampling, not as separate decoding strategies — they reshape or restrict the distribution BEFORE the sampling step actually happens.

Temperature reshapes the whole distribution via the softmax function — pushing it toward sharper/more-confident (low temperature) or flatter/more-uniform (high temperature) before any token gets sampled. Top-k restricts sampling to only the k most probable tokens, cutting off that problematic long tail entirely. Top-p (nucleus sampling) restricts sampling to the smallest set of tokens whose cumulative probability covers p, adaptively widening or narrowing the candidate pool based on how confident the model currently is.

The practical pattern engineers actually use: combine temperature with EITHER top-k or top-p (rarely both at aggressive settings simultaneously, since they compound). A common, solid default for general-purpose chat is temperature around 0.7-0.9 combined with top-p around 0.9 — sample-with-reshaping, not raw greedy and not raw unmodified sampling. This combination captures the genuine benefit of sampling (diversity, avoiding the repetition greedy decoding is prone to) while avoiding sampling's biggest risk (occasionally picking a nonsensical low-probability token from deep in the tail).

import numpy as np

def softmax_with_temperature(logits: np.ndarray, temperature: float) -> np.ndarray:
    scaled = logits / max(temperature, 1e-8)
    exp = np.exp(scaled - np.max(scaled))
    return exp / exp.sum()

def top_p_filter(probs: np.ndarray, p: float) -> np.ndarray:
    sorted_idx = np.argsort(probs)[::-1]
    cumulative = np.cumsum(probs[sorted_idx])
    cutoff = np.searchsorted(cumulative, p) + 1
    keep_idx = sorted_idx[:cutoff]
    filtered = np.zeros_like(probs)
    filtered[keep_idx] = probs[keep_idx]
    return filtered / filtered.sum()

def sample_with_temperature_and_top_p(
    logits: np.ndarray, temperature: float, top_p: float, rng: np.random.Generator,
) -> int:
    # The practical pattern: reshape with temperature, restrict with top-p,
    # THEN sample — never sample from the raw, unmodified distribution
    probs = softmax_with_temperature(logits, temperature)
    probs = top_p_filter(probs, top_p)
    return int(rng.choice(len(probs), p=probs))

rng = np.random.default_rng(0)
logits = np.array([4.0, 3.5, 1.0, 0.5, 0.2, -1.0, -2.0])
for _ in range(5):
    print(sample_with_temperature_and_top_p(logits, temperature=0.8, top_p=0.9, rng=rng))