Final Linear Layer and Softmax

After the final decoder layer processes the sequence information, incorporating masked self-attention and encoder-decoder attention, we are left with a sequence of output vectors. Each vector at position $t$ in the output sequence represents the decoder's understanding of what the token at that position should be, given the input sequence and the previously generated output tokens up to position $t-1$.

However, these output vectors are typically high-dimensional and contain continuous values. They don't directly tell us which specific word from our vocabulary is the most likely choice for that position. To get from these internal representations to actual predicted words, we need two final steps: a Linear transformation and a Softmax activation function.

The Final Linear Layer

The first step is to project the high-dimensional output vector produced by the decoder stack into a vector whose size matches our target vocabulary. For example, if our model needs to predict words from a vocabulary of 50,000 unique tokens, this linear layer will transform the decoder's output vector (say, of dimension $d_{model} = 512$) into a vector of dimension 50,000.

This is achieved using a standard fully connected linear layer, sometimes called a projection layer. Mathematically, if $x_{out} \in \mathbb{R}^{d_{model}}$ is the output vector from the top decoder layer for a specific position, and the vocabulary size is $V$, the linear layer performs:

$$\text{logits} = W_{proj} x_{out} + b_{proj}$$

Here, $W_{proj}$ is a weight matrix of shape $(V, d_{model})$ and $b_{proj}$ is a bias vector of shape $(V, 1)$. Both $W_{proj}$ and $b_{proj}$ are learned parameters, adjusted during the training process. The resulting vector, often called logits, has dimension $V$. Each element in the logits vector corresponds to a score for one specific token in the vocabulary. A higher score suggests that the corresponding token is more likely, according to the model at this stage.

The Softmax Function

The logits vector provides scores, but they aren't probabilities. The scores can be positive or negative, and they don't sum up to 1. To convert these scores into a proper probability distribution over the vocabulary, we apply the Softmax function.

The Softmax function takes the logits vector as input and outputs a probability vector of the same dimension ($V$). For a logits vector $z = [z_1, z_2, ..., z_V]$, the Softmax function calculates the probability $P_i$ for the $i$-th token in the vocabulary as:

$$P_i = \text{Softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{V} e^{z_j}}$$

This function has two desirable properties:

1. Normalization: All output probabilities sum to 1 ($\sum_{i=1}^{V} P_i = 1$).
2. Monotonicity: It preserves the order of the input scores. The token with the highest logit score will also have the highest probability after applying Softmax.

The resulting vector $P = [P_1, P_2, ..., P_V]$ represents the model's predicted probability distribution over the entire vocabulary for the current position in the output sequence.

Flow from the final decoder output vector through the linear layer and softmax function to produce a probability distribution over the target vocabulary.

During inference (when generating new text), the model typically selects the token with the highest probability from this distribution as the predicted token for the current position. This predicted token is then often fed back into the decoder as input for generating the next token in the sequence. During training, this probability distribution is compared against the actual target token using a loss function (like cross-entropy) to calculate the error and update the model's parameters.

This combination of a linear layer followed by a Softmax function is the standard way to map the internal representations learned by the Transformer's decoder back to the space of observable tokens, enabling the model to make concrete predictions.