Sampling Strategies: Greedy, Multinomial, Beam Search & Contrastive Search

Temperature, top-k, and top-p all shape the probability distribution the model samples from — but there's a separate, related question: given that distribution, what actual STRATEGY do you use to pick a token at each step? Different strategies here lead to very different output styles, independent of how the distribution itself was shaped.

Greedy strategy is the simplest: at every step, just pick the single token with the highest probability and move on (autoregress). It's fast and deterministic, but it's often not ideal — it tends to produce repetitive sentences, because once the model starts down a high-probability-but-generic path, greedy decoding has no mechanism to escape it.

Multinomial sampling instead samples from the probability distribution rather than always taking the top token — this is exactly where temperature comes in, controlling how random that sampling actually is.

Beam search addresses a problem that both greedy and multinomial sampling share: they only ever optimize for the single most immediate next token, when really you care about maximizing the probability of the WHOLE sequence. Properly maximizing the whole sequence would require knowing future conditionals — what comes after each candidate — but during decoding you only know probabilities for the immediate next step, not the downstream continuation. Beam search approximates the true global maximization by keeping the top-k partial sequences (the 'beam') alive at each step, rather than committing to just one. Some beams that started with a less-probable token can end up leading to a much higher-probability full sequence — by keeping alternatives alive, beam search explores more of the probability tree than a single greedy path ever could. This is widely used in tasks like machine translation, where correctness of the whole sequence matters more than raw creativity.

Contrastive search is the newest of the four, and it balances fluency with diversity by explicitly penalizing repetitive continuations. At each step, it considers candidate tokens and applies a penalty based on how similar a candidate is to what's already been generated, then selects the token that best balances raw probability against that diversity penalty. This directly targets the 'stuck in a loop' problem that greedy decoding is prone to, while still keeping overall coherence high — unlike multinomial sampling, which can wander incoherent purely by chance.

# Toy illustration of the 4 strategies over a tiny fixed distribution —
# real decoding runs this per-step across the model's full vocabulary.
import numpy as np

vocab = ["the", "a", "cat", "dog", "runs"]
probs = np.array([0.4, 0.3, 0.15, 0.1, 0.05])

def greedy(probs: np.ndarray) -> int:
    return int(np.argmax(probs))  # always the single highest-probability token

def multinomial(probs: np.ndarray, rng: np.random.Generator) -> int:
    return int(rng.choice(len(probs), p=probs))  # samples proportional to probability

def beam_search_step(beams: list[tuple[list[int], float]], probs: np.ndarray, k: int):
    """Expand each existing beam by the top candidates, keep only the best k overall."""
    candidates = []
    for tokens, score in beams:
        for tok_id in np.argsort(probs)[-k:]:
            candidates.append((tokens + [int(tok_id)], score + np.log(probs[tok_id])))
    return sorted(candidates, key=lambda c: c[1], reverse=True)[:k]  # keep top-k by sequence score

rng = np.random.default_rng(0)
print("Greedy pick:", vocab[greedy(probs)])
print("Multinomial pick:", vocab[multinomial(probs, rng)])
print("Beam search (k=2) after 1 step:", beam_search_step([([], 0.0)], probs, k=2))