LLM Forward Pass Explained: Attention, Embeddings & Logits

In this section, we continue exploring the transformer training process, using large language models (LLMs) as an expository vehicle and building on the tokenization chapter. Here we discuss a worked example of the forward pass, and in the next chapters we cover back propagation, and optimization.

Our objective is to manually work through a full GPT-2 style block to see exactly how embeddings, multi-head attention, residuals, layer norms, the MLP, and the projection to logits interact and then to calculate loss, compute the gradients via back propagation, and update weights using optimisation methods.

Here we cover the forward pass, which happens during inference and during the first stage of learning. The forward pass processes the text input through the network architecture to produce probabilities for the next tokens, where the targets are the same sequence shifted one position forward so that each input token is trained to predict the following token. Later sections cover back propagation, and optimization (weight updating).

💡

Alert: Here I highlight the steps and computations of the forward pass procedure with a numerical example to make the ideas concrete. The approach is consistent with our simplified GPT-2 model. These results are preliminary and will be validated with a formal mathematical verifier in due course.

We will use three dimensions, which gives us room to represent two input tokens (‘AI’ and ‘rocks’) plus an EOS token for proper token handling. It also ensures that LayerNorm and other stages do not collapse the gradient during backpropagation. This makes the toy both realistic and pedagogically useful. We use $L$ = 1 layer, $d_m = 4$ token embedding dimensions and $h=2$ heads and otherwise stay faithful to GPT-2.

‣

Vocabulary and tokenization

We’ll use the smallest meaningful vocabulary and input sequence that demonstrates the workings of the transformer model while also including EOS handling.

Vocab size $V$ =3
Tokens: "AI" → ID 0, "rocks" → ID 1, "<EOS>" → ID 2
Input sequence: [AI, rocks, <EOS>] → [0, 1, 2]
Target sequence: [rocks, <EOS>, <nothing>] → [1, 2, –]

‣

Embedding lookup from

E

The model learns an embedding matrix $E∈RV×dm=R3×4$ with random initialization

$E$ is the embedding matrix (size $V \times d_m$ ).
Each row corresponds to one token in the vocabulary.
Multiplying a one-hot vector $eₜ$ for token $t$ (length $V$ ) by $E$ selects the matching rows or equivalently the row can be selected via lookup. Here we select all rows which are dense vector representations of the tokens.
In our example, we’ll set the initial values before learning as small, zero mean decimals, the first row is for token $0$ AI, the second is token $1$ rocks, and the third row is for token 2 $“<EOS>”$

E = \begin{bmatrix} 0.120 & -0.050 & 0.030 & 0.025 \\ -0.080 & 0.090 & -0.020 & -0.015 \\ 0.050 & 0.020 & -0.070 & 0.040 \end{bmatrix}

‣

Positional embeddings

P

The model learns a positional embedding matrix $P \in \mathbb{R}^{L \times d_m}$

$P$ is the embedding matrix (size $L$ x $d_m$ ).
Initially, the length of $L$ is the maximum sequence length the model can handle.
Each of the rows corresponds to a position in the input sequence. Again, we’ll use small, zero mean decimals in our example.

P = \begin{bmatrix} 0.010 & -0.020 & 0.000 & 0.005 \\ -0.030 & 0.020 & 0.010 & -0.010 \\ 0.020 & 0.010 & -0.020 & 0.000 \end{bmatrix}

Element wise addition of $E$ and $P$ , $xi=E[token_i]+P[position_i]$ , creates the tensor that enters the first transformer block.

\begin{aligned} x_0 &= E[\text{AI}] + P[0] = [0.120,\,-0.050,\,0.030,\,0.025] + [0.010,\,-0.020,\,0.000,\,0.005] \\ &= [0.130,\,-0.070,\,0.030,\,0.030] \\ x_1 &= E[\text{rocks}] + P[1] = [-0.080,\,0.090,\,-0.020,\,-0.015] + [-0.030,\,0.020,\,0.010,\,-0.010] \\ &= [-0.110,\,0.110,\,-0.010,\,-0.025] \\ x_2 &= E[\text{EOS}] + P[2] = [0.050,\,0.020,\,-0.070,\,0.040] + [0.020,\,0.010,\,-0.020,\,0.000] \\ &= [0.070,\,0.030,\,-0.090,\,0.040] \\\\ x &= \begin{bmatrix} 0.130 & -0.070 & 0.030 & 0.030 \\ -0.110 & 0.110 & -0.010 & -0.025 \\ 0.070 & 0.030 & -0.090 & 0.040 \end{bmatrix} \end{aligned}

‣

LayerNorm 1

LayerNorm rescales each token vector to have a mean $0$ and variance $1$ . This prevents arbitrary embedding magnitudes from dominating downstream calculations due to differences in their scale.

Compute mean and variance of each vector and normalize vector elements.

\begin{equation*} \mu_i = \frac{1}{d} \sum_{j=1}^{d} x_{ij}, \qquad \sigma_i^2 = \frac{1}{d} \sum_{j=1}^{d} (x_{ij} - \mu_i)^2, \qquad \hat{x}_{ij} = \frac{x_{ij} - \mu_i}{\sqrt{\sigma_i^2 + \epsilon}} \end{equation*}

After normalization, each element $\hat{x}_{ij}$ is rescaled and shifted by learnable parameters $\gamma_j$ (scale) and $\beta_j$ (shift). They are initialized as $\gamma_j = 1$ and $\beta_j = 0$ , so LayerNorm starts as pure normalization. During training, $\gamma_j$ can grow or shrink each dimension (scaling) and $\beta_j$ can move it up or down (shifting), giving the model flexibility to adapt features for the next layers. A small constant ε is always included in the denominator to prevent division by zero when the variance happens to be exactly zero (PyTorch uses $ε=1e-5$ ), in our worked example it doesn’t affect the numbers.

