Beam Search

The problem

Greedy decoding picks the most likely token at every step. But the most likely sequence may start with a token that is not the most likely one. Beam search keeps the beam_width best partial sequences (the beam), ranked by their total log-probability, so it can recover from a first choice that looked good and turned out badly.

The algorithm

next_log_probs(seq) is the model: given a list of token ids, it returns an array with the log-probability of every possible next token.

• Start with one sequence, [bos] (the beginning-of-sequence token), with score 0.
• At each step, extend every sequence in the beam with every token: the candidate's score is the sequence's score plus that token's log-probability (adding log-probabilities is multiplying probabilities). Keep the beam_width best candidates.
• A kept candidate that ends with eos (the end-of-sequence token) is finished: move it to a list of finished sequences, and do not extend it further. The beam continues with the unfinished ones only.
• Stop after max_len steps, or when the beam is empty. Sequences still unfinished at the end count as finished.
• Return (best_sequence, best_score) among all finished sequences, without bos.

Example. Tokens A, B, C and eos; two steps.

• After bos: A has probability $0.6$ and B has $0.4$.
• After A: C has $0.4$, A has $0.3$, eos has $0.3$.
• After B: eos has $0.9$, C has $0.1$.

Greedy (beam_width = 1): A (0.6), then C (0.4): the sequence A C with probability $0.24$.

Beam of 2:

1. Step 1: both A (0.6) and B (0.4) are kept.
2. Step 2, all candidates: A C (0.24), A A (0.18), A eos (0.18), B eos (0.36), B C (0.04). The two best are B eos (0.36) and A C (0.24). B eos ends with eos: it is finished.
3. End of the steps: A C is unfinished and counts as finished. The best is B eos, with $0.36$, better than greedy's $0.24$. In log terms, $\log 0.36 = -1.02$.

The function returns ([B, eos], -1.02): greedy would have thrown B away at step 1.

What to implement

beam_search(next_log_probs, bos, eos, beam_width, max_len): return (best_sequence, best_score).

Requirements

• Sort candidates by score, highest first. On ties keep the order in which they were generated (Python's sorted is stable): beam order first, then token id.
• beam_width = 1 is greedy decoding.

Hint

Store the beam as a list of (score, seq) pairs, where seq is a list of token ids.