Beam Search
Greedy decoding picks the best next token. Beam search keeps k live hypotheses so a slightly worse token now can still win the sentence later.
Core concept
Beam search is a bounded breadth-first walk over the tree of possible output sequences.
At every decoding step the model gives a probability over the whole vocabulary. Greedy decoding takes the
argmax and never looks back — one bad early commit poisons the rest of the sentence. Exhaustive search
would consider every sequence, which is V^T and hopeless. Beam search sits between them: expand
all k current hypotheses, score every continuation, then keep only the top k and
throw the rest away.
Because probabilities multiply, we sum log probabilities instead — it keeps the arithmetic stable and turns the objective into a plain additive score.
Key components
Beam width k
How many hypotheses survive each step. k=1 collapses to greedy. Translation typically uses 4–10; beyond that, quality plateaus and often degrades.
Cumulative log-prob score
Each hypothesis carries the running sum of log P. Every candidate extension is scored as parent score + log P(next token), so ranking is a single sort.
Pruning step
k hypotheses × V vocabulary = k·V candidates. Sort, keep k, discard the rest. This prune is the entire reason the search stays tractable — and the reason it is not optimal.
Termination & length penalty
A hypothesis that emits <eos> is moved to a finished set and stops expanding. Stop when
k are finished, then re-rank the finished set with the α length penalty.
How it works
- InitializeStart with one hypothesis: the empty sequence (or
<bos>) with score 0. - ExpandRun the decoder on each live hypothesis. Every one produces a distribution over V tokens, giving k·V candidate extensions.
- Scorecandidate score = parent score + log P(token). Additive, because we are in log space.
- PruneSort all k·V candidates, keep the top k. Everything else is gone permanently — no backtracking.
- Retire finished hypothesesAny beam that emitted
<eos>leaves the beam and joins the finished pool; refill the beam from the next-best candidates. - Repeat, then re-rankLoop until k hypotheses are finished or max length is hit. Apply the length penalty to the finished pool and return the best.
A concrete trace (k = 2)
| Step | Candidate | Score (Σ log P) | Fate |
|---|---|---|---|
| t=1 | "the" | -0.22 | keep |
| t=1 | "a" | -0.51 | keep |
| t=1 | "one" | -2.90 | pruned |
| t=2 | "a cat" | -0.74 | keep |
| t=2 | "the cat" | -0.95 | keep |
| t=2 | "the dog" | -1.60 | pruned |
Read the interesting row: at t=1 greedy would have locked in "the".
Beam search kept "a" too — and by t=2 "a cat" is the leading hypothesis. That
recovery is the whole value proposition. Cost: O(k · V · T) scoring work and k parallel
decoder states in memory.
Where you actually meet it
- Machine translationWhere beam search was made famous. Beam 4–5 is the standard NMT default; the length penalty matters enormously here.
- Speech recognitionBeam search over the acoustic lattice, usually fused with a language-model score and pruning thresholds.
- Image captioning & summarizationEncoder–decoder models where a fluent full sentence matters more than per-token confidence.
- Code & structured generationConstrained beam search prunes candidates that violate a grammar, so every surviving beam stays syntactically valid.
The trap most people hit
Bigger beams are not better beams
Intuition says a wider beam searches more and must win. Empirically, translation quality peaks around k≈4–10 and then falls — the beam search curse. Wider search finds sequences with genuinely higher model probability, and those turn out to be short, generic, degenerate outputs. The search got better; the objective was flawed. This is also why open-ended text generation usually abandons beam search entirely for sampling (top-k, nucleus): for creative text, the highest-probability sequence is the boring one.
Checkpoint — answer before you move on
- Why do we sum log probabilities instead of multiplying raw probabilities, and what breaks if you multiply? Hint: think about 30 numbers below 1 multiplied together in float32.
- Without a length penalty, beam search systematically prefers one kind of output. Which kind, and why does the scoring function cause it? Hint: every additional token adds a negative term to the score.
- Beam search is not guaranteed to find the highest-scoring sequence. Point to the exact step where optimality is lost, and describe a case where the true best sequence is thrown away. Hint: it is the step that makes the algorithm affordable.