Beam search decoding

The problem with greedy decoding

When generating sequences (translation, captioning, summarization), we want:

\mathbf{y}^* = \arg\max_{\mathbf{y}} P(\mathbf{y} \mid \mathbf{x})

where $\mathbf{y} = (y_1, y_2, \ldots, y_T)$ and $P(\mathbf{y} \mid \mathbf{x}) = \prod_{t=1}^T P(y_t \mid y_{<t}, \mathbf{x})$ .

Greedy decoding picks $y_t = \arg\max P(y_t \mid y_{<t}, \mathbf{x})$ at each step. Why this fails: the product of locally-optimal choices is NOT globally optimal.

Example: Translating "Je suis étudiant" to English.

Greedy might pick: "I" (0.6) → "am" (0.7) → "student" (0.3) → $P = 0.126$
Better path: "I" (0.6) → "am" (0.65) → "a" (0.9) → "student" (0.8) → $P = 0.281$

The greedy path missed "a" because at step 2, "am" alone looked good—but including "a" later creates a much more natural phrase.

Derivation: How beam search scores sequences

Step 1: The objective function

We want to maximize the log-probability (more numerically stable than raw probability):

\text{score}(\mathbf{y}) = \log P(\mathbf{y} \mid \mathbf{x}) = \sum_{t=1}^T \log P(y_t \mid y_{<t}, \mathbf{x})

Why log? Products of small probabilities underflow; sums of logs stay in range. Also: $\arg\max P = \arg\max \log P$ since $\log$ is monotonic.

Step 2: Length normalization

Raw log-probability penalizes longer sequences (more terms in the sum, each negative). A 10-word sentence will almost always score lower than a 3-word sentence, even if it's better.

Fix: Divide by length (with smoothing):

\text{score}_{\text{norm}}(\mathbf{y}) = \frac{1}{T^\alpha} \sum_{t=1}^T \log P(y_t \mid y_{<t}, \mathbf{x})

Here $T$ is the number of generated tokens (i.e., we do not count the <START> token, since it carries no predicted probability). In the worked example below, a completed sequence like [<START>, Bonjour, <END>] has $T=2$ generated tokens (Bonjour and <END>).

$\alpha = 1$ : full normalization (average log-prob per token)
$\alpha = 0$ : no normalization (raw sum)
$\alpha \in (0.6, 0.7)$ : common in practice (slight preference for longer, less penalty)

Why $\alpha < 1$ ? Very short sequences can have deceptively high average probabilities (e.g., "I." has high avg prob but is incomplete). Partial normalization balances.

Step 3: The algorithm

Formula Beam search algorithm

Input: Encoder output $\mathbf{h}$ , beam width $k$ , length penalty $\alpha$
Output: Top $k$ decoded sequences

beams = [([<START>], 0.0)]  # (sequence, cumulative_log_prob)
for t = 1 to max_length:
    candidates = []
    for (seq, score) in beams:
        if seq[-1] == <END>:
            candidates.append((seq, score))
            continue
        for token in vocabulary:
            new_seq = seq + [token]
            new_score = score + log P(token | seq, h)
            candidates.append((new_seq, new_score))
    # Keep top k by normalized score.
    # num_generated = number of tokens AFTER <START> (do not count <START>)
    beams = top_k(candidates, k, key=λ(s, sc): sc / num_generated(s)^α)
    
    if all sequences end with <END>: break
return beams

Note on length $T$ : num_generated(s) = len(s) - 1 because the leading <START> token is a fixed marker, not a predicted token. This is exactly the $T$ used in the score formula, so the pseudocode and the worked example are consistent.

Key insight: At each step, we generate $k \times V$ candidates but only keep $k$ . This prunes the exponential search space ( $V^T$ total sequences) to $O(kTV)$ operations.

Worked example: Translating "Hello" to French

Setup: Vocabulary = {, Bonjour, Salut, }, beam width $k=2$ , $\alpha=1$ .