\begin{equation*} y_{ij} = \gamma_j \hat{x}_{ij} + \beta_j \end{equation*}

For our worked example, we now numerically demonstrate standardisation, scale and shift to produce the input tensor for self attention.

\begin{align*}x_0 &= [0.130,\,-0.070,\,0.030,\,0.030] \\\mu_0 &= \tfrac{1}{4}(0.130 + (-0.070) + 0.030 + 0.030) = 0.030 \\\sigma_0^2 &= \tfrac{1}{4}\big((0.130 - 0.030)^2 + (-0.070 - 0.030)^2 + (0.030 - 0.030)^2 + (0.030 - 0.030)^2\big) = 0.005 \\\hat{x}_0 &= \frac{[0.130,\,-0.070,\,0.030,\,0.030] - [0.030,\,0.030,\,0.030,\,0.030]}{\sqrt{0.005}}           = [1.414,\,-1.414,\,0.000,\,0.000] \\y_0 &= [1.414,\,-1.414,\,0.000,\,0.000] \\\\x_1 &= [-0.110,\,0.110,\,-0.010,\,-0.025] \\\mu_1 &= \tfrac{1}{4}((-0.110) + 0.110 + (-0.010) + (-0.025)) = -0.009 \\\sigma_1^2 &= 0.0071 \\\hat{x}_1 &= \frac{[-0.110,\,0.110,\,-0.010,\,-0.025] - [-0.009,\,-0.009,\,-0.009,\,-0.009]}{\sqrt{0.0071}}           \approx [-1.199,\,1.417,\,-0.012,\,-0.207] \\y_1 &\approx [-1.199,\,1.417,\,-0.012,\,-0.207] \\\\x_2 &= [0.070,\,0.030,\,-0.090,\,0.040] \\\mu_2 &= \tfrac{1}{4}(0.070 + 0.030 + (-0.090) + 0.040) = 0.013 \\\sigma_2^2 &= 0.0039 \\\hat{x}_2 &= \frac{[0.070,\,0.030,\,-0.090,\,0.040] - [0.013,\,0.013,\,0.013,\,0.013]}{\sqrt{0.0039}}           \approx [0.914,\,0.278,\,-1.656,\,0.464] \\y_2 &\approx [0.914,\,0.278,\,-1.656,\,0.464] \\\\Y &= \begin{bmatrix}1.414 & -1.414 & 0.000 & 0.000 \\-1.199 & 1.417 & -0.012 & -0.207 \\0.914 & 0.278 & -1.656 & 0.464\end{bmatrix}\end{align*}

‣

Multi-head self attention

Self-attention transforms the normalized input into query, key, and value vectors, uses dot products of queries and keys to calculate attention scores, normalizes these scores with softmax to produce attention weights, and then applies these weights to the values to produce the attention output. Note that in GPT-2, bias terms are added to the linear projections (Q, K, V, and the output). These are per-feature offsets that don’t affect the core tensor mechanics, so they are omitted in this worked example for clarity. In the full GPT-2 model, however, each projection includes a learned bias vector with the same output dimension as its corresponding matrix (e.g., matching the size of Q, K, or V). These biases are initialized to zeros and added immediately after each linear projection. This allows each output feature dimension to adjust independently by shifting its activation up or down during training.

‣

Q, K, V

matrices

Query, Key, and Value matrices $W_Q, W_K, W_V \in \mathbb{R}^{d_m \times d_h}$ are learned parameter matrices for each head. They are initialized with small values drawn from a normal distribution with zero mean, updated during optimization following back propagation.
We set the model dimension 4. With $h=2$ heads, the head dimension is $dh=dm/h=2$ . This means that for each head the weight matrices are $WQ∈R4×2$ , $WK∈R4×2$ , and $WV∈R4×2$ .
$Q$ , $K$ , and $V$ are usually initialized from a normal distribution $\mathcal{N}(0,\sigma^2)$ with $\sigma \sim 1/\sqrt{d_m}$ , but for our worked example we will instead use small hand-chosen values.

\begin{aligned}\text{Head 1:} \quad & W_Q^{(1)} = \begin{bmatrix}0.120 & -0.050 \\0.030 & 0.025 \\-0.020 & 0.015 \\0.010 & -0.030\end{bmatrix}, \quad W_K^{(1)} = \begin{bmatrix}0.050 & -0.020 \\-0.010 & 0.040 \\0.030 & 0.010 \\-0.020 & 0.025\end{bmatrix}, \quad W_V^{(1)} = \begin{bmatrix}0.010 & 0.020 \\0.030 & -0.010 \\-0.020 & 0.015 \\0.025 & -0.030\end{bmatrix} \\\\\text{Head 2:} \quad & W_Q^{(2)} = \begin{bmatrix}-0.020 & 0.030 \\0.015 & -0.010 \\0.025 & 0.020 \\-0.030 & 0.010\end{bmatrix}, \quad W_K^{(2)} = \begin{bmatrix}0.010 & 0.015 \\-0.025 & 0.020 \\0.005 & -0.010 \\0.020 & -0.005\end{bmatrix}, \quad W_V^{(2)} = \begin{bmatrix}0.030 & -0.010 \\0.010 & 0.020 \\-0.015 & 0.025 \\0.005 & -0.020\end{bmatrix}\end{aligned}

‣

Linear projections forward

Compute the linear projections $Q$ , $K$ , and $V$ for head 1 and for head 2. For Head 1 and Head 2, the linear projections are obtained by multiplying the normalized input $Y$ with the weight matrices. Multiplying the normalized input $Y$ by these matrices produces the linear projections $Q = YW_Q$ , $K = YW_K$ , and $V = YW_V$ .

