Beam Search, in depth

The idea in one paragraph of maths

A sequence's score is the sum of its token log-probabilities: log P(y) = Σ log P(y_t | y_<t, x). Log-probabilities are used because products of many small probabilities underflow, and sums are cheap. Every score is negative, and adding a token can only make it more negative. That one fact explains most of beam search's behaviour. Short sequences have an unfair advantage, so implementations divide by length raised to a power α (the length penalty). It also means that once enough finished hypotheses beat every live beam's current score, no live beam can ever catch up, which gives a safe stopping rule when no length normalisation is applied.

Greedy decoding is beam search with B = 1. Exhaustive search over all sequences is beam search with B equal to the vocabulary size raised to the length, which is impossible. Beam search sits between them: it costs roughly B times greedy decoding and finds sequences with higher model probability, but it is not guaranteed to find the most probable one. Whether the most probable sequence is the one you want is a separate question, and for open-ended text the answer is often no. Sampling strategies covers the alternatives.

A worked example where greedy loses

Take a toy model with a five-token vocabulary: The, A, cat, dog, sat and an end token <e>. From the start, P(The) = 0.5, P(A) = 0.4, P(<e>) = 0.1. After The, the model says cat 0.4, dog 0.35, <e> 0.25. After A, it is very confident: cat 0.9, dog 0.05, <e> 0.05. After any noun, <e> has 0.9 and sat 0.1.

Greedy takes The (0.5), then cat (0.4), then <e> (0.9), for a sequence probability of 0.5 × 0.4 × 0.9 = 0.18. Beam search with B = 2 keeps both The and A after step one. At step two it scores all six extensions: A cat = 0.36, The cat = 0.20, The dog = 0.175, The <e> = 0.125, and two tiny ones. It keeps A cat and The cat. At step three, A cat <e> finishes at 0.324 and The cat <e> at 0.18. The best answer, A cat, started with the second-best first token, which is exactly the case greedy cannot recover from. An exhaustive search over the same model confirms 0.324 is the global optimum here, followed by 0.18 and then The dog at 0.1575.

The loop as implementations write it

The reference loop below is what every production implementation does, minus batching and tensors. It was run against exhaustive search on the toy model above and returns the same top two sequences.

def beam_search(next_logprobs, prompt, eos, B, max_new, alpha=0.0):
    """next_logprobs(seq) -> {token: logprob}; in practice only the top 2B per beam."""
    live = [(0.0, list(prompt))]                 # (raw log-prob, tokens)
    done = []                                    # (normalised, raw, tokens)
    for step in range(1, max_new + 1):
        cand = []
        for score, seq in live:
            for tok, lp in next_logprobs(seq).items():
                cand.append((score + lp, seq + [tok]))
        cand.sort(key=lambda c: c[0], reverse=True)
        live = []
        for score, seq in cand[:2 * B]:          # 2B so B survivors exist even if B end
            if seq[-1] == eos:
                done.append((score / step ** alpha, score, seq))
            else:
                live.append((score, seq))
            if len(live) == B:
                break
        done.sort(key=lambda d: d[0], reverse=True)
        done = done[:B]
        # with alpha == 0 scores only fall, so a full done-list that beats the
        # best live beam can never be overtaken
        if not live or (alpha == 0 and len(done) == B and done[-1][1] >= live[0][0]):
            break
    return done or [(s / max_new ** alpha, s, q) for s, q in live]

Three details matter. First, the top 2B candidates are examined, because up to B of them might end in the end token; taking only B would leave the beam short. Hugging Face and vLLM both use the 2B rule. Second, finished hypotheses leave the beam and go to a separate list. A common bug is to leave them in the beam, where they keep being extended past the end token. Third, the stopping rule is only exact without length normalisation. With α > 0 a longer live beam can still overtake a finished one, so libraries offer a choice between a heuristic stop and running to the length limit.

Beams and the KV cache

The forward pass for B beams is a batch of B sequences, so compute grows about B times. Memory is the bigger problem. Each beam needs its own key-value cache, and at every step the survivors are a re-shuffle of the previous beams: beam 0 might continue from old beam 1, and two survivors might both descend from old beam 0. The cache must follow.

In the Hugging Face implementation the cache is a dense tensor with one row per beam, and after each step it is reordered by the selected beam_idx, a gather along the batch dimension. That is simple and correct, but it moves the whole cache every step, and it stores the shared prompt B times. For a long prompt with B = 4, prompt memory is four times what greedy needs.

The PagedAttention design stores the cache in fixed-size blocks with reference counts. Beams that share a prefix can point at the same blocks, and a block is copied only when beams diverge inside it (copy-on-write), so the shared prompt is stored once. PagedAttention explains the block table. In vLLM's own beam search, the version checked (v0.11.0) runs an outer loop in the entry point: each step it asks the engine for exactly one token per beam with SamplingParams(logprobs=2 * beam_width, max_tokens=1), then selects survivors in Python. It does not fork caches: each step submits every beam's full token sequence as a new one-token request, so reuse of shared prefixes comes from prefix caching, and each step pays a scheduling round-trip. Expect it to be noticeably slower per token than sampling on the same engine.

