After the final decoder layer has processed the sequence, incorporating information from the input sequence (via cross-attention) and the previously generated output tokens (via masked self-attention), we are left with a sequence of high-dimensional representation vectors. For each position in the output sequence, the decoder stack produces a vector of dimension $d_{model}$ . While this vector encodes rich contextual information, it isn't directly interpretable as a probability distribution over the target vocabulary, which is necessary for tasks like translation or text generation.

The final step in the Transformer's decoder architecture involves converting these output vectors into usable probabilities. This is typically achieved through two sequential operations: a final linear transformation followed by a softmax activation function.

The Final Linear Transformation

The output from the top decoder layer is a tensor of shape (batch_size, target_sequence_length, $d_{model}$ ). The purpose of the final linear layer is to project this $d_{model}$ -dimensional representation vector for each position onto a vector whose dimension equals the size of the target vocabulary, let's denote this size as $V$ .

This is a standard fully connected linear layer without an activation function (or sometimes viewed as having an identity activation). If we represent the output of the final decoder layer as $H_{dec} \in \mathbb{R}^{\text{batch\_size} \times \text{target\_sequence\_length} \times d_{model}}$ , and the weight matrix of the linear layer as $W_{out} \in \mathbb{R}^{d_{model} \times V}$ and its bias as $b_{out} \in \mathbb{R}^{V}$ , the operation can be described as:

\text{Logits} = H_{dec} W_{out} + b_{out}

Here, the matrix multiplication is applied independently to the $d_{model}$ -dimensional vector at each position along the target_sequence_length dimension. The resulting Logits tensor has a shape of (batch_size, target_sequence_length, $V$ ). Each vector of size $V$ at a specific position t in the sequence contains raw, unnormalized scores (logits) for every possible token in the target vocabulary. A higher score suggests a higher likelihood for that token at that position.

Implementation Note: A common technique, particularly in models where the source and target vocabularies are shared (or closely related), is to share the weights between the input embedding layer and this final linear layer. The input embedding layer projects from vocabulary indices (effectively one-hot vectors) to $d_{model}$ , while the final linear layer projects from $d_{model}$ back to vocabulary scores. Sharing the weight matrix $W_{out}$ with the input embedding matrix (possibly transposed) significantly reduces the number of model parameters, especially for large vocabularies, and has been shown to work well empirically.

The Softmax Function

The logits produced by the linear layer are real-valued scores and don't form a probability distribution. To convert these scores into probabilities, the softmax function is applied independently to the logit vector at each position in the target sequence.

For a specific position $t$ in the sequence, let $z_t \in \mathbb{R}^V$ be the logit vector. The softmax function computes the probability $p_i$ for the $i$ -th token in the vocabulary (where $i$ ranges from 1 to $V$ ) as follows:

p_i = \text{Softmax}(z_t)_i = \frac{e^{z_{t,i}}}{\sum_{j=1}^{V} e^{z_{t,j}}}

This operation yields a probability vector $p_t \in \mathbb{R}^V$ for each position $t$ , where:

Each element $p_i$ is non-negative ( $p_i \ge 0$ ).
The sum of all elements in the vector is 1 ( $\sum_{i=1}^{V} p_i = 1$ ).

The output of the entire Transformer decoder stack is therefore a tensor of shape (batch_size, target_sequence_length, $V$ ), where each vector along the last dimension represents a probability distribution over the target vocabulary.

Usage in Training and Inference

Training: During training, this output probability distribution is used to compute the loss, typically Cross-Entropy Loss, against the actual target token at each position. The model's parameters (including those in the final linear layer, attention mechanisms, feed-forward networks, etc.) are updated via backpropagation to minimize this loss, effectively learning to predict the correct next token given the context. Techniques like Label Smoothing are often applied to the target distribution before calculating the loss, which can improve generalization.
Inference: During inference (e.g., generating text), the model operates autoregressively. At each step $t$ $t$ , the probability distribution $p_t$ $p_{t}$ generated by the softmax layer is used to select the next token. Common strategies include:
- Greedy Decoding: Simply selecting the token with the highest probability ( $argmax_i p_i$ ).
- Beam Search: Maintaining several candidate sequences (beams) and exploring the most probable extensions at each step, leading to potentially better results than greedy decoding.
- Sampling Methods: Introducing randomness by sampling from the probability distribution (e.g., temperature sampling, top-k sampling, nucleus sampling) to generate more diverse outputs.

This final linear layer and softmax function bridge the gap between the complex internal representations learned by the Transformer and the discrete, probabilistic nature of language generation tasks. They provide the necessary mechanism to translate the model's understanding into concrete predictions over the vocabulary.