Beam search

Beam search is a way for an AI model to choose its words by keeping several drafts open at once. Instead of committing to the most likely next word every time, it follows a handful of promising partial sentences and picks the best complete one at the end. Translation and speech-to-text tools use it, because there you want the single most likely answer.

How does it work?

A language model writes one token at a time (see autoregressive generation). Greedy decoding always takes the most likely next token, and can miss a better sentence that starts with a less likely word. Checking every possible sentence would find the best one, but that is far too slow. Beam search, first proposed by Lowerre in 1976, sits in between, as Jurafsky and Martin describe in their textbook:

1. Keep the k best partial sentences, the hypotheses. The number k is the beam width, after the beam of a flashlight that can be set wider or narrower.
2. Extend every hypothesis with every possible next token, and score each result by adding up the log-probabilities of its tokens.
3. Keep only the k best results and repeat.
4. When a hypothesis ends with an end-of-sequence token, set it aside and shrink the beam by one, until k complete hypotheses remain.
5. Return the best-scoring one, or all k with their scores.

Without randomness, the same input always gives the same output. Beam search is still a heuristic: it is not guaranteed to find the most likely output.

Which settings matter?

The beam width. Jurafsky and Martin write that machine translation generally uses 5 to 10. The original transformer used 4 for translation. Google's translation system GNMT kept 8 to 12 hypotheses and found that 4 or 2 cost only a little quality. Each hypothesis needs its own pass through the model, so a wider beam costs more computation.

Length normalisation. Log-probabilities are negative and add up, so longer sentences score lower and plain beam search favours short output. The fix is to divide the score by the length, or by a power of it: GNMT found values between 0.6 and 0.7 usually best, and the transformer used 0.6.

Variants. Diverse beam search (Vijayakumar et al., 2016) counters the problem that beams "differ only slightly from each other". Grid beam search (Hokamp and Liu, 2017) forces given words or phrases into the output, such as fixed terminology in a translation.

Where does it go wrong?

Beam search suits tasks where the length of the answer is fairly predictable, such as translation and summarisation, according to the Hugging Face guide. In open-ended text such as dialogue or stories it tends to repeat itself. Holtzman et al. (2019) found that text chosen for high probability comes out "bland and strangely repetitive". For open-ended text the guide prefers sampling, with settings such as temperature and top-p, while noting that sampling can repeat as well and that no single method fits every task.

Speech recognition shows the other side. OpenAI's Whisper (2022) uses beam search with 5 beams for long recordings, precisely to reduce the repetition loops that occur more often with greedy decoding.

Example

The Hugging Face guide uses a toy example. After "The", the most likely next word is "nice" (probability 0.5), and after "The nice" it is "woman" (0.4), so greedy decoding ends at "The nice woman" with a probability of 0.5 × 0.4 = 0.2. Beam search with 2 beams also keeps the runner-up "The dog". One step later it finds that "The dog has" scores 0.36, higher than greedy's 0.2, and keeps that sentence.