Beam Search: Keeping Top-k Sequences Alive, and Its Trade-offs vs. Greedy

Both greedy decoding and simple sampling share a real limitation: they only ever focus on the most immediate token to be generated, when what you actually care about is maximizing the probability of the WHOLE sequence. Properly maximizing the whole sequence would mean knowing the future conditionals — what comes after each candidate choice — but during decoding, at any given step, you only know the probabilities for the immediate next token, not how the sequence will unfold from there.

Beam search is a practical approximation to that true global maximization. Instead of committing to a single path (as greedy does) or sampling one path at random (as multinomial sampling does), beam search expands and keeps alive the top-k PARTIAL sequences — the 'beam' — at every single step, rather than betting everything on one. This matters because some beams that started with a slightly less-probable initial token can go on to produce a much higher-probability COMPLETE sequence than a beam that greedily grabbed the best-looking token at step one. By keeping k alternatives alive simultaneously instead of pruning down to one immediately, beam search explores meaningfully more of the probability tree than any single-path strategy ever could.

This comes at a real, direct cost: beam search requires roughly k times the compute of greedy decoding at every step, since you're now running the model forward for k candidate continuations instead of just one. There's also a well-known qualitative trade-off — beam search tends toward safer, more generic, higher-average-probability sequences (which is exactly what makes it good at correctness-focused tasks), but that same safety bias can make its output feel less surprising or creative than sampling-based approaches.

Beam search is widely used in tasks like machine translation, where getting the overall sequence objectively correct matters far more than open-ended creativity — you want the single best, most faithful translation of a sentence, not a diverse set of stylistically varied options. It's a much weaker fit for chat, storytelling, or brainstorming, where the safety bias it introduces works directly against the diversity those tasks actually want.

import numpy as np

def beam_search(model_step_fn, initial_tokens: list[int], beam_width: int = 3, max_new_tokens: int = 10):
    """Keeps the top `beam_width` (tokens, cumulative_log_prob) sequences
    alive at each step, instead of committing to a single greedy path."""
    beams = [(list(initial_tokens), 0.0)]  # start with one beam, score 0

    for _ in range(max_new_tokens):
        candidates = []
        for tokens, score in beams:
            probs = model_step_fn(tokens)
            top_k_ids = np.argsort(probs)[-beam_width:]  # candidate next tokens
            for tok_id in top_k_ids:
                new_score = score + np.log(max(probs[tok_id], 1e-12))
                candidates.append((tokens + [int(tok_id)], new_score))

        # Keep only the overall top `beam_width` sequences across ALL beams —
        # this is what lets a slightly-behind beam overtake others later
        beams = sorted(candidates, key=lambda c: c[1], reverse=True)[:beam_width]

    return max(beams, key=lambda b: b[1])  # highest-scoring complete sequence

def toy_model(tokens: list[int]) -> np.ndarray:
    return np.array([0.4, 0.35, 0.15, 0.1])  # fixed toy distribution for illustration

best_sequence, best_score = beam_search(toy_model, [1], beam_width=3, max_new_tokens=5)
print("Best sequence:", best_sequence, "log-prob:", round(best_score, 3))