Beam Search Decoding: Scoring, Length Normalization, and Pruning

A decoder with a softmax over a vocabulary of $|V|$ symbols has $|V|^{T}$ possible outputs of length $T$. Choosing the most probable symbol at each step is optimal for nothing in particular: the sequence it produces is the product of per-step maxima, and the sequence of largest total probability can begin with a symbol that is not the per-step maximum anywhere. Exact search over the space is exponential, so decoding is a search problem with an approximation, and beam search is the approximation that keeps a fixed number of partial sequences alive at each step.

Keeping several hypotheses alive

Beam search maintains a set of at most $k$ prefixes, each with an accumulated score. At each step every surviving prefix is extended by every symbol in the vocabulary, giving up to $k \cdot |V|$ candidate continuations, and the best $k$ of those by score are retained while the rest are discarded. A prefix that has emitted the end-of-sequence symbol is finished and set aside rather than extended, and generation stops when the beam contains only finished hypotheses or a maximum length is reached. The output is the highest-scoring finished hypothesis, or, if the search ran out of steps, the best partial one by whatever rule the caller defines.

The score is what makes the search well defined, and it has to be a sum rather than a product. Multiplying per-step probabilities underflows: a step probability of $10^{-3}$ repeated 100 times is $10^{-300}$, at the edge of what a 32-bit float represents, and the same computation in log space is an accumulation of 100 additions of about $-6.9$. Every implementation therefore accumulates $\log p$ and compares sums. The comparison also decides which candidates survive, and ties in the floating-point score are resolved by ordering, which makes the result depend on the order in which ties are broken unless the vocabulary is processed in a fixed order.

One property of beam search is worth stating precisely: with $k$ equal to the number of possible prefixes the search is exact for a fixed length, and with $k = 1$ it degenerates to greedy decoding, but no bounded $k$ guarantees the highest-probability sequence. A hypothesis discarded at step $t$ can never be recovered, and the true optimum may pass through a low-probability prefix early on, because a prefix that looks weak at step 5 may be the only one whose continuations have high probability later.

Stopping has to be decided explicitly, because the search can end in three ways. It can exhaust the maximum length, in which case the best partial hypothesis is returned and the output is truncated by the budget rather than by the model. It can finish every hypothesis in the beam, which is the clean case. Or it can fill part of the beam with finished hypotheses while unfinished ones remain, in which case continuing is worthwhile only if a partial score plus an optimistic bound on the remaining steps can still beat the finished set. That bound is what makes early stopping safe: a partial hypothesis can improve by at most zero per remaining step, since log-probabilities are non-positive, so once the best finished score exceeds the best partial score the search can stop without losing the optimum.

The length bias

Each step adds a non-positive number to the score, so among hypotheses of different lengths the shorter ones are penalized less by construction. A search that compares raw sums therefore favours short outputs, and the end-of-sequence symbol contributes its own negative log-probability to the total, which creates an additional cost for stopping. Left uncorrected, the search tends to terminate early, and the size of the effect grows with the typical length of the output.

Length normalization divides the accumulated score by a term that grows with the hypothesis length, so that hypotheses of different lengths become comparable. The precise form varies between systems: dividing by $|y|$, by $|y|^{\alpha}$ for an exponent between 0 and 1, or by a smoother function of the length that is close to linear for long outputs and close to constant for very short ones. An exponent of 1 divides by the length exactly and is the most aggressive correction, an exponent of 0 disables normalization entirely, and the useful range for translation-style tasks sits between those extremes. The choice changes the output more than beam width does for otherwise identical models: a length penalty that is too weak produces truncated outputs, and one that is too strong produces hypotheses that keep generating past the natural end.

The deeper issue is that the objective being optimized is the model’s own likelihood, not the quality of the output. A sequence model trained with maximum likelihood assigns the highest probability to sequences that resemble the training distribution most closely, which for open-ended generation means the most predictable continuation, and the highest-likelihood sequence is often bland or repetitive rather than the one a reader would consider correct. Increasing beam width makes the search find that sequence more reliably, which is why beam search with a large width is standard for tasks with a definite correct output, such as translation, and why it is often paired with sampling, truncation of the tail of the distribution, or explicit diversity objectives for open-ended generation.

Pruning beyond the beam

Naive beam search spends a step generating $k \cdot |V|$ candidates and keeping $k$, so the arithmetic per token is proportional to the beam width. Several prunings cut that without changing the retained set in most cases. The first is to extend each prefix only by the top few symbols under the model’s distribution, since the remaining candidates have scores too low to enter a beam of the given width unless the beam is nearly empty. The second is a target-length prune, which discards hypotheses that cannot reach the shortest plausible length or have passed the longest one. The third is a score bound: once the finished hypotheses have scores better than a bound derived from the current best partial score and an upper bound on the remaining steps, the search can stop early, since no continuation of the beam can beat what has already been found.

Distinct-beam variants address a different problem, that hypotheses inside a beam are often near-duplicates of one another, which wastes width on the same output and causes the search to lose alternatives that were one symbol away. Forcing retained hypotheses to differ in some part of the sequence, or decoding several groups in parallel with cross-group penalties, keeps the effective diversity of the beam higher at the same cost.

import math

def beam_search(next_logprobs, k=4, eos="</s>", max_len=32):
    """next_logprobs(prefix) -> list of (symbol, logprob) over the vocabulary."""
    finished = []
    beam = [((), 0.0)]                      # (prefix, accumulated log-probability)
    for _ in range(max_len):
        candidates = []
        for prefix, score in beam:
            for symbol, lp in next_logprobs(prefix):
                if symbol == eos:
                    finished.append((prefix, score + lp))
                else:
                    candidates.append((prefix + (symbol,), score + lp))
        if not candidates:
            break
        candidates.sort(key=lambda item: item[1], reverse=True)
        beam = candidates[:k]               # discarded prefixes are lost for good
    pool = finished or beam
    # length-normalized comparison; exponent 1 divides by the token count
    return max(pool, key=lambda item: item[1] / max(len(item[0]), 1))

The loop shows the two decisions that dominate behaviour: what is accumulated, and how the accumulation is compared across lengths. It also shows why a decoder with per-hypothesis state needs care. Each prefix carries the decoder state that produced it, so when the beam is reordered the state rows must be gathered by the parent index of each surviving candidate rather than by the candidate’s position in the beam, and a mistake there silently mixes the state of one hypothesis into another. Batched implementations store one state row per beam entry per sequence and index it explicitly, which is the same bookkeeping that batching variable-length inputs requires, described in padding and masking for variable-length batches.

Scoring variants that are not log-probability sums

The accumulated sum is only one possible score. When the decoder is a transducer whose outputs include repeated labels and blanks, the correct quantity to accumulate is the probability of the collapsed labelling rather than of one path, so hypotheses must merge the paths that share a prefix and keep the two ending-state values separately, as described in CTC and monotonic alignment. When the decoder is a recurrent generator over a target sequence, the model’s own conditionals already account for what it has emitted, so the sum is well defined step by step; the dependence on the previous prefix is exactly what makes the search sequential and the beam state non-transferable between steps.

Three properties of the search follow from the way it is built. It is greedy with respect to a fixed width, so an increase in width can change the output in either direction as far as quality is concerned. It optimizes the model’s likelihood, which is not the same objective as the one the output is judged by. It is sequential, so the latency of producing $T$ symbols is $T$ steps of decoder evaluation over $k \cdot |V|$ candidates each, with no way to parallelize across the output positions.