Step 0: Initialize

Beams: [([<START>], 0.0)]

Step 1: Expand from <START>

Candidates:
- [<START>, Bonjour]: $\log 0.6 = -0.51$
- [<START>, Salut]: $\log 0.35 = -1.05$
- [<START>, <END>]: $\log 0.05 = -3.0$
Keep top 2: Bonjour (-0.51), Salut (-1.05)

Why this step? We try all possible first words, score them by model's predicted probabilities, keep the 2 best.

Step 2: Expand both beams

From Bonjour:
- [<START>, Bonjour, <END>]: $-0.51 + \log 0.9 = -0.51 - 0.11 = -0.62$
- [<START>, Bonjour, Salut]: $-0.51 + \log 0.08 = -0.51 - 2.53 = -3.04$
From Salut:
- [<START>, Salut, <END>]: $-1.05 + \log 0.85 = -1.05 - 0.16 = -1.21$
- [<START>, Salut, Bonjour]: $-1.05 + \log 0.1 = -1.05 - 2.30 = -3.35$

Normalization (divide by $T =$ number of generated tokens $=$ len(seq) - 1; here each completed sequence has 2 generated tokens, so $T=2$ ):

Bonjour <END>: $-0.62 / 2 = -0.31$
Salut <END>: $-1.21 / 2 = -0.605$

Keep top 2: [<START>, Bonjour, <END>] (normalized -0.31), [<START>, Salut, <END>] (normalized -0.605).

Why this step? Both sequences ended, so we compare their length-normalized scores. "Bonjour" wins because the model assigned it higher probability.

Final output: "Bonjour" (score -0.31)

Common hyperparameters

Parameter	Typical range	Effect
Beam width $k$	4-10	Larger = better quality, slower. Diminishing returns after ~10.
Length penalty $\alpha$	0.6-1.0	Higher = prefer longer sentences. 0 = no penalty.
Min length	Task-dependent	Force decoder to generate at least $n$ tokens (avoid premature `<END>`).
N-gram blocking	2-4	Prevent repeating the same $n$ -gram (useful in summarization).

Common mistake Common pitfall: Thinking larger beams always help

Wrong intuition: " $k=100$ must be better than $k=5$ because we explore more."

Why it feels right: More candidates = more thorough search = better result, right?

The fix:

Diminishing returns: After $k \approx 10$ , the top- $k$ sequences are very similar (differ in low-probability words). You're just spending compute to track near-duplicates.
Length bias amplifies: With large $k$ , you keep many mediocre long sequences that benefit from length normalization, crowding out better short sequences.
Empirical findings: Translation quality peaks around $k=4\text{-}10$ for most models (Britz et al., 2017).

Better approach: Tune $k$ on a validation set. If $k=5$ and $k=20$ give similar BLEU scores, use $k=5$ .

Common mistake Forgetting that beam search is still approximate

Wrong intuition: "Beam search finds the globally optimal sequence."

Why it feels right: We're keeping the top-scoring sequences at each step—surely we'll find the best one?

The fix:

Beam search prunes. If the true best sequence has a low-probability word early (that falls outside the top $k$ ), we'll never explore it.
Example: True best is "I have a cat" (overall prob 0.15), but "I have the" (0.12) and "I own a" (0.13) are in the beam at step 2. "I have a" scored 0.11 at step 2, so it was pruned. We'll never reach "I have a cat."
Implication: Beam search is a heuristic. It's much better than greedy (which is deterministic and bad), but not exhaustive. For critical applications, consider ensemble decoding or sampling methods.

Comparison with other decoding methods

Method	Search space	Quality	Diversity Speed
Greedy	1 path	Low	None
Beam search	$k$ paths	High	Low
Sampling (top- $p$ , temp)	Stochastic	Variable	High
Exhaustive	All $V^T$	Optimal	N/A

When to use beam search:

