Greedy and beam search decoding for a language model

The question comes in two parts. Part 1 is greedy decoding over a fixed score matrix. Part 2 is beam search, where the scores at each step depend on the tokens already chosen. Most interviewers start with Part 1 as a warm-up, then ask "why is greedy not enough?" and move to Part 2.

Part 1: greedy decoding. Every time step is independent, so the answer is an argmax per row. The only detail to get right is the tie rule: scan left to right and replace the current best only on a strictly greater score, so the smallest index wins ties. row.index(max(row)) also gives the smallest index, but it scans each row twice.

Why greedy is not enough. In a real model the distribution at step t depends on the prefix. A token that looks best now can lead to a bad continuation. Example: step 1 gives A: 0.6, B: 0.4. After A the next token is a coin flip (0.5, 0.5); after B one token has 0.9. Greedy takes A and ends at 0.6 * 0.5 = 0.30. The path B, x scores 0.4 * 0.9 = 0.36.

Brute force. Score every sequence of length T: V^T sequences. Exact, but unusable past tiny inputs.

Part 2: beam search. Keep the k best prefixes (the beam), ranked by the sum of log-probabilities. At each step:

1. Expand every alive prefix with every token: score + logp(token | prefix).
2. Keep the top k candidates (higher score first; on a tie, the lexicographically smaller sequence, so output is deterministic).
3. A candidate that ends in eos moves to a finished pool and stops expanding. The rest form the next beam.
4. Stop at max_len, when no alive beams remain, or early when k beams are finished and every alive beam already scores below the best finished one.

Return the best sequence across finished plus the surviving beams.

Invariants and reasons.

• Sum log-probabilities, never multiply probabilities: products underflow to 0.0 after a few dozen tokens.
• Log-probabilities are <= 0, so extending a prefix can never raise its score. That is what makes the early stop safe. If the scores are arbitrary (logits, rewards), drop the early stop.
• k = 1 is exactly greedy decoding. Larger k searches more of the tree but is still a heuristic: a prefix pruned at step t is gone for good.
• Use a top-k selection (heapq.nsmallest) instead of sorting all k * V candidates.

import heapq
import math
from typing import Callable, Dict, List, Optional, Sequence, Tuple

Seq = Tuple[int, ...]
StepFn = Callable[[Seq], Sequence[float]]


def greedy_decode(scores: List[List[float]]) -> List[int]:
    """Part 1: pick the best token at each step; ties go to the smaller index."""
    result: List[int] = []
    for row in scores:
        if not row:
            raise ValueError("empty score row")
        best = 0
        for j in range(1, len(row)):
            if row[j] > row[best]:  # strict >, so the first (smallest) index wins ties
                best = j
        result.append(best)
    return result


def _better(a: Tuple[float, Seq], b: Tuple[float, Seq]) -> bool:
    """Higher score wins; on a tie, the lexicographically smaller sequence wins."""
    return a[0] > b[0] or (a[0] == b[0] and a[1] < b[1])


def beam_search(
    step_fn: StepFn,
    beam_width: int,
    max_len: int,
    eos: Optional[int] = None,
) -> Tuple[List[int], float]:
    """Part 2: keep the beam_width best prefixes by summed log-probability.

    step_fn(prefix) returns one log-probability per vocabulary token for the
    next position. A beam that emits eos is finished and stops expanding.
    """
    if beam_width < 1 or max_len < 0:
        raise ValueError("beam_width must be >= 1 and max_len >= 0")
    beams: List[Tuple[float, Seq]] = [(0.0, ())]
    finished: List[Tuple[float, Seq]] = []

    for _ in range(max_len):
        candidates: List[Tuple[float, Seq]] = []
        for score, prefix in beams:
            for tok, lp in enumerate(step_fn(prefix)):
                if lp == -math.inf:
                    continue
                candidates.append((score + lp, prefix + (tok,)))
        if not candidates:
            break
        # Top-k by (score desc, sequence asc): O(C log k) instead of a full sort.
        top = heapq.nsmallest(beam_width, candidates, key=lambda c: (-c[0], c[1]))
        beams = []
        for cand in top:
            if eos is not None and cand[1][-1] == eos:
                finished.append(cand)
            else:
                beams.append(cand)
        if not beams:
            break
        # Prune: log-probs are <= 0, so an alive beam can only get worse.
        if len(finished) >= beam_width:
            best_done = max(f[0] for f in finished)
            if all(b[0] < best_done for b in beams):
                break

    pool = finished + beams
    best = pool[0]
    for cand in pool[1:]:
        if _better(cand, best):
            best = cand
    return list(best[1]), best[0]


