Top-K Sampling
Implement top-K sampling for text generation.
Top-K sampling restricts the next token to the K most likely tokens:
- Sort logits descending
- Keep only the top K
- Set the rest to -infinity
- Apply softmax to the remaining logits
Input:
- Line 1: K
- Line 2: space-separated logits
Output: The filtered probability distribution (all positions), rounded to 4 decimal places. Non-top-K positions should be 0.0000.
Example:
2 3.0 1.0 4.0 2.0
0.2689 0.0000 0.7311 0.0000
- First, we sort the logits in descending order: 4.0,3.0,2.0,1.0
- Then, we keep only the top K (K=2) logits and set the rest to −∞: 4.0,3.0,−∞,−∞
- Next, we apply the softmax function to the remaining logits: P(i)=∑j=1K​elogitj​elogiti​​, resulting in e4.0+e3.0e4.0​ and e4.0+e3.0e3.0​
- The final output is the filtered probability distribution, with non-top-K positions set to 0.0000: 0.2689,0.0000,0.7311,0.0000
Constraints:
- 1 <= K <= vocab size
- Positions not in top-K get probability 0
- Round to 4 decimal places
More from LLM 3: Applications & Evaluation
Background Knowledge
The problem of Top-K Sampling is a technique used in text generation tasks, particularly in the context of language models. It is a method to restrict the next token in a sequence to the most likely K tokens, based on the model's predictions. This is useful for controlling the output of the model, by limiting the possibilities to a subset of the most probable tokens. The core concept here involves understanding logits, which are the raw, unnormalized scores that a model outputs for each possible token. These logits are then typically passed through a softmax function to obtain a probability distribution over all possible tokens.
The softmax function is a crucial component in this process. It takes the logits as input and outputs a probability distribution, where each probability is proportional to the exponential of the corresponding logit. The softmax function is defined as ∑j=1n​exj​exi​​, where xi​ is the logit for the ith token, and n is the total number of tokens. This function ensures that the output probabilities are non-negative and sum up to 1, making them a valid probability distribution.
Understanding the Top-K Sampling technique also requires knowledge of how to manipulate and transform the logits to achieve the desired restriction. This involves sorting the logits in descending order, selecting the top K logits, and then setting the rest to −∞. This process effectively eliminates the less likely tokens from consideration, as their corresponding probabilities will become 0 after applying the softmax function.
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.