\text{Head 1:} \quad \begin{aligned}Q^{(1)} &= Y W_Q^{(1)} = \begin{bmatrix}1.414 & -1.414 & 0.000 & 0.000 \\-1.291 & 1.514 & -0.016 & -0.207 \\0.942 & 0.287 & -1.680 & 0.451\end{bmatrix}\begin{bmatrix}0.120 & -0.050 \\0.030 & 0.025 \\-0.020 & 0.015 \\0.010 & -0.030\end{bmatrix} \\&\approx \begin{bmatrix}0.1273 & -0.1061 \\-0.1113 & 0.1084 \\0.1598 & -0.0787\end{bmatrix} \\\\[2mm]K^{(1)} &= Y W_K^{(1)} = \begin{bmatrix}1.414 & -1.414 & 0.000 & 0.000 \\-1.291 & 1.514 & -0.016 & -0.207 \\0.942 & 0.287 & -1.680 & 0.451\end{bmatrix}\begin{bmatrix}0.050 & -0.020 \\-0.010 & 0.040 \\0.030 & 0.010 \\-0.020 & 0.025\end{bmatrix} \\&\approx \begin{bmatrix}0.0848 & -0.0848 \\-0.0760 & 0.0810 \\-0.0152 & -0.0129\end{bmatrix} \\\\[2mm]V^{(1)} &= Y W_V^{(1)} = \begin{bmatrix}1.414 & -1.414 & 0.000 & 0.000 \\-1.291 & 1.514 & -0.016 & -0.207 \\0.942 & 0.287 & -1.680 & 0.451\end{bmatrix}\begin{bmatrix}0.010 & 0.020 \\0.030 & -0.010 \\-0.020 & 0.015 \\0.025 & -0.030\end{bmatrix} \\&\approx\begin{bmatrix}-0.0283 & 0.0424 \\0.0050 & -0.0070 \\-0.0330 & 0.0220\end{bmatrix}\end{aligned}

\text{Head 2:} \quad \begin{aligned} Q^{(2)} &= Y W_Q^{(2)} = \begin{bmatrix} 1.414 & -1.414 & 0.000 & 0.000 \\ -1.291 & 1.514 & -0.016 & -0.207 \\ 0.942 & 0.287 & -1.680 & 0.451 \end{bmatrix} \begin{bmatrix} -0.020 & 0.030 \\ 0.015 & -0.010 \\ 0.025 & 0.020 \\ -0.030 & 0.010 \end{bmatrix} \\ &\approx \begin{bmatrix} -0.0424 & 0.0566 \\ 0.0463 & -0.0566 \\ -0.0419 & -0.0215 \end{bmatrix} \\\\[2mm] K^{(2)} &= Y W_K^{(2)} = \begin{bmatrix} 1.414 & -1.414 & 0.000 & 0.000 \\ -1.291 & 1.514 & -0.016 & -0.207 \\ 0.942 & 0.287 & -1.680 & 0.451 \end{bmatrix} \begin{bmatrix} 0.010 & 0.015 \\ -0.025 & 0.020 \\ 0.005 & -0.010 \\ 0.020 & -0.005 \end{bmatrix} \\ &\approx \begin{bmatrix} -0.0415 & -0.0071 \\ 0.0467 & 0.0209 \\ -0.0109 & -0.0240 \end{bmatrix} \\\\[2mm] V^{(2)} &= Y W_V^{(2)} = \begin{bmatrix} 1.414 & -1.414 & 0.000 & 0.000 \\ -1.291 & 1.514 & -0.016 & -0.207 \\ 0.942 & 0.287 & -1.680 & 0.451 \end{bmatrix} \begin{bmatrix} 0.030 & -0.010 \\ 0.010 & 0.020 \\ -0.015 & 0.025 \\ 0.005 & -0.020 \end{bmatrix} \\ &\approx \begin{bmatrix} 0.0283 & -0.0424 \\ -0.0269 & 0.0629 \\ -0.0216 & -0.0178 \end{bmatrix} \end{aligned}

‣

Attentions scores forward

Next compute attention scores for each head $S=QK⊤/\sqrt{d_h}$ , which measure the similarity between each query and key, indicating how much focus one token should place on another before masking and normalization.

S^{(1)} \;=\; \frac{Q^{(1)} (K^{(1)})^\top}{\sqrt{d_h}} = \frac{ \begin{bmatrix} 0.1273 & -0.1061 \\ -0.1113 & 0.1084 \\ 0.1598 & -0.0787 \end{bmatrix} \begin{bmatrix} 0.0848 & -0.0760 & -0.0152 \\ -0.0848 & 0.0810 & -0.0129 \end{bmatrix} }{\sqrt{2}} \\[6pt] \approx \frac{ \begin{bmatrix} 0.01886 & -0.01939 & -0.00056 \\ -0.01767 & 0.01820 & 0.00029 \\ 0.01953 & -0.01992 & -0.00142 \end{bmatrix} }{1.414} \;=\; \begin{bmatrix} 0.0133 & -0.0137 & -0.00040 \\ -0.0125 & 0.0129 & 0.00020 \\ 0.0138 & -0.0141 & -0.00100 \end{bmatrix}

