L8.5 · LESSON 5 OF 6

Beam search

Compute temperature and candidate filtering, and distinguish sampling from greedy and beam decoding.

THE IDEA

Beam search

Beam search keeps the best width partial sequences at every step instead of just one. It extends all of them, scores every extension (typically by the sum of log-probabilities) and keeps the best width again. With width 1 it is greedy. A wider beam can find sequences greedy misses, such as <s> b c </s> in the tiny model.

beam width2

after step 0

<s>1

after step 1

<s> a0.6

<s> b0.4

after step 2

<s> b c0.36

<s> a c0.3

after step 3

<s> b c </s>0.36

<s> a c </s>0.3

width 1 is greedy and ends with <s> a c </s> (0.3) width ≥ 2 keeps b alive long enough to find <s> b c </s> (0.36)

Beam search keeps several partial sequences and extends them all. It optimises a scoring rule such as total log-probability, so it tends to prefer shorter sequences unless the score is length-normalised, and a larger beam does not guarantee text people prefer.

Change the width and read the beams after each step.

Predict first. Can beam scoring favour repetitive or short sequences under some objectives?

Chapter context

Decoding chooses tokens from the model's output distribution. It changes generated behavior without retraining the model. Sampling settings influence diversity and concentration, but they do not independently guarantee correctness or creativity.

TRY IT YOURSELF

Part 1. Do it hands-on

Find the best whole sequence and trace a width-2 beam by hand, then implement beam search with an optional length penalty.

beam width2

after step 0

<s>1

after step 1

<s> a0.6

<s> b0.4

after step 2

<s> b c0.36

<s> a c0.3

after step 3

<s> b c </s>0.36

<s> a c </s>0.3

width 1 is greedy and ends with <s> a c </s> (0.3) width ≥ 2 keeps b alive long enough to find <s> b c </s> (0.36)

most probable full sequence (middle tokens) probability of greedy's sequence best extension for width 2 after step 2

To do: Using the same tiny model, which complete sequence has the highest probability? Type the middle tokens (like "b c")
To do: What probability does greedy's sequence get? (The best one has 0.36)
To do: Beam search with width 2 keeps the two best partial sequences after step 1: a (0.6) and b (0.4). After step 2, which of the four candidates is best? Type e.g. "b c"

Part 2. Do it in code

beam_search(lm, width, start="<s>", end="</s>", max_len=10, length_alpha=0.0): beams are (sequence, log_probability). At each step extend every unfinished beam with every next token, keep the best `width` by score, where score = log_prob / (len(sequence) ** length_alpha) (alpha 0 means no length normalisation). Finished beams (ending with `end`) are kept as they are. Stop when all beams are finished or max_len tokens are reached. Return the beams sorted by score, best first, as a list of (sequence, probability) with probability = exp(log_prob).
width must be at least 1 (ValueError).

import math

def beam_search(lm, width, start="<s>", end="</s>", max_len=10, length_alpha=0.0):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Can beam scoring favor repetitive or short sequences under some objectives?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

p_i(T)=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)},\quad T>0.

Top-p retains a probability-sorted prefix with cumulative mass at least p, then renormalizes the retained entries. Specify tie handling and any minimum-candidate rule in implementation.

Worked example

For logits [2,1,0] at temperature one, probabilities are approximately [0.665241,0.244728,0.090031]. With top-p=0.8, keep the first two because the first alone is below the threshold. Their renormalized probabilities are approximately [0.731059,0.268941].

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import exp,isclose
weights=[exp(v-2) for v in [2,1,0]]
probs=[w/sum(weights) for w in weights]
assert probs[0] < .8 < sum(probs[:2])
retained=[p/sum(probs[:2]) for p in probs[:2]]
assert isclose(probs[0],.6652409557748218)
assert isclose(retained[0],.7310585786300049)
assert isclose(sum(retained),1)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Top-p=.8 does not mean keeping 80% of vocabulary items. A high temperature does not create new factual knowledge. Filter-and-renormalize order must be explicit.