Beam search decoding

Deep dive

How it works

Parts:

• running beams — up to k (prefix, summed log-prob)
• candidates — the k·V extensions, cut to the top 2k
• finished pool — sequences that emitted EOS, with normalised scores
• scorer — log P / lp(len), lp = ((5+len)/6)^α (GNMT; α = 0 is plain log P)
• per-beam model state — KV cache or recurrent state, reordered on reshuffle

Operations per step:

1. Expand. One batched model call gives next-token log-probs for every running beam; add them to the beams' sums.
2. Select. Keep the top 2k of the k·V candidates: at most k are EOS, so at least k are not.
3. Finish or continue. EOS candidates go to the pool with their normalised score; the best k others are the next beams.
4. Reorder state. Gather each new beam's cache from its parent.
5. Stop at max length, with an empty beam, or when no running beam can beat the best finished score.
6. Return the pool sorted by score: first = answer, first n = n-best list.

 beam b1 ─┬ tok x ┐                      ┌ EOS ───────► finished pool (score / lp)
 beam b2 ─┼ tok y ┼► k·V candidates ► top 2k ┤
 ...      └ ...   ┘                      └ best k non-EOS ► beams of step t+1
 caches:  cache[new] ← cache[parent]  (gather, or share blocks copy-on-write)

Stopping. Raw log-probs only fall, so with α ≥ 0 a running beam with sum r can never beat r / lp(max_len); once the best finished score exceeds that, stop (the code does). Hugging Face early_stopping=True stops at k finished, False uses a heuristic, "never" the bound. Kasai et al. add a patience factor to the common stop-at-k-finished rule; Whisper and CTranslate2 expose it.

Scoring. GNMT: s = log P / lp + cp, cp = β · Σ_i log min(Σ_j p_ij, 1) (attention p of target word j on source word i); α = β = 0.2 in their runs. On the WMT'14 En-Fr dev set both terms raised BLEU from 30.3 to 31.4 (maximum-likelihood model). Hugging Face: sum_logprobs / len^length_penalty, default 1.0, prompt excluded.

Diverse beam search (Vijayakumar et al.): G groups decode in turn; each subtracts λ times a penalty for tokens earlier groups chose at this step. λ of 0.2-0.8 worked.

CTC prefix beam search (speech, handwriting, OCR): many alignments collapse to one output, so the beam holds prefixes with two probabilities each, ending in blank and not (Hannun, Distill). An external LM enters as p(Y|X) · p(Y)^α · L(Y)^β, L = word-insertion bonus; in translation, shallow fusion (Gulcehre et al.). See CTC loss.

Batching. All beams of all requests are rows of one batch; cache rows are gathered by parent index (Whisper rearrange_kv_cache, Hugging Face reorder_cache). A paged KV cache shares prefix blocks, copying one only when a beam writes into it.

Code

Toy bigram model, beam search, greedy and exhaustive search. Asserts: k = 1 is greedy; k = 2 beats it on a trap; a very wide beam equals exhaustive search; α changes the winner.

const EOS: usize = 0; // token 0 ends a sequence
type Seq = Vec<usize>;
/// Toy bigram model: p[prev * v + next]; row v is the start-of-sequence row (prev = None).
struct Bigram { v: usize, p: Vec<f64> } // (v + 1) rows of v next-token probabilities
impl Bigram {
    fn logp(&self, prev: Option<usize>, next: usize) -> f64 { self.p[prev.unwrap_or(self.v) * self.v + next].ln() }
}

/// GNMT length penalty lp = ((5+len)/6)^alpha; score = log P / lp. alpha = 0 is plain log P.
fn norm(raw: f64, len: usize, alpha: f64) -> f64 { raw / ((5.0 + len as f64) / 6.0).powf(alpha) }

