The Transformer Architecture

Attention and Self-Attention built the mechanism that lets each position consult every other. A mechanism is not a model. Between attention and a working GPT sit four more components, each solving a concrete problem that arises when you stack attention layers deeply, plus a question of how they are wired together.

This article assembles the block, works layer normalization by hand, and finishes by counting the parameters of GPT-2 small - arriving at the published figure.

A. What the block has to contain

Raschka enumerates the pieces required to build the GPT architecture: the GPT backbone, layer normalization, the GELU activation, the feed forward network, shortcut connections, the transformer block assembling them, and the final architecture. Attention is the mechanism inside the block; the rest is what makes a deep stack of them trainable.

B. Multi-head attention

A single attention operation produces one weighting over positions, expressing one notion of relevance. Multi-head attention runs several in parallel, each with its own $W_q, W_k, W_v$, then concatenates and projects the results.

The heads do not each get the full width. In GPT-2 small the embedding dimension is 768 and there are 12 heads, so each head works in

$$d_{\text{head}} = \frac{768}{12} = 64$$

dimensions. Twelve 64-dimensional heads concatenate back to 768. The cost is therefore the same as one 768-dimensional head, while the model gains twelve independent relevance patterns instead of one - different heads can attend to different relationships without competing for a single weighting.

C. Layer normalization

Deep stacks are unstable to train when activations drift in scale between layers. Layer normalization rescales each vector to zero mean and unit variance:

$$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad \mu = \frac{1}{n}\sum_i x_i, \qquad \sigma^2 = \frac{1}{n}\sum_i (x_i - \mu)^2 ,$$

with a small $\epsilon$ guarding against division by zero. Learned scale and shift parameters follow, so the layer can undo the normalization if that turns out to be useful.

Worked, on an 8-dimensional vector. Take

$$\mathbf{x} = \begin{bmatrix}2 & 4 & 4 & 4 & 5 & 5 & 7 & 9\end{bmatrix} .$$

The sum is $40$, so $\mu = 40/8 = 5$. The squared deviations are $9, 1, 1, 1, 0, 0, 4, 16$, summing to $32$, so $\sigma^2 = 32/8 = 4$ and $\sigma = 2$. Subtracting the mean and dividing by the standard deviation:

$$\hat{\mathbf{x}} = \begin{bmatrix}-1.5 & -0.5 & -0.5 & -0.5 & 0 & 0 & 1 & 2\end{bmatrix} .$$

The result has mean $0$ and variance $1$, as required.

Across features, not across the batch. Raschka draws the contrast with batch normalization explicitly: batch norm normalizes across the batch dimension, while layer norm normalizes across the feature dimension. The distinction matters for language models, where sequences in a batch differ in length and content - layer norm treats each position independently and so does not care what else is in the batch.

Where the normalization goes is itself a design decision with a history. Raschka notes that Post-LayerNorm, used in the original transformer, applies normalization after the attention and feed-forward sublayers, whereas Pre-LayerNorm, adopted in GPT-2 and newer LLMs, applies it before them - which can lead to more stable training dynamics.

D. Shortcut connections

Raschka's framing is direct: shortcut connections - also called skip or residual connections - were originally proposed for deep computer-vision networks (specifically residual networks) to mitigate vanishing gradients. The operation is addition: a layer's input is added to its output, creating an alternate path that bypasses the layer.

$$\mathbf{y} = \mathbf{x} + F(\mathbf{x})$$

The effect is measurable, and Raschka measures it. In a five-layer network he reports the mean absolute gradient at each layer with and without shortcuts:

Without shortcuts the gradient reaching layer 1 is about $0.0002$ - three orders of magnitude smaller than at layer 5, so the earliest layers barely learn. With shortcuts the gradients stay comparable throughout. The reason is visible in the derivative: differentiating $\mathbf{y} = \mathbf{x} + F(\mathbf{x})$ gives $1 + F'(\mathbf{x})$, and that $1$ is a path along which gradient flows undiminished no matter how small $F'$ becomes.

E. The feed-forward network

After attention mixes information across positions, a feed-forward network transforms each position independently. Its shape is distinctive: Raschka describes the inputs expanding by a factor of four, from 768 to 3,072 values, with a second layer compressing the 3,072 back down to 768.

$$768 \;\xrightarrow{\;\text{expand}\;}\; 3072 \;\xrightarrow{\;\text{GELU}\;}\; 3072 \;\xrightarrow{\;\text{contract}\;}\; 768$$

