Beam Search Decoding - LLM Generation & Decoding

Beam Search Decoding

What you'll build

Beam search over a table of per-step log-probabilities: keep the beam_width best partial sequences at every step, with optional end-of-sequence handling, and return the best sequence and its score.

Background

Greedy decoding commits to the single best token at each step and can miss a sequence whose first token looked worse. Beam search keeps several hypotheses alive. It is used where there is one "right" output (translation, speech recognition, constrained or structured generation); open-ended chat uses sampling instead, because beam search tends to produce bland, repetitive text.

In a real model, the next-token distribution depends on the prefix. To keep this problem self-contained, log_probs[t] is the same for every beam at step t. Without EOS that makes greedy already optimal, so the interesting case is EOS: a beam that ends early stops paying log-probability costs.

Algorithm

Scores are sums of log-probabilities (never multiply raw probabilities: they underflow).

1. Start with one empty beam: [([], 0.0)].
2. For each step t in 0 .. min(max_length, num_steps) - 1:
   • For each beam, in order: if it ends with eos_token, carry it to the candidate list unchanged; otherwise add one candidate per token v = 0 .. vocab_size - 1: (seq + [v], score + log_probs[t, v]).
   • Sort all candidates by score, descending, with a stable sort (ties keep candidate order: earlier beam first, then lower token id), and keep the first beam_width.
3. Return the first (best) beam.

Inputs

• log_probs: float array (num_steps, vocab_size) of log-probabilities
• beam_width: int >= 1
• max_length: int or None (None means num_steps)
• eos_token: int or None (None disables EOS handling)

Expected Return Format

A tuple (best_sequence, best_score):

• best_sequence: list of Python ints
• best_score: total log-probability of that sequence, rounded to 4 decimals

Edge cases and conventions

• Finished (EOS) beams compete with unfinished ones on their raw score: no length normalization here (see the follow-up).
• beam_width = 1 is greedy decoding.

Constraints

NumPy and plain Python lists.

Follow-up

• Raw log-probability scores favour short sequences. Add length normalization score / len(seq)^alpha (alpha about 0.6 to 0.7, Wu et al., 2016). Where in the algorithm must it be applied?
• Complexity per step is O(B · V) candidates; how do real implementations avoid sorting all of them (top-k per beam first)?

Input: beam_search(np.log(np.array([[0.6, 0.3, 0.1], [0.2, 0.7, 0.1], [0.1, 0.3, 0.6]])), beam_width=1)
Output: ([0, 1, 2], -1.3783)

Input: beam_search(np.log(np.array([[0.5, 0.45, 0.05], [0.3, 0.3, 0.4], [0.3, 0.3, 0.4]])), beam_width=2, eos_token=1)
Output: ([1], -0.7985)

import numpy as np

def beam_search(log_probs, beam_width, max_length=None, eos_token=None):
    """
    Beam search decoding over pre-computed log probabilities.

    Args:
        log_probs: Log probability matrix (num_steps, vocab_size);
                   log_probs[t] gives the log probs for step t
        beam_width: Number of beams to keep
        max_length: Maximum sequence length (defaults to num_steps)
        eos_token: End-of-sequence token id (optional); a beam ending with it
                   is carried forward unchanged

    Returns:
        best_sequence: List of token ids (ints) for the best beam
        best_score: Its total log-probability, rounded to 4 decimals
    """
    # Your code here
    pass