/// Beam search with a finished pool. Returns the n-best finished (seq incl. EOS, score), best first.
fn beam_search(m: &Bigram, k: usize, alpha: f64, max_len: usize) -> Vec<(Seq, f64)> {
    let mut beams: Vec<(Seq, f64)> = vec![(vec![], 0.0)]; // running: (tokens, raw log P)
    let mut done: Vec<(Seq, f64)> = vec![];
    for _ in 0..max_len {
        let mut cands: Vec<(usize, usize, f64)> = Vec::with_capacity(beams.len() * m.v); // (beam, token, raw)
        for (b, (s, raw)) in beams.iter().enumerate() {
            for t in 0..m.v { cands.push((b, t, raw + m.logp(s.last().copied(), t))); }
        }
        let by_raw = |a: &(usize, usize, f64), b: &(usize, usize, f64)| b.2.partial_cmp(&a.2).unwrap();
        if cands.len() > 2 * k { cands.select_nth_unstable_by(2 * k, by_raw); cands.truncate(2 * k); } // <= k EOS among the top 2k
        cands.sort_by(by_raw);
        let mut next: Vec<(Seq, f64)> = vec![];
        for (b, t, raw) in cands {
            let mut s = beams[b].0.clone();
            s.push(t);
            if t == EOS { let sc = norm(raw, s.len(), alpha); done.push((s, sc)); } // finished: leaves the beam
            else if next.len() < k { next.push((s, raw)); }
        }
        beams = next;
        // Safe stop: extending only lowers raw and lp <= lp(max_len), so norm(r, max_len) bounds any extension.
        let best_done = done.iter().map(|d| d.1).fold(f64::NEG_INFINITY, f64::max);
        if beams.is_empty() || (alpha >= 0.0 && best_done >= norm(beams[0].1, max_len, alpha)) { beams.clear(); break; }
    }
    for (s, raw) in beams { let sc = norm(raw, s.len(), alpha); done.push((s, sc)); } // still running at max_len: keep, no EOS
    done.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
    done
}

fn greedy(m: &Bigram, max_len: usize) -> Seq {
    let mut s: Seq = vec![];
    while s.len() < max_len {
        let prev = s.last().copied();
        let t = (0..m.v).max_by(|&a, &b| m.logp(prev, a).partial_cmp(&m.logp(prev, b)).unwrap()).unwrap();
        s.push(t);
        if t == EOS { break; }
    }
    s
}

fn exhaustive(m: &Bigram, alpha: f64, max_len: usize) -> (Seq, f64) { // DFS over every sequence
    fn go(m: &Bigram, s: &mut Seq, raw: f64, alpha: f64, max_len: usize, best: &mut (Seq, f64)) {
        if s.last() == Some(&EOS) || s.len() == max_len {
            let sc = norm(raw, s.len(), alpha);
            if sc > best.1 { *best = (s.clone(), sc); }
            return;
        }
        for t in 0..m.v { s.push(t); let r = raw + m.logp(s.len().checked_sub(2).map(|i| s[i]), t); go(m, s, r, alpha, max_len, best); s.pop(); }
    }
    let mut best = (vec![], f64::NEG_INFINITY);
    go(m, &mut vec![], 0.0, alpha, max_len, &mut best);
    best
}

fn table(v: usize, rows: &[(usize, &[f64])]) -> Bigram { // listed rows only; others uniform
    let mut p = vec![1.0 / v as f64; (v + 1) * v];
    for &(r, probs) in rows { p[r * v..(r + 1) * v].copy_from_slice(probs); }
    Bigram { v, p }
}

fn main() {
    // Tokens: 0 = EOS, 1 = a, 2 = b, 3 = c; row 4 = BOS.
    let trap = table(4, &[(4, &[0.01, 0.50, 0.40, 0.09]), (1, &[0.28, 0.26, 0.24, 0.22]), (2, &[0.90, 0.04, 0.03, 0.03])]);
    let g = greedy(&trap, 6);
    let (best, _) = exhaustive(&trap, 0.0, 6);
    assert_eq!(g, vec![1, 0]); // greedy: a EOS, P = 0.14
    assert_eq!(best, vec![2, 0]); // best: b EOS, P = 0.36
    assert_eq!(beam_search(&trap, 1, 0.0, 6)[0].0, g); // k = 1 is greedy
    assert_eq!(beam_search(&trap, 2, 0.0, 6)[0].0, best); // k = 2 finds it
    // Tokens: 0 = EOS, 1 = x, 2 = c, 3 = d, 4 = e. Short answer vs. a longer chain.
    let chain = table(5, &[(5, &[0.0001, 0.45, 0.5497, 0.0001, 0.0001]), (1, &[0.8, 0.05, 0.05, 0.05, 0.05]),
        (2, &[0.0375, 0.0375, 0.0375, 0.85, 0.0375]), (3, &[0.0375, 0.0375, 0.0375, 0.0375, 0.85]), (4, &[0.85, 0.0375, 0.0375, 0.0375, 0.0375])]);
    let (s0, s1) = (beam_search(&chain, 4, 0.0, 6)[0].clone(), beam_search(&chain, 4, 1.0, 6)[0].clone());
    assert_ne!(s0.0, s1.0); // the length penalty changes the winner
    for m in [&trap, &chain] {
        for alpha in [0.0, 1.0] {
            let (e, sc) = exhaustive(m, alpha, 6);
            let b = &beam_search(m, 1000, alpha, 6)[0]; // wider than the whole tree = exhaustive
            assert_eq!(b.0, e); assert!((b.1 - sc).abs() < 1e-12);
        }
    }
}

