Implement beam search for a sequence model given a next-token log-probability function.

TL;DR: Beam search holds the k highest-probability partial sequences (beams) at each step rather than greedily grabbing the single best token. Per step, expand every beam by all tokens, score by cumulative log-probability, and retain the top k overall. Stay in log space (add log-probs; multiplying probabilities underflows), set finished sequences aside, and apply length normalization so longer sequences are not unfairly penalized.

How to approach it. Contrast it with greedy decoding (greedy is beam search at k=1 and can miss a globally better sequence), describe the per-step expand-score-prune loop, and flag the three correctness details: log-space scoring, end-of-sequence handling, and length normalization. Then write it.

A strong answer.

import math, heapq

def beam_search(step_logprobs, start, eos, beam_width=3, max_len=20):
    """
    step_logprobs(seq) -> dict {next_token: log_prob} for the given prefix.
    Returns the best completed sequence by length-normalized log-prob.
    """
    beams = [(0.0, [start])]            # (cumulative_logprob, tokens)
    completed = []

    for _ in range(max_len):
        candidates = []
        for score, seq in beams:
            if seq[-1] == eos:          # already finished: carry forward
                completed.append((score, seq))
                continue
            for tok, lp in step_logprobs(seq).items():
                candidates.append((score + lp, seq + [tok]))   # SUM log-probs
        if not candidates:
            break
        # keep the top-k partial hypotheses by cumulative log-prob
        beams = heapq.nlargest(beam_width, candidates, key=lambda x: x[0])

    completed += beams
    # length normalization: divide by length so long seqs aren't penalized
    return max(completed, key=lambda x: x[0] / len(x[1]))[1]

What a strong candidate narrates:

• Log space, summing. Probabilities of a sequence multiply, and multiplying many <1 numbers underflows to 0; summing log-probabilities is numerically stable and monotonic.
• Top-k over all expansions. At each step you expand every beam by every possible next token and keep the best k overall (not k per beam), which is the search.
• End-of-sequence handling. A beam that emits eos is complete; set it aside (do not keep expanding it) and compare against other completed sequences at the end.
• Length normalization. Cumulative log-prob is more negative for longer sequences, so raw scores bias toward short outputs; dividing by length (or a length penalty) corrects that. This is a real bug in naive implementations.
• Greedy is k=1, and beam search is still a heuristic, not guaranteed globally optimal, but explores more than greedy.

How the candidate set collapses back to k each step:

Key takeaways

• Score in log space and sum: products of many sub-1 probabilities underflow to zero.
• Prune to the global top-k across all expansions, not k-per-beam, or you are not actually searching.
• Without length normalization the search collapses to short outputs, since longer sequences carry more negative log-prob.
• Beam search wins on tasks with a correct answer (translation, ASR); for open-ended text it goes bland, so sample instead.

What interviewers probe next.

• "Beam search vs sampling (top-k/top-p)?" Beam search maximizes likelihood and suits tasks with a "correct" output (translation, ASR); for open-ended generation it produces bland, repetitive text, so sampling (temperature, nucleus) is preferred there.
• "Complexity?" Each step expands beam_width × vocab candidates and keeps the top beam_width; time is O(steps × beam_width × vocab) (or × log for the top-k selection).
• "Why does large beam width sometimes hurt quality?" It can over-favor high-likelihood, generic sequences (the "beam search curse"), so bigger is not always better for fluency.
• "How does this interact with the KV cache at inference?" Each beam is a separate sequence with its own cache; you manage k caches (or share prefixes), which is part of why wide beams are expensive to serve.

Common mistakes.

• Multiplying probabilities instead of summing log-probs, underflowing to zero.
• Keeping k per beam instead of the global top-k across all expansions.
• No length normalization, so the search collapses to short sequences.
• Continuing to expand sequences after eos, or never collecting completed hypotheses.