Beam Search › Beam Search: Approximate Best-First Decoding

Beam search is a heuristic that maintains a fixed number of candidate hypotheses at each step. This number is called the beam width, written $k$. Instead of greedily committing to the single best token, beam search tracks the $k$ best partial sequences at every time step, extending each one by one step, and then pruning back to the $k$ best again.

The name "beam" is a spatial metaphor. Imagine the search space as a tree where the root is the start token and each level corresponds to one decoding step. The branches fan out to all $|V|$ possible next tokens. Greedy search follows a single path down this tree. Exhaustive search explores every path. Beam search illuminates $k$ paths at once, like a beam of light of fixed width sweeping across the tree. When a path scores poorly, it falls out of the beam. When a path scores well, it continues.

Beam search operates on log-probabilities rather than raw probabilities. Multiplying many probabilities together quickly underflows to zero in floating-point arithmetic. Log-probabilities transform multiplication into addition, solving the underflow problem while preserving the ordering (since logarithm is monotonically increasing).

The joint probability of a sequence $\mathbf{y} = (y_1, y_2, \ldots, y_T)$ is:

Taking the log:

where:

• $P(\mathbf{y})$: the joint probability of the complete sequence
• $P(y_t \mid y_1, \ldots, y_{t-1}, \mathbf{c})$: the conditional probability of token $y_t$ given previous tokens and context
• $\log P(\mathbf{y})$: the log-probability of the sequence, computed as a sum of per-step log-probabilities

Since log-probabilities are negative (the log of a probability between 0 and 1 is always negative), higher (less negative) log-probabilities indicate better sequences.

The beam search procedure is:

1. Initialize the beam with one hypothesis: the start token with log-probability 0.
2. For each step $t = 1, 2, \ldots, T$: a. For each hypothesis in the beam (up to $k$ hypotheses), run the decoder forward to get a distribution over the next token. b. Generate all $k \times |V|$ candidate extensions (each hypothesis extended by each token). c. Score each candidate: add the new token's log-probability to the hypothesis's current score. d. Prune: keep only the top $k$ candidates by score. These become the new beam. e. If a hypothesis produces the end token (EOS), move it to the "completed" list and remove it from the beam.
3. When the beam is empty or a maximum length is reached, return the completed hypothesis with the highest score (or the best hypothesis remaining in the beam if none completed).

Consider a toy vocabulary of five tokens: {A, B, C, D, EOS} and a beam width of 2.

At step 0, the beam contains one hypothesis: [<START>] with score 0.

The decoder outputs probabilities for the first token:

• A: 0.45
• B: 0.35
• C: 0.15
• D: 0.04
• EOS: 0.01

In log-space: log(0.45) = -0.80, log(0.35) = -1.05, log(0.15) = -1.90, etc.

With beam width 2, we keep A (score -0.80) and B (score -1.05). The beam now contains two hypotheses.

At step 1, each hypothesis is extended by all 5 tokens. We have 10 candidates. Suppose the decoder, after seeing A, assigns probability 0.60 to EOS (score: $-0.80 + \log 0.60 = -0.80 + (-0.51) = -1.31$) and 0.25 to B (score: $-0.80 + (-1.39) = -2.19$). After seeing B, the decoder assigns probability 0.50 to A (score: $-1.05 + (-0.69) = -1.74$) and 0.30 to C (score: $-1.05 + (-1.20) = -2.25$).

The best 2 candidates from all 10 are A-EOS (score -1.31) and B-A (score -1.74). A-EOS is a completed sequence (it produced EOS), so it moves to the completed list. The beam now contains only B-A.

The process continues until either all hypotheses complete or the maximum length is reached. The final answer is the completed sequence with the best score, which here is A-EOS.

Notice what happened: greedy decoding would have generated A at step 1, gotten EOS at step 2, and returned A. But if B-A-EOS had a higher joint probability than A-EOS, beam search with width 2 would have found it, while greedy search would not.

A key property to notice is that the beam at any given step may contain hypotheses of different lengths, all anchored at the same decoding step. The pruning compares raw cumulative log-probabilities across these hypotheses. This is why length normalization matters: a hypothesis that has accumulated more log-probability simply by being shorter is not necessarily better. We will return to this in the length normalization section.