Beam Search Decoder

Start from brute force

The obvious approach is exhaustive search: enumerate every one of the V^steps sequences that begin at start (V is the number of tokens), multiply the transition probabilities along each, and keep the best, breaking ties lexicographically. It works because it looks at everything, so it always returns the true optimum, [0, 2, 1] at 0.40 on the first example. It costs O(V^steps · steps) time and O(steps) space if you generate one sequence at a time. That is exponential: 20 steps over a 50,000-token vocabulary is more sequences than any machine can score. It is also not this task: with k = 1 the tests expect the trapped greedy answer [0, 1, 1], not the optimum.

The approach

Beam search keeps only the k best prefixes and grows them one token at a time. The insight is that a score is a running product: extending a prefix multiplies its score by one transition, so the candidates for the next step are built from the kept prefixes alone and everything outside the beam is discarded for good. That is the trade: it is fast, and it can miss the optimum when a weak prefix leads to a strong ending. The invariant: after each step, beams holds the k highest-scoring extensions of the previous beam, best first.

The reference starts with beams = [([start], 1.0)]. Each step it walks every (seq, score), reads last = seq[-1], and appends (seq + [j], score * trans[last][j]) to candidates for every token j, zero-probability ones included. One sort with the key (-c[1], c[0]) ranks by score descending and breaks ties by the smaller sequence, and beams = candidates[:k] prunes across the whole pool, not k per beam. After steps rounds it returns beams[0][0]. With k = 1 this is exactly greedy decoding.

Walk through an example

The first example’s matrix traced with k = 1 and k = 2: each step expands every kept prefix, then one sort over the whole pool keeps the top k.

Complexity

Beam search: each of the steps rounds builds at most kV candidates and sorts them; copying a prefix and comparing two tied prefixes each cost up to steps + 1 tokens, so time is O(steps² · kV · log kV) at worst and space is O(kV · steps) for the pool. Brute force: O(V^steps · steps) time and O(steps) space, because it scores every sequence one at a time.

Edge cases to test

• steps = 0: no expansion runs, so the answer is [start] alone, e.g. [1] for start = 1.
• k at least the number of candidates, e.g. k = 9 on the first example: the slice keeps everything, nothing is pruned, and the result is the exhaustive optimum [0, 2, 1].
• Uniform rows, trans = [[0.5, 0.5], [0.5, 0.5]], steps = 2, k = 2: every sequence scores 0.25, so only the tie-break decides, and the answer is [0, 0, 0].
• Ties across beams, trans = [[0.25, 0.25, 0.5], [0, 0, 1.0], [0.5, 0.5, 0]], steps = 2, k = 3: the answer is [0, 1, 2]. A sort on score alone keeps beam order and returns [0, 2, 0], which also scores 0.25 but is the larger sequence.
• A dead-end token whose row is all zeros, trans = [[0, 1], [0, 0]], steps = 2, k = 1: every extension of [0, 1] scores 0 and the answer is [0, 1, 0]. Code that drops zero-probability candidates empties the beam and crashes on beams[0].

Common mistakes

• Keeping the top-k per beam instead of the top-k across all expansions — that is the crux of beam search.
• Sorting by score alone and trusting insertion order for ties. The pool is built in beam order, which follows score, not sequence order, so a tie across two beams can resolve to the larger sequence.
• Extending with seq.append(j) instead of seq + [j]: every candidate from one beam then shares a single list, and they all end up identical.
• Filtering out zero-probability extensions. The statement keeps them, and a row of zeros then leaves the beam empty.

Follow-ups an interviewer asks next

Long sequences make the product underflow. What changes?

Sum log-probabilities instead of multiplying probabilities. The logarithm is strictly increasing, so the ranking is unchanged, but a sum of a few hundred negative numbers stays representable where the product of a few hundred small probabilities rounds to 0.0 and makes every candidate tie. Map a zero probability to negative infinity yourself, because Python’s math.log(0) raises.

Sequences can stop early at an end-of-sequence token.

Move a hypothesis that emits the end token out of the beam into a finished list, keep expanding the rest, and stop once k hypotheses have finished or a length cap is hit. Compare finished hypotheses with a length-normalized score, such as total log-probability divided by length raised to a power alpha, because every extra token adds a negative term and raw scores favor short outputs.

The vocabulary has 50,000 tokens. Do you really sort k · V candidates?

No. A candidate can only make the global top k if it is in its own beam’s top k, so take each beam’s top k tokens with a heap (O(V log k) per beam), then pick the best k of those k² candidates. On a GPU the same step is one top-k over the flattened k × V score matrix.

The k beams all end up nearly identical.

That is the usual failure of plain beam search: one strong prefix fills the beam with its own variations. Diverse beam search splits the beam into groups and penalizes a group for picking tokens that earlier groups already picked at that step. A simpler fix caps how many kept hypotheses may share a prefix.

Where this shows up

Beam search is a core decoding strategy for machine translation, speech recognition, and constrained generation — the middle ground between greedy decoding (k = 1) and infeasible exhaustive search. The “keep the top-k hypotheses and prune the rest” idea recurs in any large pruned search problem.

Reference solution

def beam_search(start, trans, steps, k):
    beams = [([start], 1.0)]
    for _ in range(steps):
        candidates = []
        for seq, score in beams:
            last = seq[-1]
            for j in range(len(trans[last])):
                candidates.append((seq + [j], score * trans[last][j]))
        candidates.sort(key=lambda c: (-c[1], c[0]))
        beams = candidates[:k]
    return beams[0][0]


TRAP = [[0, 0.6, 0.4], [0, 0.5, 0.5], [0, 1.0, 0]]
print(beam_search(0, TRAP, 2, 1))   # greedy falls in the trap
print(beam_search(0, TRAP, 2, 2))   # the beam escapes

function beamSearch(start: number, trans: number[][], steps: number, k: number): number[] {
  let beams: [number[], number][] = [[[start], 1]];
  for (let s = 0; s < steps; s++) {
    const candidates: [number[], number][] = [];
    for (const [seq, score] of beams) {
      const last = seq[seq.length - 1];
      for (let j = 0; j < trans[last].length; j++) {
        candidates.push([[...seq, j], score * trans[last][j]]);
      }
    }
    candidates.sort((a, b) => {
      if (b[1] !== a[1]) return b[1] - a[1];
      const x = a[0], y = b[0];
      for (let i = 0; i < x.length; i++) if (x[i] !== y[i]) return x[i] - y[i];
      return 0;
    });
    beams = candidates.slice(0, k);
  }
  return beams[0][0];
}

const TRAP = [[0, 0.6, 0.4], [0, 0.5, 0.5], [0, 1.0, 0]];
console.log(beamSearch(0, TRAP, 2, 1));   // greedy falls in the trap
console.log(beamSearch(0, TRAP, 2, 2));   // the beam escapes