S^{(2)} \;=\; \frac{Q^{(2)} (K^{(2)})^\top}{\sqrt{d_h}} = \frac{ \begin{bmatrix} -0.0424 & 0.0566 \\ 0.0463 & -0.0566 \\ -0.0419 & -0.0215 \end{bmatrix} \begin{bmatrix} -0.0415 & -0.0071 & -0.0109 \\ 0.0467 & 0.0209 & -0.0240 \end{bmatrix} }{\sqrt{2}} \\[6pt] \approx \frac{ \begin{bmatrix} 0.00440 & 0.00148 & -0.00090 \\ -0.00456 & -0.00151 & 0.00086 \\ 0.00074 & -0.00015 & 0.00098 \end{bmatrix} }{1.414} \;=\; \begin{bmatrix} 0.00311 & 0.00105 & -0.00064 \\ -0.00322 & -0.00107 & 0.00061 \\ 0.00052 & -0.00011 & 0.00069 \end{bmatrix}

‣

Causal mask forward

Adding the causal mask to $S$ sets future tokens to $−∞$ , ensuring that each token can only attend to itself and earlier tokens. The mask $M$ is an $L$ x $L$ matrix where each row represents a query position and each column represents a key position. Zero entries allow attention for self or past positions and $−∞$ prevents attention for future positions. With $L$ = $3$ .

M \in \mathbb{R}^{3 \times 3}, \quad M = \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix}

\begin{aligned} S^{(1)}_{\text{masked}} &= \begin{bmatrix} 0.0133 & -0.0137 & -0.00040 \\ -0.0125 & 0.0129 & 0.00020 \\ 0.0138 & -0.0141 & -0.00100 \end{bmatrix} + \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix} \\ &= \begin{bmatrix} 0.0133 & -\infty & -\infty \\ -0.0125 & 0.0129 & -\infty \\ 0.0138 & -0.0141 & -0.00100 \end{bmatrix} \\\\[2mm] S^{(2)}_{\text{masked}} &= \begin{bmatrix} 0.00311 & 0.00105 & -0.00064 \\ -0.00322 & -0.00107 & 0.00061 \\ 0.00052 & -0.00011 & 0.00069 \end{bmatrix} + \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix} \\ &= \begin{bmatrix} 0.00311 & -\infty & -\infty \\ -0.00322 & -0.00107 & -\infty \\ 0.00052 & -0.00011 & 0.00069 \end{bmatrix} \end{aligned}

Applying the softmax function row-wise to the masked scores $S_{\text{masked}}$ converts each row into a probability distribution. This ensures that the attention weights are non-negative and sum to $1$ , so each token’s output is a weighted average of the values it is allowed to attend to:

A_{ij} = \frac{\exp\!\big(S^{\mathrm{masked}}_{ij}\big)}{\sum_{k=1}^{L} \exp\!\big(S^{\mathrm{masked}}_{ik}\big)}

\text{Head 1:}\quad A^{(1)} = \begin{bmatrix} \frac{e^{0.0133}}{e^{0.0133} + e^{-\infty} + e^{-\infty}} & \frac{e^{-\infty}}{e^{0.0133} + e^{-\infty} + e^{-\infty}} & \frac{e^{-\infty}}{e^{0.0133} + e^{-\infty} + e^{-\infty}} \\[2mm] \frac{e^{-0.0125}}{e^{-0.0125} + e^{0.0129} + e^{-\infty}} & \frac{e^{0.0129}}{e^{-0.0125} + e^{0.0129} + e^{-\infty}} & \frac{e^{-\infty}}{e^{-0.0125} + e^{0.0129} + e^{-\infty}} \\[2mm] \frac{e^{0.0138}}{e^{0.0138} + e^{-0.0141} + e^{-0.00100}} & \frac{e^{-0.0141}}{e^{0.0138} + e^{-0.0141} + e^{-0.00100}} & \frac{e^{-0.00100}}{e^{0.0138} + e^{-0.0141} + e^{-0.00100}} \end{bmatrix} \approx \begin{bmatrix} 1.000 & 0.000 & 0.000 \\ 0.494 & 0.506 & 0.000 \\ 0.338 & 0.329 & 0.333 \end{bmatrix}

\text{Head 2:}\quad A^{(2)} = \begin{bmatrix} \frac{e^{0.00311}}{e^{0.00311} + e^{-\infty} + e^{-\infty}} & \frac{e^{-\infty}}{e^{0.00311} + e^{-\infty} + e^{-\infty}} & \frac{e^{-\infty}}{e^{0.00311} + e^{-\infty} + e^{-\infty}} \\[2mm] \frac{e^{-0.00322}}{e^{-0.00322} + e^{-0.00107} + e^{-\infty}} & \frac{e^{-0.00107}}{e^{-0.00322} + e^{-0.00107} + e^{-\infty}} & \frac{e^{-\infty}}{e^{-0.00322} + e^{-0.00107} + e^{-\infty}} \\[2mm] \frac{e^{0.00052}}{e^{0.00052} + e^{-0.00011} + e^{0.00069}} & \frac{e^{-0.00011}}{e^{0.00052} + e^{-0.00011} + e^{0.00069}} & \frac{e^{0.00069}}{e^{0.00052} + e^{-0.00011} + e^{0.00069}} \end{bmatrix} \approx \begin{bmatrix} 1.000 & 0.000 & 0.000 \\ 0.499 & 0.501 & 0.000 \\ 0.333 & 0.333 & 0.333 \end{bmatrix}

‣

Attention output forward

Multiplying the attention weights $A^{(h)}$ by the value matrix $V^{(h)}$ computes the attention output for each head $Z(h)$ , where each token representation becomes a weighted sum of the value vectors according to its learned attention distribution.

Z^{(h)} = A^{(h)} V^{(h)}