def table_step_fn(table: Dict[Seq, List[float]], vocab: int) -> StepFn:
    """Helper for tests: unknown prefixes get uniform log-probabilities."""
    uniform = [math.log(1 / vocab)] * vocab
    return lambda prefix: table.get(prefix, uniform)


if __name__ == "__main__":
    assert greedy_decode([[0.1, 0.7, 0.2], [0.5, 0.5, 0.0]]) == [1, 0]
    L = math.log
    fn = table_step_fn({(): [L(0.6), L(0.4)], (0,): [L(0.5), L(0.5)],
                        (1,): [L(0.9), L(0.1)]}, 2)
    assert beam_search(fn, 1, 2)[0] == [0, 0]  # greedy: 0.30
    assert beam_search(fn, 2, 2)[0] == [1, 0]  # beam finds 0.36
    print("ok")

Let T be the number of steps, V the vocabulary size, k the beam width.

• greedy_decode: O(T * V) time, O(T) space for the output.
• beam_search: per step, k * V candidates, top-k selection in O(k * V * log k), plus building each new prefix tuple in O(T). Total O(T * k * V * (log k + T)) time with tuple copies, or O(T * k * V * log k) with back-pointers (see follow-ups). Space O(k * V) for the candidates of one step, plus O(k * T) for the beams and the finished pool. It also makes O(T * k) calls to step_fn, which dominate in a real model.

• Greedy, normal: [[0.1, 0.7, 0.2], [0.5, 0.5, 0.0], [-1.0, -3.0, -0.5]] -> [1, 0, 2] (the tie in row 2 goes to index 0; negative scores work).
• Greedy, empty input: [] -> [].
• Greedy, one token: [[4.0]] -> [0].
• Beam, greedy trap: step 1 log[0.6, 0.4]; after 0 log[0.5, 0.5]; after 1 log[0.9, 0.1]; max_len = 2. With k = 1 -> [0, 0], score log 0.30. With k = 2 -> [1, 0], score log 0.36.
• Beam, EOS: vocabulary {0, 1, 2}, eos = 2; step 1 log[0.3, 0.2, 0.5]; after 0 log[0.1, 0.1, 0.8]; k = 2, max_len = 5 -> [2], score log 0.5 (the finished beam beats [0, 2] at log 0.24, and the search stops early).
• Beam, all ties: step_fn always returns [0.0, 0.0], k = 3, max_len = 3 -> [0, 0, 0], score 0.0 (lexicographic tie-break).
• Beam, max_len = 0 -> ([], 0.0).
• Beam, beam_width = 0 -> ValueError.

1. Length normalization. Summed log-probs favor short outputs, since every token adds a negative term. Rank finished beams by score / len ** alpha (GNMT uses ((5 + len) / 6) ** alpha). The pruning comparison must use the same normalized score, and the early stop is no longer exact.
2. Return the top n sequences, not one. Keep the finished pool as a min-heap of size n and return it sorted. Also usual: stop only when n beams have finished.
3. Batched / tensor version. With NumPy or PyTorch: scores (k,) plus logprobs (k, V), flatten to k * V, take topk, recover beam = idx // V, token = idx % V. Store a back-pointer per step instead of copying prefixes, and rebuild the best sequence at the end. This removes the O(T) tuple copy per candidate.
4. Constraints on output. No repeated n-grams (mask tokens whose n-gram already appears in the prefix by setting their log-prob to -inf), a minimum length (mask eos until then), or a forced vocabulary. All are masks applied before the top-k step; the code already skips -inf.
5. Sampling alternatives. Top-k sampling, top-p (nucleus) sampling and temperature. Explain when to use each: beam search suits translation and summarization; open-ended chat uses sampling, because beam output tends to be repetitive and generic.

Mid-level (L4)

• Greedy argmax with the correct tie rule, written quickly.
• Can explain with a small example why greedy misses a better sequence.
• A working beam search with a fixed width and a fixed length, using summed log-probs.

Senior (L5)

• Handles EOS and finished beams correctly, and does not keep expanding them.
• Explains log space versus underflow, and uses top-k selection instead of a full sort.
• Gives the complexity in T, k and V, and names the model call as the real cost.
• Mentions the short-sequence bias and length normalization without prompting.

Staff (L6+)

• Explains when the early stop is safe (non-positive increments) and when it is not.
• Describes the batched tensor version with back-pointers and KV-cache reuse per beam (reorder the cache by the surviving beam indices).
• Discusses the trade-off between beam search and sampling for different products, and that larger beams can lower output quality (the beam search curse).