Top-K Sampling

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 $\frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}$, where $x_i$ is the logit for the $i^{th}$ 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 $-\infty$. This process effectively eliminates the less likely tokens from consideration, as their corresponding probabilities will become 0 after applying the softmax function.