Greedy and Top-k Beam Search Decoding That Returns the Sentence Probability

A sequence model generates a sentence one token at a time. You are given its next-step probability distribution: for any prefix of tokens generated so far, it returns a probability for every candidate next token. Write a decoder that turns this distribution into an output sentence. First implement greedy decoding, then extend it to beam search that keeps the top `k` paths at each step. Both versions must also return the probability of the whole sentence. ### Constraints and Clarifications - Assume the distribution is exposed as a function `next_probs(prefix)` that takes a tuple of tokens and returns a dictionary from token to probability, with the probabilities of one call summing to 1. - Assume the vocabulary contains an end-of-sentence token `eos` that the model can emit, and that decoding also stops after `max_len` tokens. - The probability of a sentence is the product of the probabilities of its tokens, each conditioned on the tokens before it. ### Clarifying Questions - Is the distribution conditioned on the whole prefix, or is it a fixed table with one distribution per position that ignores earlier choices? - Does the probability of emitting `eos` count toward the sentence probability, and should `eos` appear in the returned sentence? - How should ties between equally probable tokens or paths be broken? - Should the result report the raw probability or its logarithm? - What should be returned if `max_len` is reached before `eos`? ### Part 1 — Greedy decoding Implement `greedy_decode(next_probs, eos, max_len)`, which at each step appends the single most probable next token and returns `(tokens, probability)`. State its time complexity in terms of the sentence length and the vocabulary size. ```hint What to carry forward Decide what running quantity lets you return the sentence probability at the end without recomputing it, and what happens to that quantity over many steps. ``` #### What This Part Should Cover - Correct stopping on `eos` or at `max_len` - A sentence probability maintained step by step that stays numerically usable - Deterministic tie handling and the cost per step ### Part 2 — Beam search with the top k paths Extend the decoder to `beam_search(next_probs, eos, max_len, k)`, which keeps the `k` most probable partial sentences at every step instead of one, and returns the best sentence together with its probability. Explain what each beam entry stores, what happens to a path that has already emitted `eos`, when the search stops, and the time and memory complexity. ```hint One step of the beam Count how many candidates a single step produces from `k` paths, then decide which of them survive and on what score. ``` #### What This Part Should Cover - Expansion of every live path and pruning to the top `k` across all of them - Handling of finished paths and a sound stopping rule - How the result relates to greedy decoding and to an exhaustive search - Time, memory and model-call cost in terms of `k`, vocabulary size and length ### Part 3 — Low probabilities for long sentences Without writing code, discuss a problem with the sentence probability you return: long sentences tend to receive low probability. Explain why this happens, how it affects which sentence beam search returns, and what you would change in the scoring or the search to address it. ```hint Compare two lengths Take a short sentence and a much longer one whose individual tokens are each more likely, and compare how the raw product ranks them. ``` #### What This Part Should Cover - The mechanism that makes the raw score shrink with length, and where in the search it matters - At least one correction to the scoring, and how it interacts with pruning and stopping - The difference between numerical underflow and a bias in which sentence is chosen - How a correction is tuned, and how it can overcorrect ### What a Strong Answer Covers - Working code for both decoders with shared conventions for ties, stopping and scoring - Probability bookkeeping that is exact and does not underflow on long outputs - A clear statement that beam search is a heuristic, with its cost and quality trade-off as `k` grows - Edge cases: zero-probability tokens, `max_len` reached without `eos`, and `k` larger than the number of candidates - A beam implementation written fast enough to leave time for the length discussion ### Follow-up Questions - How would you batch the model calls for the `k` paths at each step, and what changes when `next_probs` is an expensive neural network call? - How would you return the `n` best distinct sentences instead of one, and keep the beam from filling with near-duplicates? - When would sampling (top-k or nucleus sampling) be a better choice than beam search?