Google

The final step in autoregressive generation is to select the next token by finding the one with the highest probability. This is formally expressed using the argmax function: $$\hat{\mathbf{y}} = \underset{\mathbf{y}}{\arg\max}\, \text{Pr}(\mathbf{y}|\text{cache})$$ Here, $$\hat{\mathbf{y}}$$ is the predicted token, and $$\mathbf{y}$$ is any token in the vocabulary. The notation $$\text{Pr}(\mathbf{y}|\text{cache})$$ is used instead of $$\text{Pr}(\mathbf{y}|\mathbf{x})$$ to explicitly highlight that the decoding process relies directly on the context stored in the Key-Value (KV) cache, rather than reprocessing the original input $$\mathbf{x}$$ at each step.

Next Token Prediction Formula Using KV Cache

In the formula for selecting the next token, `ŷ = argmax_y Pr(y|cache)`, the probability is conditioned on the 'cache' rather than the original input sequence. Explain the primary reason for this distinction and what it implies about the efficiency of the generation process.

Rationale for Conditioning on the Cache

An autoregressive model has just processed a sequence of text and computed a probability for every single word in its vocabulary for the next position. In the context of the standard formula for this process, describe the specific mathematical operation used to select the single most likely next token from this set of probabilities.

Explaining the Token Selection Process

The formula `ŷ = argmax_y Pr(y|cache)` implies that the selection of the next token is a deterministic process, where the single token with the highest calculated probability is always chosen.

An autoregressive language model is generating text and needs to determine the next token. After processing the existing context (which is stored in its cache), the model's final layer outputs the following probabilities for a small subset of its vocabulary:

*   P('mat' | cache) = 0.65
*   P('floor' | cache) = 0.25
*   P('sky' | cache) = 0.05
*   P('the' | cache) = 0.05

According to the standard formula for selecting the single most likely next token, which token will be chosen?

Learn Before

Related