Beam Search — Module 11.1

How This Lesson Fits the Module & Volume

Sampling, top-k, and top-p commit to one token at a time randomly. Beam search instead keeps the B highest-scoring partial sequences (beams) and expands them, aiming for high joint sequence probability.

It dominated classic NMT and remains useful for short, constrained outputs. Open-ended chat usually prefers truncated sampling; this lecture shows when search wins and when it fails (dull / repetitive text).

Learning Objectives

By the end of this lesson, students should be able to:

Define beam search with beam width B.
Trace one expand–score–prune step on a tiny vocabulary.
Contrast beam search with greedy and with sampling.
Explain length normalization / why raw log-prob favors short sequences.
Implement a minimal beam step in PyTorch.
State when beam search is a poor fit for creative LMs.

Definition

Beam search is a heuristic search over token sequences. At each step, every active hypothesis is expanded by candidate next tokens; all expansions are scored (usually sum of log-probabilities), and only the top B hypotheses are kept. Beam width B = 1 recovers greedy decoding.

Expand, Score, Prune

Beams

B partial strings

Expand

Try next tokens

Score

Sum log p

Prune

Keep top B

Method	Stores	Stochastic?	Typical win
Greedy	1 path	No	Speed, simplicity
Beam (B>1)	B paths	No (standard)	MT, short structured text
Sampling	1 path	Yes	Chat, creative writing

Beam vs Sampling

Beam strengths

Can recover from early local mistakes.
Optimizes sequence score approximately.
Deterministic and debuggable.

Beam weaknesses

Cost scales with B.
Often bland / repetitive for open text.
Needs length penalty tricks.

Sampling strengths

Natural diversity.
Cheap (one hypothesis).
Matches chat UX expectations.

Code: One Beam Expansion Step

import torch
import torch.nn.functional as F

def beam_step(logits, beam_logprobs, beam_width=3):
    """
    logits: (B_beams, V) next-token logits for each beam
    beam_logprobs: (B_beams,) cumulative logprob so far
    returns new_logprobs (beam_width,), token_ids (beam_width,), parent_beams (beam_width,)
    """
    log_probs = F.log_softmax(logits, dim=-1)          # (B, V)
    total = beam_logprobs.unsqueeze(-1) + log_probs    # (B, V)
    flat = total.view(-1)
    vals, idx = torch.topk(flat, beam_width)
    vocab = logits.size(-1)
    parent = idx // vocab
    token = idx % vocab
    return vals, token, parent

# Toy: 2 beams, V=5
logits = torch.tensor([[2.0, 1.0, 0.5, 0.0, -1.0],
                       [0.2, 2.5, 1.0, 0.1, -0.5]])
beam_lp = torch.tensor([-0.4, -0.6])
scores, toks, parents = beam_step(logits, beam_lp, beam_width=3)
print(scores, toks, parents)
# Keep those 3 hypotheses, append toks, continue until EOS / max len

Strengths and Tradeoffs

Strengths

Better than greedy for many structured tasks.
Interpretable scoreboard of hypotheses.
Works well with constrained decoding.

Tradeoffs

Not globally optimal (heuristic).
Open-ended generation often prefers sampling.
Memory/compute grow with beam width.

Common Misconception

“Larger beam always means better text.” Bigger beams can increase BLEU in MT yet make open-ended LM text more generic or repetitive. Quality is task-dependent; bigger B is not free quality.

Related module pages: Sampling, Top-P, Top-K, Logits, Inference.

Knowledge Check

Short Answer: What does beam width B mean? Answer: How many partial hypotheses are kept each step.
True/False: Beam width 1 equals greedy decoding. Answer: True.
Multiple Choice: Hypotheses are usually scored with: (a) sum of log-probs, (b) random noise only, (c) image MSE. Answer: (a).
Short Answer: Why might raw log-prob prefer short outputs? Answer: Each extra token multiplies probability (≤1), so longer sequences accumulate lower joint prob.
True/False: Standard beam search is stochastic like top-p. Answer: False.
Multiple Choice: Open-ended chat typically prefers: (a) large beam, (b) truncated sampling, (c) no decoding. Answer: (b).
Short Answer: Name one classic domain for beam search. Answer: Machine translation (or ASR).
True/False: Beam search guarantees the globally highest-probability string. Answer: False—it is a heuristic.
Multiple Choice: Expanding beams costs roughly: (a) linear in B, (b) free, (c) only tokenizer time. Answer: (a).
Short Answer: What efficiency lecture caches past keys/values during generation? Answer: KV Cache.

Key Takeaways

Beam search keeps B best partial sequences and expands them.
B = 1 is greedy; larger B explores more paths at higher cost.
Great for short/structured tasks; often bland for open chat.
Scores use log-probs; watch length bias.
Next: KV Cache.

Trainer’s Guide

Hands-on idea: On paper, run B=2 for two steps on a 3-token vocab with made-up probs; show a path that greedy misses but beam finds.

Discussion prompt: Would you use beam search to generate a marketing slogan? Why or why not?

Recap: Beam search approximates high-scoring sequences by pruning to B hypotheses each step. Continue with KV Cache.