Beam Search › Beam Search in Large Language Models

The adoption of large language models (LLMs) trained on massive text corpora has shifted the dominant decoding paradigm in some applications. GPT-style models are typically used with sampling-based methods rather than beam search for open-ended generation, but beam search remains relevant for constrained tasks even with LLMs.

Large language models trained on diverse, massive datasets have well-calibrated distributions over text. At any given position, there are often many plausible continuations, and the model's probability is spread across a wide range of good options. In this setting, deterministically selecting the maximum-probability continuation may not produce the most natural or interesting text. Sampling from the distribution, possibly with temperature or nucleus constraints, tends to produce outputs that feel more human-like.

There is also a practical consideration: when you use an LLM as a general-purpose assistant, you often want different outputs for the same prompt on different runs. Deterministic decoding (beam search or greedy) always produces the same output for the same input and the same model weights. Stochastic sampling introduces the variability that makes repeated queries useful.

Despite the prevalence of sampling for LLMs, beam search remains the standard choice in several LLM-driven scenarios. When an LLM is used to perform a constrained task such as structured data extraction, code generation for a specific schema, or translation, beam search with a moderate beam width often outperforms sampling methods because the correct output is more narrowly defined.

For machine translation with transformer-based LLMs (such as NLLB or M2M-100), beam search is still the default inference method. The task has a right answer (the correct translation), and the correctness criteria are well-defined enough that maximizing probability is beneficial.

Code generation models like Codex or Code Llama also benefit from beam search when the task involves generating syntactically correct code for a test case. The constraints on code (it must parse, compile, and produce correct output) make the search space more restricted, and beam search explores that space more systematically than sampling.

One of the major practical challenges with beam search for large transformer models is that decoding is inherently sequential: each step requires a forward pass through the full model, and you cannot parallelize across steps because each step depends on the output of the previous one. This is in contrast to training, where all tokens in a sequence can be processed in parallel using the attention mask.

Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) addresses this bottleneck by using a small draft model to propose multiple tokens at once, then using the full model to verify the proposed tokens in a single batched forward pass. If the draft model's proposals are accepted (according to the full model's distribution), multiple tokens are committed simultaneously. If a proposal is rejected, the algorithm falls back to sampling from the full model at the rejection point.

The key insight is that verifying a sequence of $k$ tokens with the large model takes roughly the same compute as generating a single token (because the transformer can process all $k$ tokens in parallel). So if the draft model's acceptance rate is high, you generate multiple tokens per large-model forward pass, achieving significant speedups without changing the output distribution.

Speculative decoding is not a new decoding algorithm in the beam search sense; it is a computational acceleration technique that produces the same distribution over outputs as standard sampling from the large model. But it highlights how the sequential nature of autoregressive decoding is a practical bottleneck, and how the field is actively working to overcome it.