Beam Search

When choosing the first piece of a response, you may reject too early a path that would later turn out better. Checking all possible sequences is too expensive, however. You need a limited pool of alternatives.

Beam Search retains several partial responses with the highest scores. At each step it extends them, compares the new sequences, and again keeps only a specified number. The size of this pool is called the beam width.

With width two, the system can keep beginnings A and B instead of immediately choosing A. At the next step it compares their continuations. The basic score combines the probabilities of successive tokens — pieces of text — although implementations may adjust the effect of length.

A larger pool costs more memory and computation, and can still discard the best path. A high probability for the entire sequence also does not guarantee that the response is true, diverse, or useful.

Mechanism and details

In Autoregressive Language Modeling, the basic sequence score is the sum of the log-probabilities of its tokens:

$$S(x_{1:t})=\sum_{i=1}^{t}\ln p(x_i\mid x_{<i})$$

This is the logarithm of the probability product. A higher score is better; summation avoids multiplying very small numbers. The Hugging Face documentation, “Beam search” describes maintaining multiple sequences. Width 1 reduces this basic variant to Greedy Decoding.

Which prefix will you discard too early?

An original example starts with probabilities A = 0.45, B = 0.35, and C = 0.20. A beam of width 2 discards C. The best continuation of A later has probability 0.4, so the AA result is 0.18. Meanwhile, the omitted CA would achieve $0{,}20\cdot0{,}99=0{,}198$. Width 3 retains C and finds a better result. This is a complete, small tree of two tokens, without early termination or length normalization.

Beam Search remains an approximation: it does not expand a discarded prefix again. Width does not mean the number of tokens sampled from one distribution, as in Top-k Sampling.

For sequences of different lengths, the end token, stopping condition, and scoring function matter. For example, length_penalty in GenerationConfig divides the logarithmic score by the length raised to a specified power. This criterion should not be confused with pure probability. A wider beam does not guarantee better content either; the problem is discussed by Holtzman et al., §2.2–3.

I use AI-generated content as part of my daily learning process.