The activation between them is GELU rather than the ReLU from What Is a Neural Network?. Raschka contrasts them by noting ReLU is a piecewise linear function, whereas GELU is smooth - it has a non-zero gradient for negative inputs instead of ReLU's flat zero, which avoids the dead-unit failure mode.

Note the parameter asymmetry this creates: the feed-forward network holds roughly twice as many parameters as the attention it follows, which surprises people who think of attention as "the" transformer.

F. Counting GPT-2 small

Raschka scales up to the smallest GPT-2, 124 million parameters, and adds a useful bibliographic correction: while the original report mentioned 117 million, this was later corrected. The configuration table gives the family:

with vocab_size 50,257 and context_length 1,024.

We can check the headline number. Per block, with $d = 768$:

• Attention - four projections ($W_q, W_k, W_v$, output), each $768\times768$ plus a bias: $4(768^2 + 768) = 2{,}362{,}368$.
• Feed-forward - $768 \times 3072 + 3072$ expanding, $3072 \times 768 + 768$ contracting: $4{,}722{,}432$.
• Two layer norms - scale and shift, $768$ each: $4 \times 768 = 3{,}072$.

That is $7{,}087{,}872$ per block, and $\times 12$ blocks gives $85{,}054{,}464$. After the last block sits one more layer norm, the final normalization of section G, with its own scale and shift: $2 \times 768 = 1{,}536$. Adding the embeddings from Tokenization and Embeddings - token $50{,}257 \times 768 = 38{,}597{,}376$ and positional $1{,}024 \times 768 = 786{,}432$:

$$85{,}054{,}464 + 1{,}536 + 38{,}597{,}376 + 786{,}432 = 124{,}439{,}808 .$$

124.4 million - the published figure, and to the parameter the count of the released GPT-2 small checkpoint. The output layer adds nothing, because GPT-2 reuses the token embedding weights there.

Where the parameters actually live. The embeddings are 39.4 million of the 124 million - nearly a third - and they are a lookup table, not computation. Meanwhile the 12 blocks that do all the work hold 85 million. This ratio shifts sharply with scale: emb_dim and n_layers both grow in the larger models, so block parameters grow roughly quadratically while the token embedding grows only linearly.

G. The assembled block

Putting the pieces in Pre-LayerNorm order, each block is two sublayers, each normalized before and wrapped in a shortcut:

$$\mathbf{x} \leftarrow \mathbf{x} + \operatorname{MultiHeadAttention}(\operatorname{LayerNorm}(\mathbf{x}))$$ $$\mathbf{x} \leftarrow \mathbf{x} + \operatorname{FeedForward}(\operatorname{LayerNorm}(\mathbf{x}))$$

The shape never changes: a block maps a $(\text{batch},\ \text{sequence},\ 768)$ tensor to another of exactly the same shape. That invariance is what lets the blocks stack - twelve of them for GPT-2 small, forty-eight for GPT-2 XL, with no other structural change. The full model is embeddings, then $N$ identical blocks, then a final normalization and a projection to vocabulary size, giving one score per token for the next position.

Because the attention inside is causally masked, position $i$ can only attend to positions $\le i$, so a single forward pass produces a next-token prediction at every position at once - which is exactly what makes the training objective in the next article efficient.

Key takeaways

• Multi-head attention splits the embedding across heads - $768/12 = 64$ per head in GPT-2 small - so several relevance patterns are learned at one head's cost.
• Layer normalization rescales each vector to zero mean and unit variance across the feature dimension, unlike batch norm's batch dimension.
• The original transformer used Post-LayerNorm; GPT-2 and later models use Pre-LayerNorm for more stable training.
• Shortcut connections add a layer's input to its output; the resulting $1 + F'$ derivative keeps gradients alive - 0.0002 versus 0.22 at layer 1 in Raschka's measurement.
• The feed-forward network expands $768 \to 3072 \to 768$ and holds about twice the parameters of the attention beside it.
• Counting the blocks, the final layer norm and the embeddings gives 124,439,808 parameters for GPT-2 small, matching the published 124M (the original report's 117M was later corrected).
• A block preserves its input shape, which is precisely why the architecture scales by stacking.

What's next

We have an architecture with 124 million randomly initialised parameters, which predicts nothing. Turning it into a model that writes fluent text, and then into one that follows instructions, is the subject of Pretraining and Fine-Tuning.