Beam Search From Scratch in Python, With Tie Rules

Tokens are 0 to V − 1. start[v] is the log-probability that a sequence begins with v, and trans[a][b] the log-probability that b follows a. A sequence's score is the sum of its start term and its transition terms. Run beam search of width B for sequences of length T:

Keep the B best one-token sequences.
Repeat T − 1 times: extend every kept sequence by every token and keep the B best extensions.

"Best" means highest score; among equal scores, the lexicographically smaller sequence wins. Return [sequence, score] for the best sequence after the last step.

Examples

Input start = [-0.1, -2.5], trans = [[-2.5, -0.1], [-0.1, -2.5]], T = 3, B = 1

Output [[0, 1, 0], -0.3]

Width 1 is greedy decoding: 0, then 1, then 0.

Input start = [-0.7, -0.7], trans = [[-0.7, -0.7], [-0.7, -0.7]], T = 2, B = 2

Output [[0, 0], -1.4]

Every sequence ties; the lexicographic rule picks [0, 0].

Constraints

1 ≤ V ≤ 20, 1 ≤ T ≤ 8, 1 ≤ B ≤ 20
Checked to a relative tolerance of 10⁻¹²

Hints

Hint 1

Represent a beam as a pair (score, sequence). The score of an extension is the beam's score plus one transition term.

Hint 2

Sort the candidates by (−score, sequence) and keep the first B: that implements both the ranking and the tie rule.

Hint 3

With B = 1 you should get greedy decoding; with B at least V^(T−1) you should get the exact best sequence. Both make good tests.

Where candidates go wrong

Multiplying probabilities instead of adding log-probabilities: underflow for long sequences.
Keeping the best B extensions of each beam separately, which keeps B² sequences.
Breaking ties by the order candidates were generated, which depends on the implementation.

The idea

The frontier is a list of at most B (score, sequence) pairs. Each step builds every one-token extension (B · V candidates), sorts them by score descending and sequence ascending, and truncates to B. The work per step is O(BV log(BV)). Beam search is not exact: the best full sequence can start with a token that was pruned at step one, which is the follow-up interviewers ask about.

Time O(T · BV log(BV)) · Space O(BV + BT)

Reference solution

import numpy as np

def beam_search(start, trans, T, B):
    start = np.asarray(start, dtype=float)
    trans = np.asarray(trans, dtype=float)
    V = start.shape[0]
    beams = sorted(((float(start[v]), (v,)) for v in range(V)), key=lambda t: (-t[0], t[1]))[:B]
    for _ in range(T - 1):
        cand = []
        for score, seq in beams:
            ext = score + trans[seq[-1]]                 # scores of all V extensions at once
            cand.extend((float(ext[v]), seq + (v,)) for v in range(V))
        beams = sorted(cand, key=lambda t: (-t[0], t[1]))[:B]
    score, seq = beams[0]
    return [list(seq), score]

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

Every text generator decodes with some version of this. Interviewers use it to test bookkeeping under pressure: scores in log space, a frontier that is pruned each step, and a tie rule that makes the output reproducible.

The theory behind it is in the lesson LSTM and GRU cells, and beam search, from scratch.

Keep going

Medium15 min

Beam search decoding