Beam Search Math

Decoding as MAP inference over sequences

An autoregressive model factorizes the probability of an output sequence y = (y_1, …, y_T) given input x as a product of conditionals:

P(y | x) = ∏_t P(y_t | y_1..y_{t-1}, x)

Decoding asks for the single most likely sequence, y* = argmax_y P(y | x). This is a MAP inference problem, and it is intractable to solve exactly: with vocabulary size V and length T there are V^T candidate sequences, and because each token conditions on all previous ones the objective does not decompose into independent per-position choices. You cannot just pick the best token at each step — that is greedy decoding, and it is myopic. Beam search is the standard heuristic that keeps the search affordable while recovering far more of the probability mass than greedy does, without paying the exponential cost of enumerating every path.

The scoring recurrence: turn a product into a sum

Multiplying hundreds of probabilities in [0, 1] underflows to zero in floating point almost immediately, so the score is accumulated in log space. Define the running score of a partial hypothesis as

s(y_1..y_t) = Σ_{k=1}^{t} log P(y_k | y_1..y_{k-1}, x)

Because log is monotonic, the sequence that maximizes P(y | x) also maximizes s, and the product becomes an additive recurrence: s(y_1..y_t) = s(y_1..y_{t-1}) + log P(y_t | …). That single change is what makes beam search practical. Each step only adds one term, scores are comparable across hypotheses of the same length, and the arithmetic stays numerically stable. The logits from the final layer are already turned into log-probabilities by a log_softmax, so in practice you add the model’s output directly — no extra exp is needed on the hot path.

Expand, score, prune: the width-B loop

The algorithm keeps a set of B live hypotheses (the beam). At each step it expands every one of the B beams by every token in the vocabulary, producing B × V candidate extensions; scores each with the additive recurrence; and prunes back to the top B by score. Formally the new beam is

Beam_{t} = top-B over { s(h) + log P(v | h, x) : h ∈ Beam_{t-1}, v ∈ V }

The crucial subtlety is that the top-B is taken over the whole pooled set of B × V candidates, not B separate per-beam choices. That is why beam search can abandon a beam entirely: if one hypothesis has three strong continuations and another has none, the strong hypothesis is allowed to occupy multiple slots and the weak one dies. Greedy decoding is exactly the special case B = 1.

Complexity and the per-step cost

Per step, the dominant costs are one forward pass for each of the B beams and a top-B selection over B × V scores. Over T steps the search cost scales as

O(T · B · C_fwd)  +  O(T · B · V)   (scoring / selection)

where C_fwd is the cost of one model forward pass. The forward passes dominate: the top-B over B × V scalars is a cheap partial sort compared with running a multi-layer transformer B times. This is the headline fact for small models on CPU — beam search multiplies your compute by roughly B. On a GPU those B passes batch cleanly into one call; on a CPU with limited parallelism the wall-clock cost is closer to linear in B, which is why interactive CPU inference usually caps B at 3 to 5.

Length normalization: a necessary correction

Because every added token contributes a negative log-probability (log P ≤ 0), the running score s only ever decreases as a sequence grows. Raw beam search therefore has a strong bias toward short outputs — it will happily emit an early end-of-sequence token because stopping now beats accumulating more negative terms. The fix is to compare hypotheses on a length-normalized score. Google’s NMT length penalty is the common form:

lp(y) = ((5 + |y|) / (5 + 1))^α
score_norm(y) = s(y) / lp(y)        (0 ≤ α ≤ 1)

Consider two finished hypotheses with α = 0.7. Hypothesis A has 4 tokens and raw score s = -6.0; hypothesis B has 10 tokens and s = -11.0. Raw scores pick A (-6.0 > -11.0). Now normalize: lp(A) = (9/6)^0.7 ≈ 1.33, so -6.0 / 1.33 ≈ -4.51; lp(B) = (15/6)^0.7 ≈ 1.90, so -11.0 / 1.90 ≈ -5.79. A still wins here, but note how the gap shrank from 5.0 to 1.28 — the penalty pulls the longer, information-richer hypothesis back into contention. Tuning α is really tuning how aggressively you are willing to trade per-token confidence for length.

The beam-search curse: wider is not always better

Intuitively a larger beam searches more thoroughly and should find a higher-probability sequence — and it does. The paradox, documented across machine translation, is that beyond a modest width (often B ≈ 5) translation quality frequently gets worse even as model probability gets higher. This is the beam-search curse.

