Beam Search › Hypothesis Expansion and Pruning

The core mechanics of beam search involve two repeated operations: expansion and pruning.

At each step, every hypothesis in the beam is extended by every possible token. If the beam has $k$ hypotheses and the vocabulary has $|V|$ tokens, expansion produces $k \times |V|$ candidates. For $k = 5$ and $|V| = 30{,}000$, that is 150,000 candidates per step.

In modern implementations, this is done efficiently using batched matrix operations. The $k$ partial sequences are fed through the decoder in a single batch, producing a $k \times |V|$ 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 $k$ 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.

From the $k \times |V|$ candidates, pruning selects the top $k$ 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 $O(N \log k)$ time rather than $O(N \log N)$ for a full sort, where $N = k \times |V|$.

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.

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.