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.