Beam Search — When You Need The Best Sequence

The decoding problem: argmax over sequences

A trained autoregressive model gives you one thing at each step: a probability distribution over the next token, conditioned on everything so far. The probability of a whole sequence Y = (y_1, …, y_T) is the product of those per-step probabilities by the chain rule:

P(Y | x) = ∏_{t=1..T} P(y_t | y_1, …, y_{t-1}, x)

Decoding is the task of finding the sequence that maximizes this: Y* = argmax_Y P(Y | x). That sounds simple until you count the candidates. Each position can be any of V tokens (often 30,000–150,000), and a sequence can be hundreds of tokens long, so the search space is V^T — larger than the number of atoms in the universe for even a short reply. There is no way to enumerate it. Exact search (e.g. via dynamic programming) does not help either, because the model has unbounded context: the probability of y_t depends on the entire prefix, so there is no small shared state to memoize over. Every practical decoder is therefore an approximate search, and beam search is the workhorse among them.

Greedy decoding and why it is myopic

The cheapest approximation is greedy decoding: at each step take the single highest-probability token and commit to it. It costs one forward pass per token and never looks back. The problem is that a locally optimal token is not always on the globally optimal path. Greedy picks the best first move without asking whether it leads anywhere good.

Consider two first tokens: ‘The’ with probability 0.6 and ‘A’ with 0.4. Greedy takes ‘The’. But suppose the best continuation after ‘A’ is nearly certain (0.9) while the best continuation after ‘The’ is uncertain (0.3). Then the two-token sequence ‘A …’ scores 0.4 × 0.9 = 0.36 while ‘The …’ scores 0.6 × 0.3 = 0.18. Greedy chose the worse path because it could not see one step ahead. This is the exact gap beam search is built to close: it hedges by keeping several first moves alive long enough to discover which one actually pays off downstream. Beam search with width 1 is greedy decoding — greedy is simply the degenerate corner of the same algorithm.

Why log-probabilities, not probabilities

Beam search scores partial sequences by summing log-probabilities, never by multiplying raw probabilities. There are two reasons, and both matter in real implementations. First, numerical underflow: a probability is at most 1, so multiplying hundreds of them together drives the product toward zero fast. After a few dozen tokens the product is smaller than the smallest positive float, and it silently rounds to 0.0 — every candidate ties at zero and the search becomes meaningless.

Taking logs converts the product into a sum, because log(a · b) = log a + log b. So the sequence score becomes:

score(Y) = log P(Y | x) = Σ_{t=1..T} log P(y_t | y_<t, x)

Each term log P(·) is a comfortable negative number (a probability of 0.5 is -0.69, of 0.1 is -2.30), and summing them stays well within float range. Second, addition is cheaper and more stable than multiplication, and the model already emits logits that become log-probabilities after a single log_softmax — so the log domain is where the numbers naturally live. Every score below is a sum of these non-positive terms; a higher (less negative) score is better.

The algorithm: beam width B and expand-then-prune

Beam search keeps a fixed number of live hypotheses — the beam width B (also called the number of beams). The beam is a set of partial sequences, each with its accumulated log-prob score. The core loop is expand, score, prune, repeated until the beams finish:

beam ← { (empty prefix, score = 0.0) }
repeat for each step t:
  candidates ← []
  for each hypothesis h in beam:              # EXPAND
    for each token v in vocabulary V:
      new_score = h.score + log P(v | h, x)
      candidates.append( (h + v, new_score) )
  beam ← top-B candidates by score        # PRUNE
until every beam ends in <eos> (or t = max_len)

Each step momentarily grows the frontier to B × V candidates (every live beam times every possible next token), then immediately collapses it back to the best B. That expand-to-BV, prune-to-B rhythm is the whole idea: it is a breadth-limited best-first search where the ‘limit’ is the constant beam width. When a hypothesis emits the end-of-sequence token <eos>, it is set aside as a completed candidate and no longer expanded; search continues with the remaining beams until B complete hypotheses exist or the length cap is hit. The final answer is the best-scoring completed hypothesis.

A worked toy example: setting up the beams

Let the vocabulary be just {A, B, <eos>} and the beam width B = 2. All logs are natural logs. At the first step the model gives P(A) = 0.5, P(B) = 0.3, P(<eos>) = 0.2. In log space:

step 1 candidates (score = log P):
  A       →  log 0.5 = -0.693
  B       →  log 0.3 = -1.204
  <eos>   →  log 0.2 = -1.609

prune to top B=2  ⇒  beam = { A (-0.693),  B (-1.204) }

We keep A and B and drop the immediate <eos>. Notice the bookkeeping: each surviving beam carries its running score, which is just the sum of the log-probs of the tokens chosen so far. So far the search looks like greedy would — but greedy would have discarded B entirely, keeping only A. Beam search holds both open, and the next step is where that hedge starts to matter. We now expand each of the two beams over the full vocabulary, producing B × V = 2 × 3 = 6 candidates before pruning.