\text{Head 1:}\quad Z^{(1)} = A^{(1)} V^{(1)} = \begin{bmatrix} 1.000 & 0.000 & 0.000 \\ 0.494 & 0.506 & 0.000 \\ 0.338 & 0.329 & 0.333 \end{bmatrix} \begin{bmatrix} -0.012 & 0.018 \\ 0.029 & -0.054 \\ -0.030 & 0.037 \end{bmatrix} = \begin{bmatrix} 1.000\cdot[-0.012,\,0.018] + 0.000\cdot[0.029,\,-0.054] + 0.000\cdot[-0.030,\,0.037] \\ 0.494\cdot[-0.012,\,0.018] + 0.506\cdot[0.029,\,-0.054] + 0.000\cdot[-0.030,\,0.037] \\ 0.338\cdot[-0.012,\,0.018] + 0.329\cdot[0.029,\,-0.054] + 0.333\cdot[-0.030,\,0.037] \end{bmatrix} \approx \begin{bmatrix} -0.012 & 0.018 \\ 0.00875 & -0.01843 \\ -0.00451 & 0.00064 \end{bmatrix}

\text{Head 2:}\quad Z^{(2)} = A^{(2)} V^{(2)} = \begin{bmatrix} 1.000 & 0.000 & 0.000 \\ 0.499 & 0.501 & 0.000 \\ 0.333 & 0.333 & 0.333 \end{bmatrix} \begin{bmatrix} 0.028 & -0.043 \\ -0.027 & 0.063 \\ -0.022 & -0.018 \end{bmatrix} = \begin{bmatrix} 1.000\cdot[0.028,\,-0.043] + 0.000\cdot[-0.027,\,0.063] + 0.000\cdot[-0.022,\,-0.018] \\ 0.499\cdot[0.028,\,-0.043] + 0.501\cdot[-0.027,\,0.063] + 0.000\cdot[-0.022,\,-0.018] \\ 0.333\cdot[0.028,\,-0.043] + 0.333\cdot[-0.027,\,0.063] + 0.333\cdot[-0.022,\,-0.018] \end{bmatrix} \approx \begin{bmatrix} 0.028 & -0.043 \\ 0.00045 & 0.01016 \\ -0.007 & 0.00067 \end{bmatrix}

Concatenating the outputs from all heads combines their information into a single matrix. Each head output $Z^{(h)}$ has shape $L \times d_h$ , and placing them side by side along the feature axis produces $Z \in \mathbb{R}^{L \times d_m}$ . The model dimension is $d_m = h \cdot d_h$ . With $h = 2$ and $d_h = 2$ , the result has shape $3 \times 4$ .

\text{Concatenation:}\quad Z = \text{Concat}\big(Z^{(1)}, Z^{(2)}\big) = \begin{bmatrix} -0.012 & 0.018 & 0.028 & -0.043 \\ 0.00875 & -0.01843 & 0.00045 & 0.01016 \\ -0.00451 & 0.00064 & -0.007 & 0.00067 \end{bmatrix}

‣

Output projection forward

Output projection mixes the information from different heads into a single representation. After concatenation, the heads are only placed side by side without interacting. The projection uses a learned weight matrix $W_O \in \mathbb{R}^{d_m \times d_m}$ to linearly combine the head outputs, producing $O = Z W_O$ . In practice, $W_O \in \mathbb{R}^{d_m \times d_m}$ is initialized with small random values, often from $\mathcal{N}(0, 1/d_m)$ . For our toy example, we will use small hand-chosen values for $W_O$ to keep the calculations simple and easy to follow:

W_O = \begin{bmatrix} 0.200 & -0.100 & 0.050 & 0.000 \\ 0.050 & 0.150 & -0.050 & 0.100 \\ -0.100 & 0.050 & 0.200 & -0.050 \\ 0.000 & -0.100 & 0.050 & 0.150 \end{bmatrix}

\text{Output projection:}\quad O = Z W_O = \begin{bmatrix} -0.012 & 0.018 & 0.028 & -0.043 \\ 0.00875 & -0.01843 & 0.00045 & 0.01016 \\ -0.00451 & 0.00064 & -0.007 & 0.00067 \end{bmatrix} \begin{bmatrix} 0.200 & -0.100 & 0.050 & 0.000 \\ 0.050 & 0.150 & -0.050 & 0.100 \\ -0.100 & 0.050 & 0.200 & -0.050 \\ 0.000 & -0.100 & 0.050 & 0.150 \end{bmatrix} = \begin{bmatrix} -0.004 & 0.010 & 0.002 & -0.006 \\ 0.001 & -0.005 & 0.002 & -0.000 \\ -0.002 & 0.000 & -0.002 & 0.001 \end{bmatrix}

‣

Residual connection

Adding the residual connection ensures that the original input to the block is preserved and combined with the transformed output. This helps stabilize training by allowing gradients to flow more easily and prevents important information from being lost through the attention and projection layers.

\text{Residual connection:}\quad R = O + Y = \begin{bmatrix} -0.004 & 0.010 & 0.002 & -0.006 \\ 0.001 & -0.005 & 0.002 & -0.000 \\ -0.002 & 0.000 & -0.002 & 0.001 \end{bmatrix} + \begin{bmatrix} 1.414 & -1.414 & 0.000 & 0.000 \\ -1.199 & 1.417 & -0.012 & -0.207 \\ 0.914 & 0.278 & -1.656 & 0.464 \end{bmatrix} \;\approx\; \begin{bmatrix} 1.410 & -1.404 & 0.002 & -0.006 \\ -1.198 & 1.412 & -0.010 & -0.207 \\ 0.912 & 0.278 & -1.658 & 0.465 \end{bmatrix}

‣

LayerNorm 2

The second LayerNorm works the same way as the first one: it normalizes each token’s vector (now the output of the attention + residual step) so that its values have mean zero and unit variance, which helps stabilize the scale of features before sending them into the MLP.

