Greedy and beam search: the best word or the best sentence?

Tokenization & decoding · Part 6 of 8

A model gives you probabilities for the next word. Turning those into text means choosing. The obvious choice, always take the most likely word, has a flaw that a tiny example shows clearly.

Greedy: always take the best next word

Greedy decoding picks the highest-probability token at every step. It is fast and deterministic. Its flaw: a choice that looks best now can lead somewhere worse, and greedy never goes back.

A two-word example

Suppose a model starts a sentence with “I” 55% of the time and “We” 45% of the time. After “I” it is unsure (“am” 40%, “think” 30%, “guess” 30%). After “We” it is nearly certain (“are” 90%). Greedy takes the more likely first word, “I”, then “am”: total probability 0.55 × 0.40 = 0.22. But “We are” has probability 0.45 × 0.90 = 0.405, almost twice as likely.

Beam search: keep several candidates

Beam search with beam size k keeps the best k partial sentences at every step. Each survivor is extended by every possible next token, all the candidates are scored by their total probability (in practice the sum of log-probabilities), and only the best k survive. Greedy decoding is just the case k = 1. Nothing is sampled, so the result is deterministic. (This k has nothing to do with the k of top-k sampling in the next post.)

first = {"I": .55, "We": .45}
then = {"I": {"am": .40, "think": .30, "guess": .30}, "We": {"are": .90, "were": .05, "have": .05}}

w1 = max(first, key=first.get); w2 = max(then[w1], key=then[w1].get)          # greedy
print("greedy:", w1, w2, round(first[w1] * then[w1][w2], 3))                 # greedy: I am 0.22

beams = sorted(first.items(), key=lambda kv: -kv[1])[:2]                      # keep the best 2
cands = [((a, b), p * q) for a, p in beams for b, q in then[a].items()]       # extend each by every word
best = sorted(cands, key=lambda kv: -kv[1])[:2]
print("beam:", [(" ".join(s), round(p, 3)) for s, p in best])                 # beam: [('We are', 0.405), ('I am', 0.22)]

The catch

Beam search finds likelier sequences, but the likeliest sentence is often the safest and most generic one. For open-ended writing that can sound flat and repetitive. That is why chat models usually sample instead, which is the subject of the next post.

Takeaways

• Greedy picks the most likely token each step and never looks back.
• A locally best choice can lead to a worse whole sentence.
• Beam search keeps k candidate sequences and scores them by total probability.
• Beam search is deterministic and favours safe, generic text.

Next in the series (Part 7): Temperature, top-k and top-p: how models pick the next word