Beam Search › 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.

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 $k$ beam hypotheses, so the effective batch size grows from $B$ to $B \times k$. 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.

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 $k$ beams, you maintain $k$ 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 $k$ 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.

For large models with long context windows, the memory required to maintain $k$ 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 $k = 5$, needs to store:

$$2 \times 32 \times 32 \times 128 \times 512 \times 5 \approx 670 \text{ MB}$$

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.

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.