\mu_i = \frac{1}{d} \sum_{j=1}^{d} R_{ij}, \qquad \sigma_i^2 = \frac{1}{d} \sum_{j=1}^{d} (R_{ij} - \mu_i)^2 \hat{R}_{ij} = \frac{R_{ij} - \mu_i}{\sqrt{\sigma_i^2 + \epsilon}}, \qquad Y^{(2)}_{ij} = \gamma_j \hat{R}_{ij} + \beta_j \quad (\gamma=[1,1],\ \beta=[0,0])

\begin{aligned} \mu_0 &= \frac{1}{4}(1.410 + (-1.404) + 0.002 + (-0.006)) \approx 0.0005 \\ \sigma_0^2 &\approx 0.9905 \\ \hat{r}_0 &\approx [1.414, -1.412, 0.002, -0.006] \\ \\ \mu_1 &= \frac{1}{4}(-1.198 + 1.412 + (-0.010) + (-0.207)) \approx -0.00075 \\ \sigma_1^2 &\approx 0.868 \\ \hat{r}_1 &\approx [-1.280, 1.517, -0.011, -0.223] \\ \\ \mu_2 &= \frac{1}{4}(0.912 + 0.278 + (-1.658) + 0.465) \approx -0.00075 \\ \sigma_2^2 &\approx 0.9685 \\ \hat{r}_2 &\approx [0.923, 0.283, -1.687, 0.473] \end{aligned}

Y^{(2)} \approx \begin{bmatrix} 1.414 & -1.412 & 0.002 & -0.006 \\ -1.414 & 1.666 & -0.012 & -0.244 \\ 0.952 & 0.288 & -1.728 & 0.485 \end{bmatrix}

‣

Multilayer perceptron

The MLP sublayer takes the normalized residual output and passes it through two linear layers with an activation (non-linear function) applied in between (GPT-2 uses GELU), the first expands the feature dimension, and the second projects it back down to $d_m$ , producing a transformed $L×d_m$ matrix for each token.

H = \text{GELU}\!\big(Y^{(2)} W_1 + b_1\big), \quad U = H W_2 + b_2

The parameters are as follows:

W_1 \in \mathbb{R}^{d_m \times d_{ff}}, \quad b_1 \in \mathbb{R}^{d_{ff}}

W_2 \in \mathbb{R}^{d_{ff} \times d_m}, \quad b_2 \in \mathbb{R}^{d_m}

For this worked example $d_m = 4, \qquad d_{ff} = 4 d_m = 16$ . We choose small, hand-picked values to keep the arithmetic explicit:

W_1 = \begin{bmatrix} 0.100 & -0.050 & 0.020 & 0.080 & -0.030 & 0.040 & 0.010 & -0.020 & 0.030 & -0.040 & 0.020 & 0.050 & -0.010 & 0.030 & 0.040 & -0.020 \\ 0.050 & -0.020 & 0.010 & -0.030 & 0.040 & 0.020 & -0.050 & 0.030 & -0.010 & 0.020 & 0.030 & -0.040 & 0.010 & -0.020 & 0.050 & 0.030 \\ -0.030 & 0.040 & 0.020 & -0.010 & 0.030 & -0.050 & 0.040 & 0.020 & -0.020 & 0.010 & 0.050 & -0.030 & 0.040 & 0.010 & -0.020 & 0.030 \\ 0.020 & -0.010 & 0.030 & 0.050 & -0.020 & 0.010 & 0.040 & -0.030 & 0.020 & -0.050 & 0.010 & 0.030 & -0.040 & 0.020 & 0.050 & -0.010 \end{bmatrix},\quad b_1 = \mathbf{0} \in \mathbb{R}^{16}

W_2 = \begin{bmatrix} 0.050 & -0.100 & 0.020 & 0.030 \\ 0.100 & 0.050 & -0.020 & 0.010 \\ -0.050 & 0.150 & 0.040 & -0.030 \\ 0.200 & -0.050 & 0.010 & 0.050 \\ -0.100 & 0.050 & 0.030 & -0.020 \\ 0.050 & 0.100 & -0.010 & 0.020 \\ 0.150 & -0.100 & 0.020 & 0.040 \\ -0.050 & 0.050 & 0.010 & -0.030 \\ 0.030 & -0.020 & 0.050 & 0.010 \\ -0.010 & 0.040 & -0.050 & 0.030 \\ 0.020 & 0.010 & 0.030 & -0.040 \\ 0.040 & -0.030 & 0.020 & 0.050 \\ -0.020 & 0.050 & 0.010 & -0.010 \\ 0.010 & -0.010 & 0.040 & 0.020 \\ 0.050 & 0.020 & -0.030 & 0.010 \\ -0.030 & 0.030 & 0.020 & -0.020 \end{bmatrix},\quad b_2 = \mathbf{0} \in \mathbb{R}^{4}

‣

Linear transformation forward

The first linear transformation multiplies the normalized input $Y^{(2)}$ by the weight matrix $W_1$ and adds the bias $b_1$ . This produces the pre-activation matrix for the hidden layer.

H_{\text{pre}} = Y^{(2)} W_1 + b_1

