Module 15 — Output Projection

The Output Projection is the final stage of a Transformer model. After the encoder or decoder generates hidden representations, these vectors are projected into the vocabulary space to predict the next token.

The output layer converts hidden features into logits, applies the Softmax function to obtain probabilities, and selects the most likely token.

This process is used in almost every Transformer-based language model, including GPT, BERT, T5, LLaMA, Gemma, Qwen, and Mistral.

Topics

Vocabulary Projection
Logits
Softmax
Token Prediction

1. Hidden Representation

The final hidden representation from the Transformer is

Formula

H \in \mathbb{R}^{n\times d_{model}}

where

$n$ = Sequence Length
$d_{model}$ = Hidden Dimension

2. Vocabulary Projection

The hidden vectors are projected into the vocabulary dimension using a linear layer.

Formula

Logits = HW^T+b

where

W \in \mathbb{R}^{|V|\times d_{model}}

b \in \mathbb{R}^{|V|}

Here,

$|V|$ = Vocabulary Size

The output shape is

Logits \in \mathbb{R}^{n\times |V|}

3. Logits

Logits are the raw prediction scores before applying Softmax.

Formula

z = HW^T+b

Each value in the logits vector corresponds to one vocabulary token.

4. Softmax

Softmax converts logits into a probability distribution.

Formula

P_i = \frac {e^{z_i}} {\sum_{j=1}^{|V|}e^{z_j}}

or equivalently,

P = Softmax(Logits)

Properties

0 \le P_i \le 1

and

\sum_{i=1}^{|V|} P_i = 1

5. Token Prediction

The predicted token is the vocabulary token with the highest probability.

Formula

\hat y = argmax(P)

where

P = Softmax(Logits)

6. Complete Output Projection

The complete prediction pipeline is

Logits = HW^T+b

P = Softmax(Logits)

\hat y = argmax(P)

7. Language Modeling Probability

The probability of predicting the next token is

P(x_t | x_{<t}) = Softmax ( HW^T )

where

$x_{<t}$ represents all previously generated tokens.

8. Matrix Dimensions

Suppose

Sequence Length = $n$
Hidden Dimension = $d_{model}$
Vocabulary Size = $|V|$

Hidden Representation

H \in \mathbb{R}^{n\times d_{model}}

Vocabulary Matrix

W \in \mathbb{R}^{|V|\times d_{model}}

Logits

Logits \in \mathbb{R}^{n\times |V|}

Probability Distribution

P \in \mathbb{R}^{n\times |V|}

Predicted Tokens

\hat y \in \mathbb{R}^{n}

9. Weight Tying (Optional)

Many modern language models share the input embedding matrix with the output projection matrix.

Formula

W_{output} = W_{embedding}

This technique reduces the number of parameters and often improves performance.

10. Complete Output Layer Pipeline

The final prediction process in a Transformer model is

H \rightarrow HW^T+b \rightarrow Softmax \rightarrow Token

or mathematically,

P = Softmax ( HW^T+b )

\hat y = argmax(P)

11. Applications

Output Projection is used in

GPT
GPT-2
GPT-3
GPT-4
BERT (Masked Language Modeling Head)
RoBERTa
T5
BART
LLaMA
Gemma
Qwen
DeepSeek
Mistral
Phi
Vision-Language Models
Modern Large Language Models (LLMs)

Summary

Concept	Formula
Hidden Representation	$H\in\mathbb{R}^{n\times d_{model}}$
Vocabulary Projection