Working the example: expand, then prune

Suppose the model’s conditional probabilities are: after A, P(A|A)=0.2, P(B|A)=0.7, P(<eos>|A)=0.1; after B, P(A|B)=0.6, P(B|B)=0.3, P(<eos>|B)=0.1. We add the new log-prob to each beam’s running score:

from A (-0.693):          from B (-1.204):
  A A  -0.693 + log0.2 = -2.302    B A  -1.204 + log0.6 = -1.715
  A B  -0.693 + log0.7 = -1.050    B B  -1.204 + log0.3 = -2.408
  A <eos> -0.693 + log0.1 = -2.996  B <eos> -1.204 + log0.1 = -3.507

all 6 sorted:  A B (-1.050) > B A (-1.715) > A A (-2.302)
               > B B (-2.408) > A <eos> (-2.996) > B <eos> (-3.507)

prune to top B=2  ⇒  beam = { A B (-1.050),  B A (-1.715) }

Here is the payoff. AB now leads with -1.050, and it descends from A, the token greedy also liked — but BA at -1.715 survives too, and it descends from B, which greedy had already thrown away. The beam has quietly rescued a path greedy could never reach. One more expansion (say AB then <eos> with P=0.6, giving -1.050 + log0.6 = -1.561) completes the winning hypothesis AB<eos> with final score -1.561.

Reading the beam tree

It helps to picture the search as a tree that is repeatedly widened and then trimmed back to B trunks. The diagram traces the toy run above: solid nodes are the kept beams, faded nodes are candidates that were expanded but pruned away.

The width B is exactly how many trunks the tree is allowed to keep. Larger B explores more of the tree and can only ever find an equal-or-better sequence — but it costs proportionally more compute, and, as we will see, past a point it stops helping and can even hurt.

The short-sequence bias in raw scores

There is a bug hiding in the plain score, and it is not a coding bug — it is baked into the math. Every log-probability term is ≤ 0, because probabilities are ≤ 1. So every token you append can only make the running score smaller (more negative), never larger. A sequence’s score is a running sum of non-positive numbers:

score(y_1..y_T) = Σ log P(y_t | ·)    where every term ≤ 0
⇒ score is monotonically non-increasing in T

The consequence: among candidate completions, shorter sequences almost always outscore longer ones, purely because they have fewer negative terms to add up — not because they are better answers. Left unchecked, beam search develops a strong bias toward terse, truncated outputs; in machine translation it was observed to emit sentences that stop too early, and larger beam widths made it worse, because a wider beam is more likely to surface a high-scoring short <eos> and declare victory. The raw argmax is quietly optimizing for brevity. This is the single most important correction beam search needs, and it is what length normalization addresses.

Length normalization: the length-penalty formula

The fix is to score by something closer to per-token quality rather than total accumulated log-prob. The simplest version divides the score by the length, giving the mean log-probability:

normalized(Y) = ( Σ log P(y_t | ·) ) / |Y|

The widely used generalization (from Google’s GNMT) introduces a tunable penalty with strength α:

score(Y) = logP(Y) / lp(Y)     with   lp(Y) = ( (5 + |Y|)^α ) / (5 + 1)^α

α = 0   → lp = 1        (no penalty; raw log-prob, short-biased)
α = 1   → ≈ divide by length (strong; mean log-prob)
α ∈ [0.6, 0.7]         (typical sweet spot in practice)

Watch it flip a decision. Take a short completion scoring -1.56 at length 2 and a longer, more complete one scoring -3.20 at length 5. Raw scores prefer the short one (-1.56 > -3.20). Divide by length: -1.56/2 = -0.78 versus -3.20/5 = -0.64 — now the longer sequence wins, because its average token is more probable. α is the dial between these regimes; you tune it on a validation set to stop the model both from rambling and from cutting itself off.

Coverage and repetition penalties

Length is not the only pathology. Two more penalties are commonly bolted onto the beam score. The coverage penalty comes from translation and summarization, where the model can ignore parts of the source or attend to the same source tokens over and over. It adds a term rewarding hypotheses whose attention has covered the input, penalizing under-translation:

cp = β · Σ_i log( min( Σ_t a_{t,i} , 1.0 ) )

where a_{t,i} is the attention paid to source position i at output step t; the min(·, 1) caps the reward so over-attending is not encouraged. The repetition penalty attacks the tendency of high-probability search to loop — ‘the the the’ or a sentence repeated verbatim. It discounts the logit of any token already generated before the softmax, so previously seen tokens become less likely; a related hard-constraint variant simply bans repeating any n-gram (no_repeat_ngram_size) by setting its probability to zero. Treat both as pointers rather than core machinery: they are corrective terms added to the objective, tuned per task, and they exist precisely because searching hard for the highest-probability string surfaces the model’s degenerate, repetitive modes that sampling would have stepped around.

Complexity: O(B · V · steps)