\begin{aligned} Y^{(2)} &\approx \begin{bmatrix} 1.414 & -1.412 & 0.002 & -0.006 \\ -1.198 & 1.517 & -0.011 & -0.223 \\ 0.923 & 0.283 & -1.687 & 0.473 \end{bmatrix}, \\ W_1 &= \begin{bmatrix} 0.100 & -0.050 & 0.020 & 0.080 & -0.030 & 0.040 & 0.010 & -0.020 & 0.030 & -0.040 & 0.020 & 0.050 & -0.010 & 0.030 & 0.040 & -0.020 \\ 0.050 & -0.020 & 0.010 & -0.030 & 0.040 & 0.020 & -0.050 & 0.030 & -0.010 & 0.020 & 0.030 & -0.040 & 0.010 & -0.020 & 0.050 & 0.030 \\ -0.030 & 0.040 & 0.020 & -0.010 & 0.030 & -0.050 & 0.040 & 0.020 & -0.020 & 0.010 & 0.050 & -0.030 & 0.040 & 0.010 & -0.020 & 0.030 \\ 0.020 & -0.010 & 0.030 & 0.050 & -0.020 & 0.010 & 0.040 & -0.030 & 0.020 & -0.050 & 0.010 & 0.030 & -0.040 & 0.020 & 0.050 & -0.010 \end{bmatrix}, \\ b_1 &= \mathbf{0} \in \mathbb{R}^{16} \end{aligned}

\begin{aligned} H_{\text{pre}} &= Y^{(2)} W_1 + b_1 \\ &\approx \begin{bmatrix} 0.250 & -0.179 & 0.050 & 0.120 & -0.069 & 0.080 & 0.020 & -0.030 & 0.060 & -0.050 & 0.030 & 0.070 & -0.020 & 0.040 & 0.050 & -0.010 \\ -0.255 & 0.220 & -0.055 & -0.145 & 0.098 & -0.120 & -0.030 & 0.050 & -0.040 & 0.060 & -0.020 & -0.050 & 0.030 & -0.040 & 0.070 & 0.020 \\ 0.122 & 0.028 & -0.198 & 0.178 & 0.030 & -0.068 & 0.050 & -0.020 & 0.010 & 0.040 & -0.060 & 0.020 & 0.050 & 0.010 & -0.030 & 0.040 \end{bmatrix} \end{aligned}

‣

GELU activation forward

The activation function used in GPT-2 is the Gaussian Error Linear Unit (GELU), which smoothly weights inputs by their probability of being positive under a standard normal distribution. This provides a smoother nonlinearity than ReLU or tanh and improves performance in transformer models.

\begin{aligned}\text{GELU}(x) &= x \cdot \Phi(x) \\&= \tfrac{1}{2}x\Big(1 + \mathrm{erf}\!\big(\tfrac{x}{\sqrt{2}}\big)\Big) \\\text{Approximation:}\quad\text{GELU}(x) &\approx 0.5x \left(1 + \tanh\!\Big(\sqrt{\tfrac{2}{\pi}}\,(x + 0.044715x^3)\Big)\right)\end{aligned}

Worked numbers (tanh approximation applied elementwise to $H_pre$ and rounded to 3 d.p.):

\begin{aligned} H_{\text{pre}} &\approx \begin{bmatrix} 0.250 & -0.179 & 0.050 & 0.120 & -0.069 & 0.080 & 0.020 & -0.030 & 0.060 & -0.050 & 0.030 & 0.070 & -0.020 & 0.040 & 0.050 & -0.010 \\ -0.255 & 0.220 & -0.055 & -0.145 & 0.098 & -0.120 & -0.030 & 0.050 & -0.040 & 0.060 & -0.020 & -0.050 & 0.030 & -0.040 & 0.070 & 0.020 \\ 0.122 & 0.028 & -0.198 & 0.178 & 0.030 & -0.068 & 0.050 & -0.020 & 0.010 & 0.040 & -0.060 & 0.020 & 0.050 & 0.010 & -0.030 & 0.040 \end{bmatrix} \\[1em] H &\approx \begin{bmatrix} 0.136 & -0.090 & 0.027 & 0.119 & -0.034 & 0.043 & 0.011 & -0.016 & 0.032 & -0.026 & 0.016 & 0.038 & -0.010 & 0.022 & 0.027 & -0.005 \\ -0.152 & 0.118 & -0.031 & -0.076 & 0.050 & -0.061 & -0.015 & 0.025 & -0.020 & 0.030 & -0.010 & -0.025 & 0.016 & -0.022 & 0.038 & 0.011 \\ 0.065 & 0.027 & -0.108 & 0.096 & 0.016 & -0.037 & 0.027 & -0.011 & 0.005 & 0.022 & -0.033 & 0.011 & 0.027 & 0.005 & -0.016 & 0.022 \end{bmatrix} \end{aligned}

‣

Linear projection forward

The activated hidden matrix $H$ is multiplied by the weight matrix $W_2$ and the bias $b_2$ is added. This projects the representation back down from the expanded dimension $d_{ff}$ to the model dimension $d_m$ .

U = H W_2 + b_2

H \approx \begin{bmatrix} 0.136 & -0.091 & 0.027 & 0.065 & -0.038 & 0.043 & 0.011 & -0.016 & 0.033 & -0.027 & 0.016 & 0.038 & -0.011 & 0.022 & 0.027 & -0.005 \\ -0.152 & 0.119 & -0.032 & -0.077 & 0.051 & -0.061 & -0.015 & 0.025 & -0.020 & 0.030 & -0.010 & -0.025 & 0.016 & -0.022 & 0.038 & 0.011 \\ 0.065 & 0.027 & -0.109 & 0.097 & 0.016 & -0.038 & 0.027 & -0.011 & 0.005 & 0.022 & -0.033 & 0.011 & 0.027 & 0.005 & -0.016 & 0.022 \end{bmatrix}, \\ W_2 = \begin{bmatrix} 0.05 & -0.10 & 0.02 & 0.03 \\ 0.10 & 0.05 & -0.02 & 0.01 \\ -0.05 & 0.15 & 0.04 & -0.03 \\ 0.20 & -0.05 & 0.01 & 0.05 \\ -0.10 & 0.05 & 0.03 & -0.02 \\ 0.05 & 0.10 & -0.01 & 0.02 \\ 0.15 & -0.10 & 0.02 & 0.04 \\ -0.05 & 0.05 & 0.01 & -0.03 \\ 0.03 & -0.02 & 0.05 & 0.01 \\ -0.01 & 0.04 & -0.05 & 0.03 \\ 0.02 & 0.01 & 0.03 & -0.04 \\ 0.04 & -0.03 & 0.02 & 0.05 \\ -0.02 & 0.05 & 0.01 & -0.01 \\ 0.01 & -0.01 & 0.04 & 0.02 \\ 0.05 & 0.02 & -0.03 & 0.01 \\ -0.03 & 0.03 & 0.02 & -0.02 \end{bmatrix}, \\ b_2 = \mathbf{0} \in \mathbb{R}^{4}

