Beam search decoding, contrasted against greedy

The problem, from first principles

[07-greedy-decoding]'s core weakness: choosing the single best-looking token at EVERY step independently doesn't always add up to the best OVERALL sequence. A token that looks slightly less likely right now might open the door to a MUCH more likely continuation later, a tradeoff greedy decoding can never see, since it commits irreversibly at every single step and never looks back. Beam search fixes this by keeping several candidate sequences (beam_width-many "beams") alive SIMULTANEOUSLY: at every step, every surviving beam gets expanded by its own top candidates, all of those expansions get pooled together, and only the globally best beam_width (ranked by CUMULATIVE log-probability across the whole sequence so far, not just the newest token) survive into the next round, before expanding again.

This genuinely widens the search: a sequence that looked second-best after one step, but whose continuation turns out to be excellent, gets a real chance to survive and eventually WIN, something [07-greedy-decoding]'s single irreversible choice could never recover from. beam_width=1 collapses this exactly back down to greedy decoding (only one beam ever survives, and its top-1 expansion IS the argmax), giving greedy decoding a precise, mathematically clean special case of the more general algorithm.

From theory to code

Implement sequence_log_prob(token_ids, ...) (the total log-probability the model assigns to an already-complete sequence, summed over every next-token prediction it makes), and beam_search_decode(token_ids, ..., num_new_tokens, beam_width), maintaining beam_width candidate sequences and their cumulative log-probabilities across num_new_tokens generation steps.

Constraints

• Scoring uses SUMMED log-probabilities (log-probabilities are additive across independent steps; raw probabilities are NOT, since multiplying many probabilities together underflows toward 0 for long sequences).
• At every step, EVERY current beam gets expanded by its own top beam_width candidates FIRST, and only THEN are all beams' expansions pooled together and globally re-ranked, never ranking within one beam's expansions in isolation.
• beam_width=1 must produce EXACTLY [07-greedy-decoding]'s output: with only one beam ever alive, its top-1 expansion at every step is, by definition, the single highest-scoring next token, the same argmax greedy decoding uses.
• The final answer is the single highest CUMULATIVE-log-probability sequence among the surviving beams after all num_new_tokens steps.

Hints

candidates = []
for sequence, score in beams:
    log_probs = log_softmax(next_token_logits_for(sequence))
    top_indices = np.argsort(log_probs)[-beam_width:]
    for idx in top_indices:
        candidates.append((sequence + [idx], score + log_probs[idx]))
candidates.sort(key=lambda c: c[1], reverse=True)
beams = candidates[:beam_width]