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:
- Start with an empty sequence
- At each step, expand each beam by all vocab tokens
- Keep the top beam_width candidates (by total log-probability)
- Stop after max_length steps or when all beams end with EOS
For simplicity, log-probabilities are provided as a matrix: logprobs[step][token_id].
Input:
- Line 1: vocab_size beam_width max_length eos_id
- Next max_length lines: log-probability rows (vocab_size floats each)
Output: The best sequence (space-separated token IDs, excluding EOS) and its score rounded to 4 decimal places.
Example:
3 2 3 2 -0.5 -1.0 -2.0 -0.8 -0.3 -3.0 -1.0 -0.5 -0.1
0 1 -0.8000
- 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.
Constraints:
- 1 <= beam_width <= 10, 1 <= max_length <= 10
- 0 <= eos_id < vocab_size
- If beam hits EOS, it's a complete beam (don't extend further)
- Pick the best complete beam. If none, pick best incomplete beam
More from LLM 3: Applications & Evaluation
Background Knowledge
Text Generation and Decoding are crucial components of many natural language processing (NLP) tasks, including machine translation, text summarization, and chatbots. In these tasks, a model is trained to predict the next token in a sequence, given the context of the previous tokens. Beam Search is a popular decoding algorithm used to generate text by exploring the most promising candidates in a vast search space. It is particularly useful when the search space is too large to be exhaustively searched.
The Beam Search algorithm works by maintaining a set of candidate sequences, known as the beam, and iteratively expanding each sequence by adding the next possible token. The beam width determines the number of candidate sequences to keep at each step. By selecting the top candidates based on their log-probabilities, the algorithm can efficiently explore the search space and find the most likely sequence. The use of log-probabilities instead of probabilities helps to avoid underflow issues when dealing with very small values.
In the context of this problem, the log-probability matrix provides the necessary information to perform beam search decoding. The matrix has a size of n×m, where n is the number of steps and m is the vocabulary size. Each entry in the matrix represents the log-probability of a particular token being the next token in the sequence, given the context of the previous tokens. The EOS (End-of-Sequence) token is used to indicate the end of a sequence.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.