Beam search in Hugging Face and vLLM

Hugging Face transformers. Beam search is selected by num_beams > 1 with do_sample=False in generate. The main knobs:

out = model.generate(
    **inputs,
    num_beams=4,
    do_sample=False,
    length_penalty=1.0,
    early_stopping=True,
    num_return_sequences=4,
    max_new_tokens=64,
    return_dict_in_generate=True,
    output_scores=True,
)
for seq, score in zip(out.sequences, out.sequences_scores):
    print(round(score.item(), 3), tok.decode(seq, skip_special_tokens=True))

Diverse (group) beam search and constrained beam search have been moved out of the core decoding loop into Hub-hosted generation methods under transformers-community, loaded with custom_generate=... and trust_remote_code=True. Check which applies to your installed version before relying on num_beam_groups or force_words_ids. The generate() deep dive covers the surrounding machinery.

vLLM. Beam search is not a sampling parameter. It has its own entry point and parameter class:

from vllm import LLM
from vllm.sampling_params import BeamSearchParams

llm = LLM(model="your-model")
params = BeamSearchParams(beam_width=4, max_tokens=64, length_penalty=1.0)
outputs = llm.beam_search([{"prompt": "Translate to German: The cat sat."}], params)
for seq in outputs[0].sequences:
    print(seq.cum_logprob, seq.text)

The documented fields are beam_width, max_tokens, ignore_eos, temperature, length_penalty and include_stop_str_in_output. Read the output field names from the version you deploy; they have changed before. Hosted chat APIs generally do not offer beam search at all.

When to use it, and when not to

Beam search finds high-probability sequences. That is the right goal when there is roughly one correct output and the model's probability tracks correctness: translation, speech-to-text, grapheme-to-phoneme, short structured extraction, code completion of a single line, and generating candidates for a reranker. It is the wrong goal for open-ended text. In chat and story generation the highest-probability continuation is often generic and repetitive; human text is not the mode of the distribution. Instruction-tuned models are also tuned and evaluated with sampling-based decoding, so beam search runs them in a regime they were not checked in.

Failure modes

Empty or truncated outputs. The end token is often a high-probability early choice, and without length normalisation the shortest finished hypothesis wins. If outputs are suspiciously short, check length_penalty and that the end token is not appearing at step one.

Repetition loops. Beam search amplifies loops because a repeated phrase becomes more probable each time it appears. no_repeat_ngram_size stops it bluntly; repetition penalties are gentler. Both can break legitimate repetition such as code and lists.

Wider is worse. Raising B past 5-10 often lowers task quality while raising model probability, the beam-search curse. Tune B on a dev set; do not assume more is better.

Near-duplicate n-best lists. The B beams often differ by one token or punctuation. If you need diversity, deduplicate after detokenising, or use group beam search or sampling.

Memory blow-up under load. A request with B = 8 occupies up to eight sequences of KV cache. In a continuously batched server one beam request can evict or delay several normal requests. Cap B per tenant and account beam requests as B requests in admission control.

Streaming does not work naturally. The best prefix can change until the end, so you cannot stream tokens as they are chosen without risking retractions. Stream only the prefix shared by all live beams, or do not stream.

Batch nondeterminism. Ties between candidates are broken by tiny floating-point differences that depend on batch composition, so the same prompt can return different beams under different loads. Log sequences_scores and treat very small score gaps as ties.

Operating it in production

Measure what matters before turning beam search on. Run greedy, beam 2, beam 4 and sampling on a dev set with your task metric, BLEU or chrF for translation, word error rate for speech, exact match for extraction, and record p50 and p99 latency and peak KV-cache memory alongside. If beam 4 does not beat greedy by more than run-to-run noise, keep greedy. Gains above B = 4 are rare.

In serving, treat B as a multiplier on capacity planning: a beam request costs roughly B times the decode compute and, with dense caches, B times the memory including the prompt. With paged caches, prompt memory is shared but generated tokens still grow B-fold. Put beam traffic on its own pool or limit concurrent beam requests, and expose the beam width in request logs so a latency spike can be traced to it.

What to do next

1. Write down whether your task has one correct output. If it does not, use sampling and stop here.
2. Run greedy, B = 2 and B = 4 on a fixed dev set; record the task metric, latency and peak cache memory for each.
3. If beam wins, tune length_penalty on the dev set and check average output length against references.
4. Add a short-output and a repetition check to your evaluation, and set no_repeat_ngram_size only if loops appear.
5. In Hugging Face, log sequences_scores; in vLLM, use LLM.beam_search with BeamSearchParams and benchmark it against sampling on the same engine.
6. Count beam requests as B requests in admission control and alert on beam-width use per tenant.
7. If you need diverse candidates for a reranker, compare beam n-best against n samples before committing to either.