The cost of beam search factors cleanly. At each of T steps you run the model for each of the B live beams (in practice batched into one forward pass of batch size B), and each produces a distribution over all V tokens. Scoring the frontier is therefore B × V additions per step, and selecting the top B is a partial sort of those BV values:

scoring:   O(B · V)         per step
top-B:     O(B · V · log B)  per step (or O(BV) with a heap)
total:     O(B · V · T)       over T steps

vs greedy: O(V · T)         (the B = 1 special case)

So beam search is a constant factor B times more expensive than greedy — roughly B times the FLOPs, memory, and (if not batched) wall-clock. The V factor is unavoidable because the model emits a full vocabulary distribution every step; the interesting knob is B. On a CPU-bound small language model this matters directly: a beam of 4 means four forward passes’ worth of compute and four times the KV-cache memory per generated token. That is often the reason interactive, latency-sensitive SLM deployments quietly default to greedy or sampling with B = 1 and reserve beam search for offline, quality-critical batch jobs where the extra B× is affordable.

When beam search beats sampling

Beam search and sampling answer different questions. Sampling (with temperature, top-k, or nucleus/top-p) draws a token randomly in proportion to its probability, so it produces diverse, often more natural-sounding text and is the right default for open-ended generation — stories, chat, brainstorming — where there is no single correct answer and repeated runs should differ. Beam search does the opposite: it is deterministic and it tries to find the single most probable sequence.

That determinism-plus-optimality is exactly what you want for tasks with a correct or near-correct target, where fidelity beats creativity: machine translation, grammatical error correction, speech recognition transcription, constrained code or data-to-text generation, and short factual answers. In these, a slightly higher sequence probability really does correlate with a better output, and you do not want the randomness of sampling to occasionally derail a translation for the sake of variety. The rule of thumb: if the task has a ‘right answer’ you are trying to recover, prefer beam search; if it wants a plausible, varied continuation, prefer sampling. Many production systems even combine them — diverse beam search, or sampling within beam groups — to buy some of both.

When beam search degenerates

More search is not always better text, and this is the counterintuitive result practitioners keep rediscovering. Pushing the beam width up does find higher-probability sequences by construction — but past a modest width (often 4–10) the outputs frequently get worse: blander, shorter, more generic, and prone to repetition. This is the ‘curse of beam search’ or beam-search degradation.

The reason is a mismatch between the objective and reality: the highest-probability sequence under a trained model is often not the highest-quality one. Models place surprising probability mass on safe, empty, repetitive strings (the famous failure where the globally most likely translation is the empty string). A wider beam is simply a better optimizer of a flawed objective, so it finds those degenerate high-probability modes that a narrower beam or a sampler would have walked past. This is why open-ended generation moved toward nucleus sampling, and why beam search in practice is always paired with the corrections above — length normalization, coverage and repetition penalties, n-gram blocking. Used with a small width and those guards on constrained tasks, it shines; used with a huge width on open-ended text and no penalties, it collapses into confident, probable, lifeless output.

Practical notes for CPU and small models

On a CPU-hosted small language model, beam search’s B× cost is felt in three places at once, and it is worth naming them. Compute: B forward passes per token; on a CPU with no spare parallelism this can be close to a linear B× slowdown in latency. Memory: each beam carries its own KV cache, so peak memory scales with B × context length — the same quadratic attention and linear-in-context KV growth discussed elsewhere in this series, now multiplied by the beam count. Batching: the saving grace is that the B beams of one request batch naturally into a single forward pass, so on hardware with headroom the wall-clock cost is far below B×; on a saturated CPU it is not.

The pragmatic pattern is to reserve beam search for the requests that justify it. Use greedy or light sampling for interactive turns, and switch to a small beam (B = 3 or 4) with length normalization for offline, quality-sensitive jobs — translation passes, structured extraction, re-ranking candidates. And always measure: a beam of 2 often captures most of the quality gain over greedy at half the cost of a beam of 4.

Common pitfalls

A handful of mistakes recur when implementing or tuning beam search. Forgetting length normalization is the classic one: raw scores silently bias toward short outputs, and the symptom — the model stopping too early — is easy to misdiagnose as a training problem when it is really a decoding one. Comparing scores across different lengths without normalizing (for example, ranking a completed short hypothesis against a longer in-progress one on raw log-prob) makes the same error inside the search loop. Mishandling <eos>: a finished beam must be moved to the completed set and stop being expanded, and you keep searching until you have B completed hypotheses, not until the first one finishes — otherwise you throw away better completions still in progress.

Two more: assuming bigger B is strictly better — it improves the probability of the found sequence but can degrade quality, so sweep it rather than maxing it; and using beam search for open-ended generation, where its determinism and mode-seeking produce repetitive, generic text that a sampler would avoid. Beam search is a precision instrument for finding likely sequences — powerful when the task rewards likelihood, actively wrong when it rewards diversity.