The Transformer Block

The architecture diagram, piece by piece

Open a diagram of a large language model and it looks like a wall of boxes stacked impossibly high, GPT-3 stacks 96 of them. It's tempting to assume each box does something fundamentally different, that layer 40 is computing something layer 41 isn't. It isn't. Every one of those boxes is the exact same unit, repeated, each copy with its own independently learned weights. That repeating unit is the transformer block, and once you've understood one, you've understood the architecture.

A transformer block has exactly two computational sublayers. Multi-head attention, the mechanism from the previous article, which lets every token pull information from every other token. And a feedforward network, covered later in this one, which processes what attention gathered. Each sublayer gets wrapped in the same packaging, a layer norm before it, and a residual connection around it that adds the sublayer's input back onto its output. Follow one token's vector through a single block and it passes through the same eight steps in the same order: layer norm, attention, add the residual, layer norm again, feedforward, add the residual again.

The diagram below is that exact sequence, laid out top to bottom. Click through each stage before the rest of this article covers attention and feedforward in depth, it's worth having the whole shape in view first.

Notice two stages show up twice, layer norm and the residual add. That repetition isn't an accident. It's the identical fix applied around each sublayer independently, and the last section of this article covers exactly why a stack this deep needs that fix at all.

Multi-head attention: running several attentions in parallel

The self-attention from the previous article used one query matrix, one key matrix, one value matrix, learned once and applied to the whole embedding at a time. That's a real, working attention mechanism. It's also a bottleneck: a single set of weights can only learn a single notion of “relevant.” Force every relationship in a sentence, syntax, coreference, negation, tone, through one shared lens and most of them get blurred together.

Multi-head attention fixes this by splitting the embedding dimension into several smaller slices, called heads, and running an independent attention computation on each one. A model with a 512-dimensional embedding and 8 heads doesn't run 8 full 512-dimensional attentions, it slices each token's vector into 8 pieces of 64 dimensions, gives each slice its own learned query, key, and value matrices, and runs the same scaled dot-product attention from the previous article on each slice separately. The 8 results, each still 64-dimensional, get concatenated back into one 512-dimensional vector, then passed through one more learned matrix that mixes the heads back together into a single output per token.

Each head ends up specializing on its own, nobody assigns it a job. One head might learn to track which word modifies which noun, a syntax pattern. Another might learn coreference, tracking which pronoun points back to which earlier noun. A third might latch onto something with no clean grammatical name at all. The specialization falls out of gradient descent the same way embedding clusters did in an earlier article, driven purely by what reduces the loss.

1def multi_head_attention(x, num_heads, d_model):
2    d_k = d_model // num_heads
3    heads = []
4    for h in range(num_heads):
5        q = x @ W_q[h]   # project to this head's (seq_len, d_k) slice
6        k = x @ W_k[h]
7        v = x @ W_v[h]
8        heads.append(attention(q, k, v))
9    concatenated = concat(heads, axis=-1)   # back to (seq_len, d_model)
10    return concatenated @ W_o   # mix the heads back together

The feedforward sublayer: where the model actually “thinks”

It's tempting to assume attention is where the real thinking happens, it's the piece with the evocative name and the worked example you can trace by hand. But attention only moves information between positions. Underneath the softmax weighting, it's a weighted average of value vectors, a linear combination. Recall from the neural network foundations article: stacking linear operations with nothing nonlinear between them collapses into a single linear operation, no matter how many attention heads you run in parallel. Attention alone can gather, it can't transform.

The feedforward sublayeris where that transformation happens. It's two linear layers with a nonlinear activation between them, usually ReLU in older models or GELU in most modern ones, the same activation functions from the neuron playground in an earlier article. The first layer expands each token's vector up, typically to about 4 times the model's dimension, the second layer projects it back down to the original size. A 512-dimensional model routes each token through a 2,048-dimensional hidden layer and back.

Think of attention as a meeting: every token gathers notes from everyone else in the room. The feedforward layer is each token going back to its own desk afterward and doing something with what it heard, alone. Every desk runs the identical process, the same two weight matrices, applied independently and in parallel to every position. That's the “position-wise” part of the name you'll see in papers: no token's feedforward computation ever looks at another token's vector, all the cross-token mixing already happened one step earlier, in attention.

1def feedforward(x, d_model, d_ff):
2    # x: one token's vector, shape (d_model,). d_ff is usually 4 * d_model.
3    hidden = relu(x @ W1 + b1)   # up-project: (d_model,) -> (d_ff,)
4    return hidden @ W2 + b2      # down-project: (d_ff,) -> (d_model,)

In most transformers, this sublayer holds the majority of the model's parameters, two matrices of roughly d_model by 4 × d_model each, comfortably outweighing the attention weights sitting right next to them in the same block. Scaling a model up, a later article in this series covers this directly, is largely a story about scaling this piece.

Residual connections and layer norm: why deep stacks don't collapse

The obvious way to wire two sublayers together is to just replace: feed the input into attention, take attention's output, feed that straight into feedforward, pass the result to the next block. Stack that 96 times and training breaks. A gradient flowing backward through 96 consecutive transformations shrinks or explodes long before it reaches the earliest blocks, the same vanishing-gradient problem that made very deep networks untrainable before this fix existed.

The fix is the residual connection: add the sublayer's input back onto its output instead of replacing it, so output = input + sublayer(input). The sublayer now only has to learn the correction to make, an increment on top of what came in, rather than learning to reconstruct and preserve everything unrelated while also computing something new. And that addition has a gradient of exactly 1 with respect to the input, so the gradient signal always has a direct, unobstructed path backward through every one of those 96 additions. The sublayers themselves become optional detours on the gradient's shortest way home, not the only road.

Layer normalizationhandles a different problem the residual stream creates. Every addition piles another sublayer's output on top of the running total, and unnormalized, that total drifts to arbitrarily large magnitudes 96 additions deep. Layer norm takes one token's vector, at that point in the stack, and rescales it to zero mean and unit variance across its own feature dimension, then applies a small learned scale and shift so it isn't stuck at a fixed statistical shape forever. It normalizes per token, across features, not across a batch of examples the way batch normalization does, which is what makes it work regardless of how many tokens happen to be in the sequence.

One ordering choice is worth naming: the original transformer paper normalized after each sublayer's addition, now called post-norm. Most modern large models, and the diagram earlier in this article, normalize before each sublayer instead, pre-norm, because it keeps that gradient path even cleaner at very large depths and trains more stably without careful learning-rate warmup.

Everything in this article assumed each token's vector already carries some notion of where it sits in the sequence. It doesn't, not on its own, an embedding by itself has no position baked in, and attention treats the whole sequence as an unordered set unless something adds order back in. Where that positional information actually comes from is next.