Autoregressive Generation › Worked Example: Tracing Generation Step by Step

To solidify the mechanics, let's trace a concrete numerical example through the full generation pipeline. We will use a simplified vocabulary of six tokens: "The", "cat", "sat", "on", "the", and "mat", and a toy model with hidden dimension $d = 4$. The goal is to show exactly how each token is produced, from raw hidden states through logits and softmax to the final token selection.

Setup. Our prompt is "The cat" (2 tokens). After the prefill phase, the model has produced hidden states $h_1$ and $h_2$ for these tokens. We want to generate the third token.

Step 1: Get the logit vector. The language model head projects the hidden state at the last position to vocabulary logits:

$$\ell = W_{\text{lm}} \cdot h_2 \in \mathbb{R}^{6}$$

Suppose the resulting logit vector is $\ell = [-2.1, \; 0.3, \; 3.1, \; 0.8, \; 0.1, \; -0.5]$, corresponding to the six vocabulary tokens in order.

Step 2: Apply softmax. We convert logits to probabilities using the softmax function. The maximum logit is $3.1$ (for "sat"), so we subtract this for numerical stability before exponentiating:

$$\begin{aligned} e^{\ell - 3.1} &= [e^{-5.2},\; e^{-2.8},\; e^{0},\; e^{-2.3},\; e^{-3.0},\; e^{-3.6}] \\ &\approx [0.006,\; 0.061,\; 1.000,\; 0.100,\; 0.050,\; 0.027] \end{aligned}$$

The sum of these values is approximately $1.244$. Dividing each by this sum gives the probability distribution:

$$P = [0.005, \; 0.049, \; 0.804, \; 0.080, \; 0.040, \; 0.022]$$

Step 3: Select the next token. With greedy decoding, we pick the token with the highest probability. Token index 2 ("sat") has probability $0.804$, so it is selected. The sequence is now "The cat sat".

Step 4: Update the KV cache. The model computes keys and values for "sat" at each layer and appends them to the existing cache. The prompt's keys and values for "The" and "cat" remain unchanged in the cache.

Step 5: Repeat for the next token. The input to the next forward pass is the single token "sat" (not the entire sequence). The model computes its query vector, attends to the full key cache (covering "The", "cat", and "sat"), and produces logits for the fourth token position. If the logit for "on" is now the highest, the sequence becomes "The cat sat on".

This trace makes several things concrete. The probability assigned to "sat" was $0.804$, meaning the model was quite confident given the context "The cat". The other tokens received small but non-zero probability: "on" received $0.08$, which would have been correct if the sentence were "The cat sat on". Greedy decoding committed to "sat" because it was most likely, but sampling could have selected "on" or another token. The step-by-step nature is clear: each token is a separate computation, the result is committed immediately, and the cache is updated before the next step begins.

Why does this formula make sense? Notice that the softmax denominator ensures all probabilities sum to 1, making the output a valid probability distribution. The exponential function amplifies differences in logit values: "sat" with logit 3.1 receives probability 0.804, while "a" with logit 0.3 receives only 0.049, even though the logit difference is only 2.8. This amplification is what makes the model's high-confidence predictions dominant under greedy decoding.