The explanation separates two error sources. Search error is failing to find the true argmax; model error is the model assigning high probability to bad sequences. A wider beam reduces search error but exposes model error: with enough width the search discovers degenerate, often empty or ultra-short, high-probability sequences that a narrower beam never reached. In other words, the model’s argmax is not actually the output you want. This is why length normalization and coverage penalties matter so much — they are corrections to the objective, compensating for a model whose raw MAP solution is subtly miscalibrated toward brevity.

Finished hypotheses and early stopping

Beams do not all end at the same step. When a hypothesis emits the end-of-sequence token it is complete and moves to a finished set with its final normalized score; a fresh live beam takes its place so the search still explores B active continuations. The early_stopping flag then decides how patient to be: the aggressive setting stops as soon as B finished hypotheses exist, while the safe setting keeps going until no still-live beam could possibly beat the best finished one — which is provable because a live beam’s score can only decrease. Stopping early is faster but can miss a longer hypothesis that would have won after normalization, so the two settings trade a little quality for latency.

Beam search on CPU: the KV-cache multiplier

The cost that surprises people running small models on CPU is memory, not arithmetic. Each live beam is a distinct sequence, so each needs its own KV cache — the stored keys and values for every layer and head. Running B beams multiplies KV-cache memory by B:

kv_bytes ≈ B × 2 × L × T × d_model × bytes_per_elem

for L layers and current length T (the factor 2 is keys plus values). A width-5 beam is 5× the KV footprint of greedy decoding, which on a memory-constrained CPU box can be the binding constraint well before compute is. There is a further subtlety: when a beam is pruned, its cache must be carried forward for the surviving continuations, so implementations reorder or gather the cache by parent-beam index each step — bookkeeping that is invisible on paper but is where real beam-search code spends its complexity.

Constrained beam search

Sometimes you need guarantees the plain objective cannot give: an output that must contain a specific phrase, obey a grammar, or stay inside a JSON schema. Constrained beam search modifies the pruning step so that hypotheses making progress toward required constraints are protected from being pruned away by higher-probability but non-compliant candidates — lexically constrained decoding, for instance, groups the beam into ‘banks’ by how many constraints each hypothesis has satisfied so partial-constraint sequences survive long enough to complete.

Grammar- or schema-constrained variants go further: at each step a validator masks the logits of any token that would violate the grammar (setting them to -∞) before selection, so every path in the beam is guaranteed well-formed by construction. The math is unchanged — still additive log-probs and top-B — but the candidate set is filtered by a symbolic constraint each step. This is how you get reliable structured output from a model that only ever emits probabilities.

Beam search versus sampling and MBR

Beam search is mode-seeking: it chases the highest-probability sequence. That is the right goal for tasks with a single correct-ish answer — translation, summarization, code — where you want the model’s most confident output. It is the wrong goal for open-ended generation, where the argmax is bland and repetitive; there, sampling methods (top-k, nucleus/top-p, temperature) deliberately trade probability for diversity by drawing from the distribution instead of maximizing over it.

A third option sits between them. Minimum Bayes Risk (MBR) decoding samples many candidate sequences, then picks the one that is most similar on average to all the others under a quality metric — argmax_y Σ_{y'} metric(y, y') P(y') — rather than the one with highest raw probability. MBR sidesteps the beam-search curse precisely because it does not trust the model’s argmax; it trusts the consensus of samples. It costs more forward passes but often beats beam search, which is why the field has partly moved past pure beam search.

Practical defaults and pitfalls

A short field guide. Use beam widths of 3–5 for CPU SLMs; going wider rarely helps and can trigger the curse. Always length-normalize (α ≈ 0.6–0.7) if outputs skew short, and add a no_repeat_ngram_size of 2 or 3 for tasks prone to loops. Remember the two independent multipliers: beam search costs roughly B× the compute and B× the KV-cache memory of greedy decoding, so on a tight box the width you can afford may be set by RAM, not speed.

The classic mistakes: comparing hypotheses of different lengths on raw scores (truncated output); forgetting to reorder the KV cache after pruning (silently corrupt continuations); assuming higher model probability means higher quality (the curse says otherwise); and reaching for beam search on open-ended generation where sampling is the better tool. Beam search is a precise instrument for the mode-seeking tasks it was built for — and the wrong instrument everywhere else.