Beam Search Decoder

Implement beam search decoding for text generation.

Given a vocabulary and a function that returns log-probabilities for the next token given a sequence, perform beam search:

For simplicity, log-probabilities are provided as a matrix: logprobs[step][token_id].

Input:

Output: The best sequence (space-separated token IDs, excluding EOS) and its score rounded to 4 decimal places.

Input

3 2 3 2
-0.5 -1.0 -2.0
-0.8 -0.3 -3.0
-1.0 -0.5 -0.1

Output

0 1
-0.8000

Reasoning

We start with an empty sequence and expand it by all vocab tokens (0, 1, 2) with their corresponding log-probabilities: $-0.5$ , $-1.0$ , $-2.0$ .
We keep the top 2 candidates (by total log-probability): tokens 0 and 1 with scores $-0.5$ and $-1.0$ .
At the next step, we expand each beam by all vocab tokens and calculate their total log-probabilities, e.g., for beam [0], we have $-0.5 + (-0.8) = -1.3$ , $-0.5 + (-0.3) = -0.8$ , $-0.5 + (-3.0) = -3.5$ .
We keep the top 2 candidates: beams [0, 1] with a total log-probability of $-0.8$ and select this as the best sequence since further expansion will not improve the score due to the $-0.8$ being the highest among all possible expansions.

Background Knowledge