Katkat
Back

Implement Autoregressive Decoding Strategies

CohereMachine Learning Engineer · Onsite
MediumMachine Learning EngineerOnsiteCoding & Algorithms

Last updated

This problem tests your ability to implement common sequence decoding strategies used in large language models by manipulating token probability distributions.

Problem Statement

In autoregressive models, decoding involves iteratively predicting the next token in a sequence given the previous context. You are provided with a 1D array logits, where logits[i] represents the raw score for vocabulary index i.

Implement a function decode_next_token(logits, method, k=None, p=None) that selects the next token index based on the chosen decoding strategy:

  1. Greedy: Select the index with the highest logit.
  2. Top-K: Mask all but the k highest-scoring tokens, re-normalize their probabilities, and sample from the resulting distribution.
  3. Top-P (Nucleus): Sort tokens by probability in descending order. Select the smallest set of tokens whose cumulative probability exceeds or equals p. Sample from this subset.

Function Signature

def decode_next_token(logits, method, k=None, p=None):

Constraints

  • logits: List of floats, size N (vocab size).
  • method: String ("greedy", "top_k", "top_p").
  • Assume a standard softmax function exp(x) / sum(exp(x)) is available.
  • For sampling, use random.choices with the calculated probability distribution.

Examples

Example 1: Greedy Input: logits = [0.1, 2.5, 0.8], method = "greedy" Output: 1 (Index 1 has the highest logit value).

Example 2: Top-K Input: logits = [1.0, 2.0, 3.0, 4.0], method = "top_k", k=2 Calculation: The top 2 indices are 2 and 3 (logits 3.0 and 4.0). Convert these to probabilities and sample between them.

Edge Cases

  • What if multiple tokens have the same highest logit?
  • How do you handle top_p when the highest probability token alone already exceeds p?

Hints

Hint 1

To handle sampling, first apply the softmax function to convert logits into a probability distribution summing to 1.0.

Hint 2

For Top-K, use `np.argsort` or equivalent to find indices of the largest values before applying softmax to the subset.

Hint 3

For Top-P, sort the probabilities descending and use `itertools.accumulate` to find the cutoff point where the cumulative sum crosses the threshold `p`.

Follow-up questions

  • How would you optimize the Top-K selection if the vocabulary size is in the millions?
  • Temperature scaling is often used to sharpen or flatten the distribution before sampling; where would you insert that in the calculation?
  • Explain the trade-offs between Top-K and Top-P sampling in terms of output diversity versus coherence.