Google

The process of autoregressive generation in a decoder-only Transformer involves a step-by-step architectural flow. At each generation step `i`, the input sequence is formed by concatenating the initial prompt `x` with all previously generated tokens `y_{<i}`. This combined sequence is first converted into embeddings by an embedding layer. The embeddings are then processed through a stack of `L` decoder layers, each containing self-attention and feed-forward network (FFN) modules. The output from the final decoder layer undergoes a linear mapping and is then passed to a Softmax layer. This produces the conditional probability distribution, `Pr(·|x, y_{<i})`, which is used to select the next token `y_i`.

Processing Flow of Autoregressive Generation in a Decoder-Only Transformer

In Transformer-based language models, logits are the raw, unnormalized scores that are output by the model's final linear layer before the application of a Softmax function. They are represented as a sequence of vectors, $\{z_0, ..., z_{m-1}\}$, where each vector corresponds to a token position in the sequence. These vectors are generated by projecting the final hidden states ($h^L$) into the vocabulary space, with each element in a vector representing the score for a potential token.

Logits in Transformer Language Models

In a Transformer-based language model with $L$ layers, the final hidden states are the sequence of output vectors from the last Transformer block, denoted as $\{\mathbf{h}_0^L, \dots, \mathbf{h}_{m-1}^L\}$. Each vector $\mathbf{h}_i^L$ represents the contextualized embedding of the $i$-th token after processing through the entire stack of $L$ layers. This sequence of vectors encapsulates the model's final understanding of the input sequence and is used as the basis for subsequent predictions, such as generating logits for the next token.

Final Hidden States in a Transformer Language Model

In an autoregressive decoder, the probability for the next token is calculated at each step `i` by conditioning on the input `x` and all previously generated tokens `y_{<i}`. The process begins by concatenating `x` and `y_{<i}` and passing them through an embedding layer. This sequence of embeddings is then processed by a stack of decoder layers (which typically include self-attention and feed-forward networks) to produce a sequence of hidden states, `H`. A final linear transformation (using an output weight matrix `W^o`) is applied to these hidden states to get logits, followed by a Softmax function. The probability distribution for the next token, `y_i`, is the probability vector taken from the final position of the output sequence. The formula is: `Pr(·|x, y_{<i}) = (\text{Softmax}(H W^o))_{\text{last}}` where `H = \text{Decoder}([x, y_{<i}])`.

Next-Token Probability Calculation in Autoregressive Decoders

The decoding phase of a Transformer, as illustrated in its diagram, operates sequentially to generate one token at a time. In this step-by-step process, the model uses the token from the previous step as input to an embedding layer, which then generates a new query vector. This query attends to an expanding set of keys and values, comprising those from the initial prompt (prefilling phase) and all previously generated tokens. The output from this self-attention mechanism is processed by a Softmax layer to calculate the conditional probability for the next token, such as `Pr(yn|x, y<n)`. This autoregressive cycle is repeated for each new token in the output sequence.

Diagram of the Decoding Phase

This diagram illustrates the sequential forward pass of a Transformer-based language model, breaking down the process into several key stages:
1.  **Input Tokens**: The model begins with an input sequence of tokens, denoted as $(x_0, x_1, ..., x_{m-1})$.
2.  **Token Embeddings**: These tokens are converted into a sequence of numerical embeddings $(e_0, e_1, ..., e_{m-1})$.
3.  **Transformer Blocks**: The embedding sequence is processed through a stack of $L$ identical Transformer blocks. Each block contains a self-attention sub-layer and a feed-forward network (FFN) sub-layer, and can utilize either a post-norm or pre-norm architecture.
4.  **Final Hidden States**: The final (L-th) block outputs a sequence of contextualized hidden states, $(h^L_0, h^L_1, ..., h^L_{m-1})$.
5.  **Logits**: These hidden states are linearly transformed into a sequence of unnormalized scores called logits, $(z_0, z_1, ..., z_{m-1})$.
6.  **Conditional Probabilities**: A Softmax function converts the logits into conditional probability distributions for each position, such as $\text{Pr}(x_2|x_0, x_1)$.
7.  **Output Tokens**: These probabilities are then used to predict the subsequent tokens in the sequence, $(x_1, x_2, ..., x_m)$.

Diagram of the Transformer Language Model Forward Pass

This diagram illustrates the architectural flow for a single step of autoregressive generation in a decoder-only Transformer. The process begins with the input sequence, formed by concatenating the prompt tokens (`x0...xm`) and any previously generated tokens (`y1...yi-1`). This sequence is fed into an embedding layer. The resulting embeddings are then processed through a stack of `L` decoder layers, each comprising self-attention and feed-forward network (FFN) modules. The output from the final layer is passed through a linear mapping and a Softmax layer to compute the conditional probability distribution, `Pr(·|x, y_{<i})`, for the next token.

Diagram of the Autoregressive Generation Architectural Flow

A decoder-only language model generates text one token at a time in a step-by-step process. Arrange the following steps in the correct chronological order for generating a single new token, given an initial prompt and any previously generated tokens.

In the step-by-step generation process of a decoder-only language model, consider a hypothetical modification at generation step `i`. Instead of using the initial prompt combined with all previously generated tokens as input, the model is only given the initial prompt. What is the most likely consequence of this change on the generated text?

Based on the following case study, identify the component(s) in the processing flow that directly cause the uniform probability distribution and explain the mechanism that leads to this outcome.

Learn Before

Related