Complexity and constants

• Per step: k model rows, k·V additions, a partial selection of 2k from k·V. Total O(L·k·V) scoring plus L model calls on k-row batches; k KV caches unless shared. The model call dominates.

Efficiency

• Cost: a step is one forward pass over k rows per request, k times the greedy batch. Decoding is mostly memory-bandwidth bound, so k rows cost less than k times the latency until compute saturates (reasoning).
• Levers:
  • small k — beam 2 gave most of the gain: 34.50 vs 34.81 BLEU at 12 (Sutskever et al., ensemble of 5) — translation
  • shared prefix KV blocks — no cache copy per reshuffle (paged KV cache) — long prompts, many beams
  • score-threshold pruning — GNMT drops hypotheses far below the best — large k
• Naive → standard → state of the art:
  • full sort, cache copy per beam · O(kV log kV) per step
  • top-2k selection, finished pool, cache gather · O(kV) per step
  • shared paged prefixes, rows batched across requests (continuous batching) · cache about one decode plus the divergent tails

Implementation notes

• Take 2k candidates (Hugging Face: (EOS count + 1) · beams) so k live beams survive.
• Normalise by generated length, not prompt length.
• Gather cache rows only when the parent list is not the identity (Whisper checks).

Pitfalls

• Wider is not better: Koehn and Knowles saw BLEU fall at large beams, most without length normalisation; Yang et al. report beams above 5 hurt. Stahlberg and Byrne: for over 50 % of sentences the model's global best was the empty translation, which beam search rarely found even at beam 100. The cause is the model or objective (a length reward helps: Murray and Chiang), not search. Meister et al. show beam search exactly solves a different, regularised objective.
• No length normalisation: short outputs win; α too large: rambling.
• CTC: sum alignments per prefix; one entry per alignment wastes the beam.
• Hugging Face moved group and constrained beam search to transformers-community Hub repos.

Variants and alternatives

In production

• Speech: Harpy (1976), an early beam search system; pyctcdecode CTC decoder with an n-gram LM; Whisper (beam 5 at temperature 0).
• Translation: GNMT, CTranslate2, Hugging Face num_beams.
• OCR: Tesseract RecodeBeamSearch over LSTM output.
• Chemistry: Molecular Transformer reaction prediction.

Implementations

Sources

• Lowerre, The HARPY Speech Recognition System, PhD thesis, CMU, 1976 (no link); Lowerre and Reddy, The Harpy Speech Understanding System, in Readings in Speech Recognition, 1990. DOI
• Sutskever, Vinyals, Le, Sequence to Sequence Learning with Neural Networks, NIPS 2014. paper
• Wu et al., Google's Neural Machine Translation System, 2016, §7. paper
• Vijayakumar et al., Diverse Beam Search, 2016. paper
• Koehn, Knowles, Six Challenges for Neural Machine Translation, 2017. paper
• Murray, Chiang, Correcting Length Bias in Neural Machine Translation, 2018. paper
• Yang, Huang, Ma, Breaking the Beam Search Curse, EMNLP 2018. paper
• Stahlberg, Byrne, On NMT Search Errors and Model Errors: Cat Got Your Tongue?, EMNLP 2019. paper
• Meister, Vieira, Cotterell, If beam search is the answer, what was the question?, EMNLP 2020. paper
• Kasai et al., A Call for Clarity in Beam Search, 2022. paper
• Gulcehre et al., On Using Monolingual Corpora in NMT, 2015. paper
• Hannun, Sequence Modeling with CTC, Distill 2017. article
• Schwaller et al., Molecular Transformer, 2018. paper
• Code read: Hugging Face generation, Whisper decoding.py, CTranslate2 docs, Tesseract recodebeam.h