U = \begin{bmatrix} 0.02279 & -0.01995 & 0.00886 & 0.00985 \\ -0.02222 & 0.02376 & -0.01032 & -0.01044 \\ 0.02966 & -0.03084 & -0.00157 & 0.01248 \end{bmatrix}

Note that GPT-2 used dropout after attention and in the MLP, but we omit it here for clarity. In modern models, dropout is uncommon, as scale and other regularization methods suffice to avoid overfitting.

‣

Residual connection (post MLP)

The MLP output is added back to its input $Y^{(2)}$ , preserving the original signal while allowing the new transformation to affect it. This gives the final output of the transformer block.

R^{(2)} = U + Y^{(2)}

U = \begin{bmatrix} 0.02279 & -0.01995 & 0.00886 & 0.00985 \\ -0.02222 & 0.02376 & -0.01032 & -0.01044 \\ 0.02966 & -0.03084 & -0.00157 & 0.01248 \end{bmatrix}, \quad Y^{(2)} = \begin{bmatrix} 1.414 & -1.412 & 0.002 & -0.006 \\ -1.414 & 1.666 & -0.012 & -0.244 \\ 0.952 & 0.288 & -1.728 & 0.485 \end{bmatrix} \\ R^{(2)} = U + Y^{(2)} \approx \begin{bmatrix} 1.437 & -1.432 & 0.0109 & 0.0039 \\ -1.436 & 1.689 & -0.0223 & -0.2544 \\ 0.9817 & 0.2572 & -1.7296 & 0.4975 \end{bmatrix}

‣

Logit projection

The final hidden states are projected into the vocabulary space by multiplying with the embedding matrix transpose, producing logits for each token position.

\text{logits} = R^{(2)} E^{\top}

\begin{aligned} R^{(2)} &\approx \begin{bmatrix} 1.437 & -1.432 & 0.0109 & 0.0039 \\ -1.436 & 1.689 & -0.0223 & -0.2544 \\ 0.9817 & 0.2572 & -1.7296 & 0.4975 \end{bmatrix}, \\[1mm] E &= \begin{bmatrix} 0.120 & -0.050 & 0.030 & 0.025 \\ -0.080 & 0.090 & -0.020 & -0.015 \\ 0.050 & 0.020 & -0.070 & 0.040 \end{bmatrix}, \\[1mm] E^\top &= \begin{bmatrix} 0.120 & -0.080 & 0.050 \\ -0.050 & 0.090 & 0.020 \\ 0.030 & -0.020 & -0.070 \\ 0.025 & -0.015 & 0.040 \end{bmatrix}, \\[1mm] \text{logits} &= R^{(2)} E^\top \approx \begin{bmatrix} 0.244465 & -0.244116 & 0.042603 \\ -0.263799 & 0.271152 & -0.046635 \\ 0.065493 & -0.028259 & 0.195201 \end{bmatrix} \in \mathbb{R}^{3 \times 3} \end{aligned}

The output projection ties weights with the input embedding matrix, so logits are computed as $\text{logits} = R^{(2)} E^{\top}$ instead of learning a separate output matrix.

‣

Softmax

The softmax function converts the logits into probabilities over the vocabulary, ensuring each row sums to 1; this makes the model output interpretable as next-token probabilities for inference and usable in cross-entropy loss during training.

\begin{aligned} p_{ij} &= \frac{\exp(\text{logits}_{ij})}{\sum_{k=1}^{3} \exp(\text{logits}_{ik})}, \quad i = 1,2,3,\ j = 1,2,3 \\[1mm] p &\approx \begin{bmatrix} 0.4114 & 0.2524 & 0.3362 \\ 0.2532 & 0.4323 & 0.3146 \\ 0.3280 & 0.2986 & 0.3734 \end{bmatrix} \end{aligned}

First token (AI) → next token probabilities [0.41, 0.25, 0.34] correspond to ["AI", "rocks", "<EOS>"]. The model most likely predicts AI (0.411).
Second token (rocks) → next token probabilities [0.25, 0.43, 0.31] suggest the model is most confident that the next token is rocks (0.432).
Third token (<EOS>) → probabilities [0.33, 0.30, 0.37] indicate the next token is most likely <EOS> (0.3362).
During inference, sampling or argmax is applied to the softmax probabilities. During training, cross-entropy loss is computed between the softmax distribution and the target tokens.
With the softmax probabilities computed, the forward pass is complete, and the next step is to introduce the loss function and trace gradients backward through the network during backpropagation.

References

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. In Advances in neural information processing systems (Vol. 30).

Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., & Sutskever, I. (2019). Language models are unsupervised multitask learners. OpenAI Blog, 1(8), 9.

Next page

How LLMs learn (3 of 4): Back propagation

Last page

How LLMs learn (1 of 4): Tokenization

Home

Psychometrics.ai

This work is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0).

How LLMs learn (2 of 4): Forward pass