Part of Language AI Handbook
Explains how beam search balances quality and compute in sequence generation, covering greedy decoding limits, length normalization, diverse beam search.
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Beam Search
When a sequence model generates text, it faces a combinatorial explosion. At every time step, the vocabulary might contain 30,000 tokens. After just 10 steps, the number of possible sequences is , a number so large it dwarfs the atoms in the observable universe. You obviously cannot evaluate all of them. The decoder must make choices, and those choices determine whether the final output is fluent, accurate, or coherent.
As we covered in the Encoder-Decoder Framework chapter, the decoder operates autoregressively: it consumes its own previous output to predict the next token. And in the Teacher Forcing chapter, we saw that during training the decoder receives ground-truth tokens as input. But at inference time, ground-truth tokens are unavailable. The decoder must choose which token to feed itself next based purely on its probability estimates.
This chapter is about how to make that choice well. We will start with the simplest strategy, show why it systematically fails, and then build up to beam search, the approximate best-first algorithm that has powered machine translation, summarization, and caption generation for years. Along the way, we will examine length normalization, diversity mechanisms, stochastic alternatives, and the practical engineering realities of running beam search at scale. Understanding beam search deeply means understanding how it works, why those design choices exist, and what tradeoffs they reflect.
Greedy Decoding and Why It Fails
The most natural strategy is greedy decoding: at each step, pick the highest-probability token and move on.
At time step , the decoder outputs a probability distribution over the vocabulary . Greedy decoding selects:
where:
- : the token selected at step
- : a candidate token from vocabulary
- : the model's probability for token given all previously generated tokens and the context vector
This is fast and simple. The problem is that it maximizes local probability at each step, but local optima do not stack into a global optimum.
Consider translating the German sentence "Ich habe gestern das Buch gelesen" (I read the book yesterday). Suppose at step 1, the model assigns probability 0.45 to "I" and probability 0.40 to "Yesterday". Greedy decoding commits to "I" immediately. But "Yesterday I read the book" and "I read the book yesterday" are both valid translations with similar quality. If choosing "Yesterday" at step 1 leads to a sequence where all subsequent tokens are extremely high probability, then the full "Yesterday I read the book" sequence might have higher joint probability than "I read the book yesterday". Greedy decoding, having committed to "I" on step 1, will never discover this.
The fundamental issue is that greedy decoding makes irrevocable decisions without considering the downstream consequences. A choice that looks locally optimal may close off globally better paths. The metaphor of hiking comes to mind: if you always walk in the direction that appears to go uphill right now, you will quickly reach a local peak, but you may have missed the real summit by choosing poorly at an early fork in the trail.
The Exposure Bias Connection
There is a deeper pathology here that connects to the exposure bias problem we discussed in the Teacher Forcing chapter. The model was trained on human-generated sequences, where high-probability next tokens also tend to lead to high-probability continuations. But once the model starts generating its own output, errors compound. A greedy choice at step 3 that was slightly suboptimal might, by step 7, have led the model into a distribution that looks very different from anything in training. The model then has no good options and starts producing incoherent text. Greedy decoding is the strategy most vulnerable to this failure mode.
During training, every prefix the model sees is drawn from the training data. The decoder never observes what happens when it makes a mistake because the teacher-forced input always provides the correct previous token. At test time, the model is operating from a distribution it was never explicitly trained to handle: sequences that contain its own errors. Greedy decoding amplifies this mismatch because each step further diverges the model's current state from anything it encountered during training. Beam search partially mitigates this by keeping several paths alive, which provides a form of error recovery: if one hypothesis drifts into a bad region, another hypothesis in the beam may still be in a high-quality region.
Exhaustive Search Is Infeasible
The alternative would be to consider all possible sequences. Given a vocabulary of size and a maximum output length of , the total number of sequences is . For typical NLP settings (, ), this is , which is completely intractable.
What we need is an algorithm that explores more than one path without requiring exponential computation. Beam search achieves this by committing to a controlled level of exploration, keeping the best hypotheses at each step and discarding the rest. It does not guarantee finding the globally optimal sequence, but in practice it finds sequences that are significantly better than greedy decoding at a fraction of the cost of exhaustive search.
Beam Search: Approximate Best-First Decoding
Beam search is a heuristic that maintains a fixed number of candidate hypotheses at each step. This number is called the beam width, written . Instead of greedily committing to the single best token, beam search tracks the best partial sequences at every time step, extending each one by one step, and then pruning back to the best again.
In beam search, a hypothesis is a candidate partial sequence: a sequence of tokens generated so far, paired with its cumulative log-probability. The beam is the set of hypotheses currently under consideration.
The name "beam" is a spatial metaphor. Imagine the search space as a tree where the root is the start token and each level corresponds to one decoding step. The branches fan out to all possible next tokens. Greedy search follows a single path down this tree. Exhaustive search explores every path. Beam search illuminates paths at once, like a beam of light of fixed width sweeping across the tree. When a path scores poorly, it falls out of the beam. When a path scores well, it continues.
The Algorithm
Beam search operates on log-probabilities rather than raw probabilities. Multiplying many probabilities together quickly underflows to zero in floating-point arithmetic. Log-probabilities transform multiplication into addition, solving the underflow problem while preserving the ordering (since logarithm is monotonically increasing).
The joint probability of a sequence is:
Taking the log:
where:
- : the joint probability of the complete sequence
- : the conditional probability of token given previous tokens and context
- : the log-probability of the sequence, computed as a sum of per-step log-probabilities
Since log-probabilities are negative (the log of a probability between 0 and 1 is always negative), higher (less negative) log-probabilities indicate better sequences.
The beam search procedure is:
- Initialize the beam with one hypothesis: the start token with log-probability 0.
- For each step : a. For each hypothesis in the beam (up to hypotheses), run the decoder forward to get a distribution over the next token. b. Generate all candidate extensions (each hypothesis extended by each token). c. Score each candidate: add the new token's log-probability to the hypothesis's current score. d. Prune: keep only the top candidates by score. These become the new beam. e. If a hypothesis produces the end token (EOS), move it to the "completed" list and remove it from the beam.
- When the beam is empty or a maximum length is reached, return the completed hypothesis with the highest score (or the best hypothesis remaining in the beam if none completed).
A Worked Example
Consider a toy vocabulary of five tokens: {A, B, C, D, EOS} and a beam width of 2.
At step 0, the beam contains one hypothesis: [<START>] with score 0.
The decoder outputs probabilities for the first token:
A: 0.45B: 0.35C: 0.15D: 0.04EOS: 0.01
In log-space: log(0.45) = -0.80, log(0.35) = -1.05, log(0.15) = -1.90, etc.
With beam width 2, we keep A (score -0.80) and B (score -1.05). The beam now contains two hypotheses.
At step 1, each hypothesis is extended by all 5 tokens. We have 10 candidates. Suppose the decoder, after seeing A, assigns probability 0.60 to EOS (score: ) and 0.25 to B (score: ). After seeing B, the decoder assigns probability 0.50 to A (score: ) and 0.30 to C (score: ).
The best 2 candidates from all 10 are A-EOS (score -1.31) and B-A (score -1.74). A-EOS is a completed sequence (it produced EOS), so it moves to the completed list. The beam now contains only B-A.
The process continues until either all hypotheses complete or the maximum length is reached. The final answer is the completed sequence with the best score, which here is A-EOS.
Notice what happened: greedy decoding would have generated A at step 1, gotten EOS at step 2, and returned A. But if B-A-EOS had a higher joint probability than A-EOS, beam search with width 2 would have found it, while greedy search would not.
A key property to notice is that the beam at any given step may contain hypotheses of different lengths, all anchored at the same decoding step. The pruning compares raw cumulative log-probabilities across these hypotheses. This is why length normalization matters: a hypothesis that has accumulated more log-probability simply by being shorter is not necessarily better. We will return to this in the length normalization section.
Beam Width: The Quality-Compute Tradeoff
The beam width is the central hyperparameter of beam search. It controls a direct tradeoff between the quality of the search and the computational cost.
With , beam search reduces exactly to greedy decoding. With , it becomes exhaustive search. In practice, values between 4 and 10 are common for machine translation; some summarization models use 4, and neural machine translation systems often use 5 or 6.
Why Larger Beams Are Not Always Better
Increasing beam width has diminishing returns in practice. The gain from going from to is typically much larger than from to . Beyond a certain width, the model has already found most of the high-probability paths, and adding more beams just increases computation without improving output quality.
More surprisingly, very large beam widths can degrade output quality on some tasks. Researchers have found that in neural machine translation, BLEU scores sometimes peak around to and then decline for larger beams. The reason is related to the model's training distribution: a model trained with maximum likelihood estimation (MLE) learns to assign high probability to sequences that start similarly to the training data. When the beam is large enough to explore low-probability starts, the model encounters territory far from its training distribution, and the output quality degrades. This phenomenon is sometimes called the "beam search curse".
The beam search curse reveals a deep tension in how neural sequence models are trained and how they are evaluated. During training, the model's objective is to assign high probability to the correct next token given the ground-truth prefix. This is a local, next-step objective. At test time, beam search tries to find the globally most probable sequence, which is a different objective. The model was not trained to optimize global sequence probability; it was trained to optimize step-wise conditional probabilities. For well-distributed training data, these two objectives are roughly aligned. But at the edges of the distribution, they diverge, and beam search with a very wide beam will faithfully optimize the wrong objective.
A related concern is that for certain models, specifically language models trained on extremely large corpora, even greedy decoding may produce excellent outputs because the model has learned such a well-calibrated distribution over text. For these models, beam search still helps but the margin over greedy is smaller than for older, weaker models.
Computational Cost
Each step of beam search requires decoder forward passes (one per hypothesis in the beam). With a vocabulary of size , you compute scores at each step, then sort to find the top . The total computation is where is the cost of a single decoder forward pass. In practice, the forward passes are batched together, so the overhead relative to greedy is roughly the cost per step.
The memory overhead is also non-trivial. Each beam hypothesis maintains its own key-value cache (in transformer models, this stores the attention key and value tensors for all previous tokens). With hypotheses, you hold separate KV caches simultaneously. For long sequences and large models, this can dominate memory usage. Some systems reduce this overhead by sharing computation across hypotheses when they share a common prefix, a technique called hypothesis sharing or tree-structured KV caching.
Length Normalization
Raw log-probability scores are biased against longer sequences. Since every additional token adds a negative log-probability term, longer sequences accumulate more negative contributions and systematically score lower than shorter ones. A model using raw scores will prefer short outputs, sometimes generating incomplete or truncated translations.
To see why this matters, imagine two candidate translations of a German sentence: a five-word English translation and a ten-word English translation. If each token has an average log-probability of -0.7 (a reasonable value for a good translation), the five-word translation scores -3.5 and the ten-word translation scores -7.0. The model will always prefer the five-word translation, even if the ten-word one is linguistically better. The model effectively penalizes every word added to the output.
The length penalty addresses this by normalizing the score by sequence length. One common formulation, introduced by Wu et al. in Google's Neural Machine Translation system, uses:
where the length penalty is:
where:
- : the length of the generated sequence (number of tokens)
- : a hyperparameter controlling how strongly length is penalized, typically set to 0.6 or 0.7
- The constant 5 is a smoothing term that prevents the penalty from being too aggressive for very short sequences
When , the penalty is 1.0 for all lengths, and there is no length normalization (equivalent to raw log-probability). When , the penalty is approximately linear in length, giving nearly equal footing to sequences of different lengths. Values in between, such as 0.6, provide a moderate preference for longer outputs without going all the way to strict length normalization.
Why not simply divide by directly? The constant 5 in the numerator and denominator prevents very short sequences from being penalized too heavily and provides a smoother transition. Without it, the difference between length 1 and length 2 would be enormous (a factor of 2), while the difference between length 50 and length 51 would be negligible.
The choice of has practical consequences. A value of was calibrated empirically on production MT systems. For summarization, where you generally want outputs shorter than the source, a lower (or even ) may be more appropriate. For caption generation, where captions should have a natural length, around 0.5 to 0.8 tends to work well. Tuning on a validation set is standard practice.
Some systems also add a coverage penalty alongside the length penalty. The coverage penalty rewards sequences where attention has covered all source tokens, addressing the tendency of seq2seq models to ignore parts of the input. The combined score balances fluency (via the model's probability), adequate length, and adequate coverage.
Hypothesis Expansion and Pruning
The core mechanics of beam search involve two repeated operations: expansion and pruning.
Expansion
At each step, every hypothesis in the beam is extended by every possible token. If the beam has hypotheses and the vocabulary has tokens, expansion produces candidates. For and , that is 150,000 candidates per step.
In modern implementations, this is done efficiently using batched matrix operations. The partial sequences are fed through the decoder in a single batch, producing a score matrix. The scores are added to the current hypothesis scores to get candidate scores.
One important subtlety is that the decoder has state. For RNN-based models, the state is the hidden vector; for transformer-based models, the state is the key-value cache. When you batch the hypotheses together, you need to maintain separate states for each hypothesis. If two hypotheses share a common prefix (which is common early in decoding, when the beam is initialized from a single start token), you can share their cached keys and values to save memory. As the hypotheses diverge, the caches diverge too.
Pruning
From the candidates, pruning selects the top by score (accounting for length normalization if used). This is a partial sort, which is more efficient than a full sort of 150,000 elements. Libraries like NumPy and PyTorch provide highly optimized topk operations that perform this partial sort in time rather than for a full sort, where .
When a candidate ends with EOS, it is moved from the beam to the completed hypotheses list. This is important because once a hypothesis completes, it should not continue accumulating log-probability terms. If a completed hypothesis stayed in the beam, its score would keep growing more negative with each additional step (since log-probabilities are negative), even though the sequence is done. Removing completed hypotheses from the beam ensures they are scored at their natural endpoint.
The beam continues with the remaining hypotheses until either the beam is full of EOS-completing sequences, the maximum length is reached, or the beam is empty. In some implementations, the search terminates early when the best completed hypothesis has a higher score than the best remaining hypothesis in the beam, since no future extension can improve upon it.
Tie-Breaking and Minimum Length
Several practical considerations modify the basic algorithm. A minimum output length can be enforced by setting the probability of EOS to zero until a minimum length is reached. This prevents degenerate outputs where the model immediately predicts EOS. Without this guard, models that have learned to assign high probability to EOS from short sequences (a training artifact) will produce one- or two-token translations.
Similarly, certain repeated phrases can be suppressed using n-gram blocking: if a bigram or trigram has already appeared in a hypothesis, its continuation tokens are penalized or zeroed out. This is particularly effective for summarization, where models sometimes fall into repetition loops, generating the same phrase multiple times. N-gram blocking was popularized by the FAIRSEQ library and is now a standard option in most sequence generation frameworks.
Another common modification is early stopping: once the first completed hypothesis is produced, the search continues for a fixed number of additional steps (rather than until the beam is exhausted). This balances thoroughness with efficiency in production deployments.
Diverse Beam Search
Standard beam search tends to produce hypotheses that are very similar to each other. When the model is confident about the first few tokens, all beam hypotheses will start with the same token, diverging only in later positions. The beam is homogeneous rather than diverse.
This is a problem in applications that need multiple distinct outputs, such as generating several translation candidates for a human post-editor to choose from, or generating multiple image captions and selecting the best one. When a human reviewer receives five near-identical outputs, the "top-k best" functionality provides little benefit. What is needed is outputs that represent clearly different ways of completing the task.
The Grouping Approach
Diverse beam search (Vijayakumar et al., 2018) addresses this by partitioning the beams into groups, each of size . Within each group, standard beam search runs. Between groups, a diversity penalty discourages groups from producing similar tokens.
At each step , for group , the score for token is:
where:
- : the model's log-probability for token
- : the diversity strength hyperparameter
- : a diversity penalty measuring how many hypotheses in earlier groups already chose token at this step
The penalty encourages each group to select tokens that prior groups have not chosen, spreading the search across different parts of the output space. The first group runs without penalty. The second group is penalized for tokens chosen by the first group. The third group is penalized for tokens chosen by the first two, and so on.
With , diverse beam search reduces to standard beam search. Larger values of force more diversity but may reduce the individual quality of each output.
The grouping mechanism has an intuitive interpretation. Think of each group as a different "opinion" about how the sequence should start. Group 1 follows the most probable path. Group 2 is forced to explore a different starting direction. Group 3 explores a third direction. By the end of decoding, you have distinct hypotheses, one from each group, that cover a broader range of the output space than hypotheses from standard beam search.
Stochastic Beam Search
Another approach to diversity is stochastic beam search, which introduces randomness into the candidate selection step. Instead of deterministically keeping the top candidates, stochastic beam search samples the survivors according to their scores, using a Gumbel noise perturbation. Concretely, it adds i.i.d. Gumbel noise to each candidate score before selecting the top :
This corresponds to sampling without replacement from the categorical distribution defined by the candidates' scores. The result is a diverse set of hypotheses while still being biased toward high-scoring candidates. Stochastic beam search interpolates between deterministic beam search (no noise) and pure sampling (infinite noise).
Beam Search for Machine Translation
Machine translation is the canonical application of beam search. The encoder reads the source sentence and produces hidden states. The decoder generates the translation token by token, and beam search determines which partial translations to keep exploring. The combination of a strong encoder-decoder model and beam search was the backbone of Google Translate, Microsoft Translator, and virtually every production neural MT system from 2016 onward.
In practice, the beam search procedure is tightly integrated with the encoder-decoder architecture.
At inference time, the encoder runs once to produce the source representations. These are then queried by the decoder's attention mechanism (which we will explore in depth in the Attention Intuition chapter) at every step for every beam hypothesis. Since all beam hypotheses share the same source representation, the encoder computation is not repeated. Only the decoder must be run repeatedly.
The decoder, however, must run times per step (once per beam hypothesis), or equivalently, the decoder is batched with a batch size of . Each batch element corresponds to a different beam hypothesis. The score matrix is computed in one batched forward pass, exploiting GPU parallelism to keep the overhead manageable.
At the end, the beam contains up to completed translations (or partial translations if the maximum length was reached). After applying length normalization, the highest-scoring translation is returned as the system output. Some systems also return the top-N translations (n-best list) for downstream reranking: a separate reranker model can score the n-best candidates using additional features (e.g., a language model score, a roundtrip translation score) and select the best one.
The importance of beam search to machine translation is hard to overstate. The first widely deployed neural MT systems achieved competitive quality in part because beam search found translations that greedy decoding missed. The gap between greedy and beam search output quality was large enough that early practitioners considered beam width a first-class hyperparameter alongside learning rate and model size.
BLEU Score Sensitivity to Beam Width
Beam width has a measurable impact on BLEU scores in machine translation. The relationship is roughly: wider beams improve BLEU up to a point, then plateau or slightly decline. For English-German translation with a strong transformer model, going from (greedy) to typically improves BLEU by 1-2 points. Going from to may add another 0.2-0.5 BLEU. The marginal gain diminishes rapidly.
BLEU is a reference-based metric that compares the model's output to human-written reference translations. The fact that beam search improves BLEU does not mean it produces the "best" translation in some objective sense; it means it produces translations more similar to human references. For tasks where the reference distribution is highly multimodal (many equally valid outputs), BLEU may not fully capture quality differences between decoding strategies, and human evaluation often shows a different picture than automatic metrics.
Beyond Translation: Summarization, Captioning, and Q&A
Beam search extends naturally to other seq2seq tasks, though the optimal hyperparameter settings vary. For abstractive summarization, models like BART and PEGASUS typically use beam widths of 4 to 8, with no-repeat-ngram-size set to 3 to prevent repetitive phrases. The length penalty is often set lower than in translation, because summaries should be concise and the model may need less encouragement to generate longer outputs.
For image captioning, beam search with width 3 to 5 is standard. The encoder here is a visual feature extractor (a convolutional network or a vision transformer) rather than a text encoder, but the decoder runs identically to the text case. Beam search ensures that the caption is coherent and grammatically complete rather than trailing off after the first few high-probability words.
For question answering with extractive or generative models, the situation differs. Extractive QA systems that select spans from a passage do not need beam search. Generative QA systems do, but the beam is often narrow (width 2 to 4) because the answer space is more constrained and the model is typically more confident.
Nucleus Sampling and Top-k Sampling: Stochastic Alternatives
Beam search maximizes probability. For some applications, this is exactly what you want: the most likely translation, the most probable answer. But for creative text generation, the most probable text is often the least interesting. When a language model assigns high probability to the same few words repeatedly, the output becomes repetitive and bland.
Consider generating a short story. The highest-probability continuation of "She walked into the dark room and saw" might be "a light" because those two words appear frequently in the training data after similar contexts. But "a light" is boring. The story might be much more interesting if the model picked a lower-probability but more unexpected continuation, like "her reflection looking back at her from a shattered mirror". Maximizing probability produces the average, the statistically most typical continuation. Creative writing, by definition, should not be average.
Sampling-based methods introduce stochasticity to break this pattern.
Top-k Sampling
Top-k sampling restricts the sampling pool to the most probable tokens at each step, then samples from that restricted distribution:
where:
- : the model's original probability for token
- : the set of the tokens with highest probability
- : the renormalized probability used for sampling
The problem with top-k is that the appropriate value of varies with context. If the model is very confident (the top-1 probability is 0.99), then sampling from the top-50 includes many very unlikely tokens and can produce incoherent output. If the model is very uncertain (probabilities are nearly uniform), then restricting to top-5 throws away most of the valid continuations.
Nucleus Sampling (Top-p Sampling)
Nucleus sampling (Holtzman et al., 2020) addresses this by sampling from the smallest set of tokens whose cumulative probability exceeds a threshold :
where:
- : the nucleus, the minimal set of tokens covering cumulative probability mass at least
- : a candidate subset of the vocabulary, sorted by descending probability
- : the probability threshold, typically 0.9 or 0.95
The set is constructed by sorting all tokens in descending probability order and including tokens one by one until the running sum of probabilities reaches .
Tokens outside the nucleus receive zero probability. Those inside are renormalized.
The key advantage is adaptivity. When the model is confident, the nucleus is small (a few tokens cover 0.9 of the probability mass). When the model is uncertain, the nucleus is large (many tokens are needed to accumulate 0.9). The sampling pool automatically adjusts to the local probability distribution.
Nucleus sampling is the default choice for open-ended text generation tasks: story writing, dialogue systems, creative applications. Beam search is better suited for constrained generation tasks where there is a specific correct answer: translation, summarization, question answering. The choice between them reflects a deeper question about what "good output" means for a given task.
Temperature Sampling
Temperature is a parameter that reshapes the probability distribution before sampling:
where:
- : the temperature parameter
- : the unnormalized probability with temperature scaling applied
For , the distribution sharpens: high-probability tokens become even more likely, low-probability tokens shrink. This makes output more focused and repetitive. For , the distribution flattens: probabilities become more uniform, and sampling becomes more exploratory and unpredictable. At , sampling is from the unmodified distribution. Temperature is often combined with top-k or nucleus sampling.
In modern large language model APIs, you typically encounter all three of these controls: temperature, top-k, and top-p. A common practical recipe for creative generation is temperature around 0.8 to 1.0 combined with top-p around 0.9. For factual generation, temperature near 0.2 to 0.5 with greedy or small beam search tends to produce more reliable outputs.
Contrastive Decoding
A more recent alternative, contrastive decoding (Li et al., 2022), generates text by contrasting the outputs of a strong model and a weak model. At each step, the score for a token is computed as the log-probability difference between the expert (strong) model and the amateur (weak) model:
The intuition is that tokens that are highly probable under both models are likely to be generic, high-frequency words that appear frequently in training data regardless of context. Tokens that are much more probable under the expert than the amateur are the ones that reflect the expert's superior understanding of the context. By amplifying this contrast, contrastive decoding produces output that is more informative and contextually appropriate, avoiding the repetition and generic phrasing that characterize maximum-probability decoding.
Contrastive decoding has shown promising results on long-form question answering and knowledge-intensive tasks, where avoiding generic outputs is especially important.
Implementation with Priority Queues
A clean beam search implementation uses a priority queue (min-heap) to efficiently maintain the top-k hypotheses. The data structures involved are:
- A hypothesis object containing the current token sequence and the cumulative log-probability.
- A priority queue indexed by negative log-probability (since Python's heapq is a min-heap, we negate scores so the highest-scoring hypotheses are at the top when popped).
- A completed list for hypotheses that have produced EOS.
Let's implement beam search from scratch on a toy language model.
from typing import List
import numpy as np
# A simple toy language model using random but reproducible probabilities
np.random.seed(42)
VOCAB = [
"<EOS>",
"the",
"cat",
"sat",
"on",
"mat",
"dog",
"ran",
"away",
"quickly",
]
VOCAB_SIZE = len(VOCAB)
EOS_ID = 0
def toy_lm_step(prev_tokens: List[int]) -> np.ndarray:
"""
Toy language model: returns log-probabilities over the vocabulary.
In a real system, this would be a transformer or RNN forward pass.
"""
rng = np.random.RandomState(seed=sum(prev_tokens) % 2**31)
logits = rng.randn(VOCAB_SIZE)
# Suppress EOS early, encourage it later
if len(prev_tokens) < 3:
logits[EOS_ID] -= 5.0
elif len(prev_tokens) >= 6:
logits[EOS_ID] += 3.0
# Softmax
logits -= logits.max()
probs = np.exp(logits)
probs /= probs.sum()
return np.log(probs)class Hypothesis:
"""A beam search hypothesis: token sequence + cumulative log-prob."""
def __init__(self, tokens: List[int], score: float):
self.tokens = tokens
self.score = score
def __lt__(self, other: "Hypothesis") -> bool:
# Sort descending by score: higher score is "less than" for min-heap inversion
return self.score > other.score
def __eq__(self, other: "Hypothesis") -> bool:
return self.score == other.scorefrom typing import List, Tuple
def length_penalty(length: int, alpha: float = 0.6) -> float:
"""Google's length penalty from Wu et al. 2016."""
return ((5 + length) ** alpha) / (6**alpha)
def beam_search(
start_token: int,
beam_width: int = 3,
max_length: int = 8,
alpha: float = 0.6,
) -> List[Tuple[List[int], float]]:
"""
Run beam search and return the top completed hypotheses.
Returns list of (token_sequence, normalized_score) tuples.
"""
# Initialize beam with a single hypothesis containing the start token
beam = [Hypothesis(tokens=[start_token], score=0.0)]
completed = []
for step in range(max_length):
candidates = []
for hyp in beam:
# Get log-probs from the language model
log_probs = toy_lm_step(hyp.tokens)
# Expand this hypothesis by every vocabulary token
for token_id, log_prob in enumerate(log_probs):
new_tokens = hyp.tokens + [token_id]
new_score = hyp.score + log_prob
candidates.append(Hypothesis(new_tokens, new_score))
# Sort all candidates, keep top beam_width
candidates.sort() # Sorts by neg_score ascending = score descending
beam = candidates[:beam_width]
# Move completed hypotheses (ending in EOS) to the completed list
still_active = []
for hyp in beam:
if hyp.tokens[-1] == EOS_ID:
# Apply length normalization (exclude start token from length)
seq_len = len(hyp.tokens) - 1
normalized = hyp.score / length_penalty(seq_len, alpha)
completed.append((hyp.tokens, normalized))
else:
still_active.append(hyp)
beam = still_active
if not beam:
break
# Also collect remaining beam hypotheses with length normalization
for hyp in beam:
seq_len = len(hyp.tokens) - 1
normalized = hyp.score / length_penalty(seq_len, alpha)
completed.append((hyp.tokens, normalized))
# Sort by normalized score (descending)
completed.sort(key=lambda x: x[1], reverse=True)
return completedBeam Search Results (start='the'):
============================================================
Beam width k=1:
Best sequence: the dog away the ran quickly <EOS>
Normalized score: -4.9220
Beam width k=3:
Best sequence: the on cat <EOS>
Normalized score: -2.8314
Top 3 hypotheses:
the on cat <EOS> score=-2.8314
the dog away the <EOS> score=-4.3814
the dog away on dog ran <EOS> score=-4.8244
Beam width k=5:
Best sequence: the on cat <EOS>
Normalized score: -2.8314
Top 3 hypotheses:
the on cat <EOS> score=-2.8314
the dog sat <EOS> score=-3.3536
the dog away the <EOS> score=-4.3814The output shows how beam width affects the search. With (greedy), we get the single locally optimal path. With larger beams, we explore more of the probability space and often find sequences with higher normalized scores.
Effect of Length Normalization (alpha parameter): ============================================================ alpha=0.0 (no normalization) - top 3: [3 tokens] the on cat <EOS> score=-3.3648 [4 tokens] the dog away the <EOS> score=-5.5882 [6 tokens] the dog away on dog ran <EOS> score=-6.9405 alpha=0.6 (moderate normalization) - top 3: [3 tokens] the on cat <EOS> score=-2.8314 [4 tokens] the dog away the <EOS> score=-4.3814 [6 tokens] the dog away on dog ran <EOS> score=-4.8244
With no normalization (), the algorithm often prefers shorter sequences because each additional token adds a negative contribution. With , sequences of moderate length are rewarded, and the top hypotheses tend to be longer and more complete.
Key Parameters
The key parameters for beam search are:
- beam_width (k): Number of hypotheses maintained at each step. Larger values explore more of the search space but increase computation by a factor of . Common values: 4-6 for translation, 1-4 for large language model inference.
- alpha (length penalty): Controls how strongly length is penalized. 0 means no normalization; 1 means full normalization; 0.6 is a common practical default.
- max_length: Maximum output length. Usually set to 1.5 to 2 times the source length for translation.
- min_length: Minimum output length. Set to prevent degenerate EOS predictions.
- no_repeat_ngram_size: Block repeated n-grams to reduce repetition. Setting this to 3 or 4 prevents the model from repeating trigrams or 4-grams.
Visualizing Beam Search
Let's build a visualization of the beam search tree to see how hypotheses are expanded and pruned.




These plots illustrate the beam search quality-compute curve directly. BLEU improves rapidly for small but plateaus and slightly reverses for large . Computation grows approximately linearly. The sweet spot for most production systems is to .

The visualization shows how each decoding strategy selects from the probability distribution. Beam search takes only the peak. Top-k always samples from the top 5, even if those tokens span a wide range of probabilities. Nucleus sampling adapts: here, with a peaked distribution, only 3 tokens are needed to cover 90% of the mass, so the nucleus is smaller than top-k.
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.
Why LLMs Often Prefer Sampling
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.
When LLMs Still Use Beam Search
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.
Speculative Decoding: Accelerating Sequential Generation
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 tokens with the large model takes roughly the same compute as generating a single token (because the transformer can process all 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.
Practical Engineering Considerations
Deploying beam search in production requires attention to several engineering details that are easy to overlook when studying the algorithm in isolation.
Batched Inference
In training, sequences are batched together to exploit parallelism. At inference time, beam search complicates batching in two ways. First, each input sequence in a batch generates beam hypotheses, so the effective batch size grows from to . Second, different sequences in a batch may complete at different steps, leading to variable-length computation. Libraries like Hugging Face Transformers handle this through careful batch management, padding finished beams with EOS tokens so they can remain in the batch without contributing to score updates.
KV Cache Management
Transformer models use key-value caches to avoid recomputing attention keys and values for all previous tokens at each decoding step. For beam search with beams, you maintain separate KV caches. When the beam is pruned, the surviving hypotheses may come from different parent hypotheses; the corresponding KV caches must be reindexed (not just truncated). This reindexing step is a nontrivial engineering concern in high-performance implementations.
Some frameworks implement "beam search with shared KV caches" where all hypotheses share prefixes of the KV cache and only diverge where the hypotheses diverge. This is especially effective early in decoding when all hypotheses have the same start token. As the beams diverge, the shared portion shrinks, but the early savings can still be significant.
Memory Overhead
For large models with long context windows, the memory required to maintain KV caches can be substantial. A model with 32 layers, 32 attention heads, and a head dimension of 128, decoding a sequence of length 512 with beam width , needs to store:
of KV cache (2 for keys and values, in float16). For very large models or very long sequences, this memory constraint can be the binding factor limiting beam width in practice.
N-Best Lists and Reranking
In production MT systems, beam search is often used to generate an n-best list of candidate translations rather than a single output. A second-stage reranker, which may be a separate model or a heuristic scoring function, then selects the best candidate from the list. The reranker has access to the full candidate sequence and can apply features that are expensive to compute during decoding (e.g., roundtrip translation, language model score, length-to-source-ratio check). This two-stage pipeline often outperforms single-stage beam search with a wider beam, because the reranker can apply richer features than the sequence probability alone.
Limitations
Beam search has served the NLP community well, but it carries several important limitations.
The most fundamental is the label bias problem. At each step, beam search scores sequences using a locally normalized probability (the decoder outputs a softmax, which sums to 1 over all next tokens). This means a model can "hide" probability by distributing it among many similar continuations, making it harder for beam search to discover sequences that diverge significantly from common patterns. Maximum Entropy Markov Models (MEMMs) have a well-known variant of this problem, and neural seq2seq models are not fully immune.
A second limitation is that beam search finds the maximum a posteriori (MAP) sequence under the model. For tasks where the model has a sharp, well-peaked distribution (like factual question answering), MAP decoding gives the right answer. But for tasks where many good outputs exist (like open-ended dialogue or creative writing), MAP decoding returns the one sequence the model assigns highest probability to, which tends to be generic and safe rather than interesting and varied. The sampling methods discussed above address this but introduce their own challenges.
A third limitation is computational: beam search is sequential. Each step depends on the previous one because the decoder must consume its own output. This makes it incompatible with the kind of fully parallelized computation that made transformers so efficient during training. Speculative decoding and parallel decoding methods attempt to address this, but they require significant engineering and approximate the beam search guarantee rather than preserving it exactly.
Finally, beam search is sensitive to the model's calibration. If the model systematically assigns incorrect probabilities (overconfident or underconfident), beam search faithfully optimizes those incorrect probabilities. A poorly calibrated model with wider beams may produce confidently wrong outputs. Sampling from a well-calibrated model might outperform beam search from a miscalibrated one.
These limitations are not just academic concerns. The NLP community has moved toward sampling-based decoding for many applications partly because beam search's commitment to maximum probability produces outputs that, while internally coherent, feel stilted and over-optimized. The "degenerate text" problem, where very long beam-search outputs fall into repetition, was a major practical problem in early neural MT systems. Solutions like no-repeat-ngram blocking were developed specifically because the MAP sequence under many models includes repeated n-grams, which are penalized by beam search's length normalization but not eliminated.
Despite these limitations, beam search remains a foundational algorithm. Understanding its mechanics is essential for understanding why models behave the way they do at inference time, and for diagnosing problems when output quality is poor. Every practitioner who works with sequence generation models will encounter beam search in some form, whether tuning the beam width for a translation task, debugging a repetition problem in a summarization model, or choosing between beam search and sampling for a generative application.
Summary
Beam search sits between the extremes of greedy decoding (single path, no backtracking) and exhaustive search (all paths, computationally infeasible). By maintaining a beam of candidate hypotheses and pruning at each step, it explores the probability space efficiently and usually finds significantly better sequences than greedy decoding.
The key ideas to take away are:
- Greedy decoding fails because locally optimal choices lead to globally suboptimal sequences. Irrevocable step-wise decisions close off globally better paths.
- Beam search maintains hypotheses, expands each by the full vocabulary, and keeps the top at every step, trading compute for a more thorough search.
- Beam width trades quality for compute; values of 4-6 work well for translation, larger values show diminishing returns and can even hurt quality due to the beam search curse.
- Length normalization corrects the score bias against longer sequences using the penalty , with being a common default.
- Diverse beam search partitions the beam into groups with inter-group diversity penalties to produce varied outputs, useful when multiple distinct hypotheses are needed.
- Nucleus sampling is the preferred alternative for creative generation, adapting the sampling pool to the local probability distribution and producing more natural, varied text.
- Large language models often use sampling rather than beam search for open-ended generation, but beam search remains the standard for constrained tasks like translation and structured generation.
- Engineering details such as KV cache management, batched inference, and n-best reranking matter significantly in production deployment and often determine whether beam search is practical at a given scale.
In the next chapter, we will explore the Attention Intuition that allows seq2seq models to align source and target sequences more effectively, moving away from the fixed-size context vector bottleneck that constrains the encoder-decoder models we have studied so far.
Quiz
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about beam search and sequence decoding strategies.
Beam Search Quiz
Reference
Citation details
Cite or share this article.
Continue with the full handbook
This chapter is part of Language AI Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.
Explore Language AI HandbookStay up to date
Get articles, book updates, and news delivered to your inbox.
No spam, unsubscribe anytime.
Join the community
Sign in to remove popups, track your reading progress, and join the discussion.

Comments
No comments yet. Be the first to share your thoughts!