Machine translation, summarization, captioning (want the single best output)
When fluency and correctness matter more than creativity

When NOT to use beam search:

Story generation, dialogue (want diversity; beam search produces generic responses)
Use sampling with temperature or nucleus sampling instead

Connections

Seq2seq models — Beam search is the standard decoding method for seq2seq
Attention mechanism — Compute $P(y_t \mid y_{<t}, \mathbf{x})$ using attention over encoder states
Temperature sampling — Alternative to beam search for creative generation
BLEU score — Metric used to evaluate beam search outputs in translation

Recall Explain like I'm 12

Imagine you're playing a word game where you build a sentence one word at a time, and each word has a "goodness score" (how likely it sounds).

Greedy is like only remembering your current sentence and always picking the best next word. But sometimes, a word that sounds "okay" now leads to a much better sentence later! Like if you're making "I want to eat pizza", and you pick "I want cookies" because "cookies" scored 90 vs "to" scored 80. But "to eat pizza" would've scored 90 × 95 = 8550 total, way better than "cookies" (90) ending there.

Beam search is like having a sticky note for your top 3 sentences at each step. You try adding words to all 3, get 3 × 26 = 78 new sentences, then keep only the best 3 again. So you're exploring multiple "what if I said THIS instead" paths, but not so many that your brain explodes.

The "beam width" is how many sticky notes you have. 1 note = greedy. 100 notes = your desk is a mess and you're not getting much benefit after 10 notes anyway.

#flashcards/ai-ml

What is the key difference between greedy decoding and beam search? :: Greedy picks the single highest-probability token at each step (myopic). Beam search maintains $k$ candidate sequences and explores multiple paths, improving global sequence quality.

Why do we use log-probabilities instead of raw probabilities in beam search?

To avoid numerical underflow (products of small numbers become zero in floating point) and because

\log

converts products to sums, which are easier to compute and compare.

What is length normalization in beam search and why is it needed?

Dividing the cumulative log-probability by

T^\alpha

, where

T

is the number of generated tokens (excluding <START>). Without it, longer sequences always score lower (more negative terms), so the model would favor short, incomplete outputs.

If beam width

k=1

, what does beam search reduce to?

Greedy decoding—only the single highest-probability sequence is kept at each step.

What is the computational complexity of beam search for a sequence of length $T$ with vocabulary size $V$ and beam width $k$ ? :: $O(kTV)$ —at each of $T$ steps, we expand $k$ sequences by considering $V$ tokens each.

Why doesn't beam search guarantee finding the globally optimal sequence?

It prunes the search tree—if the true best sequence has a low-probability token early that falls outside the top

k

, that path is discarded and never explored, even if later tokens would make it optimal.

What is a typical range for beam width in production translation systems?

k = 4

10

. Larger values have diminishing returns and can amplify length bias.

When should you use sampling methods instead of beam search?

For creative generation tasks (stories, dialogue, diverse responses) where you want variety and personality, not the single most probable (often generic) output.

What does the length penalty hyperparameter

\alpha

control?

How strongly we normalize by sequence length.

\alpha=0

: no normalization (favors short sequences).

\alpha=1

: full normalization (average per token).

\alpha \in (0.6, 0.7)

: common compromise.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Beam search ek aisa decoding technique hai jisse hum sequence generation mein better quality ke outputs nikalte hain. Socho translation kaam hai—agar hum sirf sabse zyada probable next word uthaye har step pe (greedy decoding), toh local level pe best choice ho sakta hai, lekin overall sentence awkward ya incomplete ban sakta hai. Jaise "I am student" grammatically galat hai compared to "I am a student."

Beam search ka idea sim

The problem with greedy decoding

Derivation: How beam search scores sequences

Step 1: The objective function

Step 2: Length normalization

Step 3: The algorithm

Worked example: Translating "Hello" to French

Common hyperparameters

Comparison with other decoding methods

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Sequence Models

Connections