Next-Token Decoding: Greedy and Beam Search

Context

You are provided with a probabilistic next-token dictionary D. This dictionary maps each token t to a collection of candidate next tokens and their conditional probabilities. Specifically, D[t] is a dictionary where keys are candidate next tokens u, and values are their conditional probabilities p(u | t). Assume that for any given t, the probabilities p(u | t) for all u in D[t] sum to 1.

The special token '<START>' signifies the beginning of a sequence, and '.' denotes the end of a sequence. If a token does not exist as a key in D, it implies that there are no continuations for that token (a dead-end).

Tasks

1. Greedy Decoding

Implement a greedy decoding strategy:

• Begin from the '<START>' token.
• Repeatedly select the next token with the highest conditional probability from D.
• Continue this process until either the '.' (end-of-sequence) token is produced, or a dead-end is encountered (no continuations exist for the current token).
• Provide both a recursive and an iterative implementation of this algorithm.
• The function should return the sequence of tokens generated, starting from the token immediately after '<START>' and including '.' if it was produced.
• Analyze the time and space complexity of your implementations.

2. Beam Search Decoding (with beam width k)

Implement a beam search decoding strategy:

• Employ a Breadth-First Search (BFS)-style expansion.
• At each decoding step, maintain a set of up to k highest-scoring partial sequences. These sequences should be ranked by their cumulative log-probability.
• Expand each active partial sequence (beam) by considering its top-k next tokens from D.
• Continue this process until all active beams have terminated with the '.' token.
• The function should return the single best sequence found, and optionally, the top-k sequences.
• Explain how your implementation handles the following specific scenarios:
  • Ties in scores between candidate sequences.
  • Tokens that are missing from D (dead-ends).
  • The potential for generating cyclical sequences.
  • Score normalization techniques (e.g., length normalization).
  • The criteria for terminating the search.
• Compare greedy search versus beam search in terms of:
  • Quality of generated sequences.
  • Computational complexity.
  • Typical use cases where each approach is preferred.
• Demonstrate your implementations using a small, illustrative example dictionary D.