Skip to content
Back to Generative AIAll lessons

Generative AI

Beam Search

O(steps · beamWidth · vocabularySize) time. At every step, every surviving beam is extended by every word in the vocabulary, before being trimmed back down to the beam width.

The idea, in plain English

Imagine you choose a hiking route by only ever taking whichever single trail looks best at each fork. This is called 'greedy'. Now imagine instead sending out a small search party that keeps a handful of the most promising routes alive at once, and only commits to one at the very end. Beam search is a language model's version of that search party. Instead of locking in the single best next word at every step, which can back itself into a dead end, it keeps the top few candidate sequences, called 'beams', alive at each step. It only picks the overall best one once generation is done.

How it works

  1. 1Start with one empty sequence. At each step, extend every sequence you are currently tracking with every possible next word. Add each new word's score to that sequence's running total.
  2. 2Sort all these extended candidates by their total score. If there is a tie, break it alphabetically, so the result stays predictable. Keep only the top 'beam width' number of them, and discard the rest.
  3. 3Repeat for a fixed number of steps. At the very end, the highest-scoring sequence among the surviving beams is the answer.

When you'd use it

Use this to generate a genuinely good overall sequence, like a full translated sentence, rather than one that only ever looks good one word at a time. It is a step up from always keeping just one running sequence, and it shares an idea with Top-k Selection: keep more than one option alive.

Common beginner mistakes

  • Don't assume beam search always finds the mathematically best possible sequence. It does not. It only ever keeps a small number of candidates alive, so a sequence that looked mediocre early on can get discarded before it has a chance to prove itself.
  • Don't confuse beam search with greedy decoding, which is a beam width of 1. Greedy commits to a single best-looking choice at every step, and can end up stuck with a worse overall sequence. This example demonstrates exactly that.

Try it — edit and run

Click the code to edit · press ⌘/Ctrl+↵ to run

Editable code. Tab and Shift+Tab indent. Press Escape, then Tab, to move focus out of the editor.

Expected output — hit Run to try it
Beam search with beam width 2:
Step 1 — top 2 beams kept:
  [cat] score=5
  [dog] score=4
Step 2 — top 2 beams kept:
  [dog fish] score=13
  [cat fish] score=8
Step 3 — top 2 beams kept:
  [dog fish cat] score=18
  [dog fish dog] score=16

Beam search result: dog fish cat (score 18)
Greedy (beam width 1) result: cat fish cat (score 13)

Beam search found a better overall sequence than greedy.

Not sure this is the right topic? See the learning paths → or where this leads →