The Transformer Architecture

Quick Context

The last modules built the parts: query, key and value, self-attention, multiple heads, positional encoding. This module assembles them.

The surprise, when you finally see the block written out, is how little there is. Two sublayers, each wrapped in the same two-line pattern, repeated N times.

The first sublayer is multi-head self-attention, which lets each spot in the sequence attend to each other spot and collect data weighted by relevance. The second sublayer is a feed-forward network. Two linear layers with a ReLU (or GELU) in between. Applied on its own to each spot. The attention sublayer mixes data across spots. The feed-forward sublayer processes each spot in isolation. Stacking the two is what gives the transformer both global context and per-spot work.

Each sublayer is wrapped in a residual connection and a layer normalisation. The residual connection adds the sublayer's input back to its output. This keeps gradients flowing through deep stacks. The layer normalisation rescales the result so that each layer receives inputs in a stable range. Without these two wrappers, a transformer with more than a few blocks becomes difficult to train.

The first 2017 paper ("Attention Is All You Need") described an encoder-decoder design with six blocks on each side. Modern models are often encoder-only (BERT, for classification and embeddings) or decoder-only (GPT, for generation). The decoder-only variant uses masked self-attention, where each spot can attend only to earlier spots. This is what makes left-to-right generation doable. The architectural gap between BERT and GPT is whether the attention mask allows the model to look forward.

The transformer replaced recurrent networks because it processes all spots in parallel during training, rather than sequentially. That parallelism is why transformers train faster on modern hardware, and why they have scaled to billions of parameters where RNNs could not.

The block

Each transformer block is these four lines:

x = LayerNorm( x + MultiHeadAttention(x) )

x = LayerNorm( x + FeedForward(x) )

That is the entire thing. Two sublayers, and each is wrapped identically: run the sublayer, add the input back, then normalise. The pattern is called Add & Norm. And it is applied so uniformly that the block is easier to recall as "sublayer, add, norm" repeated twice than as a diagram.

What each piece is for

• Attention moves data between spots. It is the only part of the block that lets one token see another.
• The feed-forward network processes each spot on its own — the same small two-layer MLP applied to each token apart. It is where the model does per-token work on any attention just gathered.
• The residual connection gives the gradient a path that skips the sublayer fully.
• Layer normalisation keeps the triggers at a stable scale as depth grows.

The division of labour between the first two is worth holding on to. Attention mixes across tokens, the FFN thinks about each token. A transformer alternates between the two, over and over.

The block, repeated

A transformer is one block stacked many times. The block has two sublayers, each wrapped in a residual connection and a normalisation:

x → LayerNorm → Multi-head attention → (+x)

x → LayerNorm → Feed-forward → (+x)

Attention moves data between spots. Each token gathers a content-weighted mix of the others.

The feed-forward network processes each spot on its own. It expands to four times the model dimension, applies a non-linearity, and projects back. This is where most of the parameters live. Roughly two thirds of them in a typical model. And it is a common misconception that attention does all the work.

Residual connections give gradients a path back through dozens of blocks. LayerNorm keeps triggers stable, and placing it before each sublayer rather than after is what makes 80-layer stacks trainable.

All of it else about a transformer is the arrangement of these blocks and what goes in and out.

Three arrangements, three jobs

Encoder-only models see the whole input at once, so each token's form is informed by both directions. That makes them excellent at grasp and useless at making. There is no notion of "next".

Decoder-only models mask future spots, so each token depends only on what precedes it. That constraint is what makes generation doable. And it turned out to scale so well that most current large models are decoder-only.

Encoder-decoder models encode the source, then decode while attending to it through cross-attention. Natural for translation, where input and output are distinct sequences.

The full input path

Below a sentence through a decoder-only model:

1. Tokenise — text becomes integer ids.
2. Embed — each id looks up a vector from the embedding table.
3. Add positional data — sinusoidal, learned, or rotary applied inside attention.
4. N transformer blocks — each refining the form.
5. Final LayerNorm.
6. Output projection to vocabulary size, giving one logit per doable next token.
7. Softmax (inside the loss during training; by name at generation time).

Step 6 is often tied to the embedding table. The same matrix used in reverse. That saves many parameters and often improves quality slightly.

The shape of the tensors flowing through: (batch, sequence, model_dim) from step 2 onwards, unchanged by each block. That constancy is what makes the stack composable.

A whole block, run end to end

A transformer block is attention, a feed-forward network, two residual connections and two layer norms. All of it is here, followed by the parameter count that explains where the compute really goes.

import numpy as np

rng = np.random.default_rng(0)
T, D, H, FF = 6, 32, 4, 128
D_HEAD = D // H
x = rng.normal(0, 1.0, (T, D))

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def layernorm(a, g=None, b=None, eps=1e-5):
    mu = a.mean(-1, keepdims=True)
    sd = a.std(-1, keepdims=True)
    out = (a - mu) / (sd + eps)
    return out if g is None else out * g + b

Wq = rng.normal(0, 0.3, (H, D, D_HEAD))
Wk = rng.normal(0, 0.3, (H, D, D_HEAD))
Wv = rng.normal(0, 0.3, (H, D, D_HEAD))
Wo = rng.normal(0, 0.3, (D, D))
W1 = rng.normal(0, 0.2, (D, FF)); b1 = np.zeros(FF)
W2 = rng.normal(0, 0.2, (FF, D)); b2 = np.zeros(D)

def gelu(z):
    return 0.5 * z * (1 + np.tanh(np.sqrt(2 / np.pi) * (z + 0.044715 * z ** 3)))

print("PRE-NORM block, which is what modern models use:")
print("    x = x + Attention(LayerNorm(x))")
print("    x = x + FeedForward(LayerNorm(x))")
print()
print("input %s, %d heads of %d dims each, feed-forward width %d."
      % (x.shape, H, D_HEAD, FF))
print()

Building on the same setup, the next block shows the left cases.

import numpy as np

def block(x):
    trace = {}
    # --- sublayer 1: attention
    a = layernorm(x)
    heads = []
    for h in range(H):
        Q, K, V = a @ Wq[h], a @ Wk[h], a @ Wv[h]
        A = softmax(Q @ K.T / np.sqrt(D_HEAD))
        heads.append(A @ V)
    attn = np.concatenate(heads, axis=1) @ Wo
    trace["attn_out"] = attn
    x = x + attn                                   # residual 1
    trace["after_attn"] = x
    # --- sublayer 2: feed-forward
    f = layernorm(x)
    ff = gelu(f @ W1 + b1) @ W2 + b2
    trace["ff_out"] = ff
    x = x + ff                                     # residual 2
    trace["after_ff"] = x
    return x, trace

out, tr = block(x)
print("%22s %14s %14s %14s"
      % ("stage", "shape", "mean", "sd"))
for name, arr in (("input x", x), ("attention output", tr["attn_out"]),
                  ("after residual 1", tr["after_attn"]),
                  ("feed-forward output", tr["ff_out"]),
                  ("after residual 2", tr["after_ff"])):
    print("%22s %14s %14.4f %14.4f" % (name, str(arr.shape), arr.mean(), arr.std()))
print()
print("the shape never changes. that is what lets blocks stack -- the output")
print("of one is a valid input to the next, forever.")
print()

The below block covers the left steps.

import numpy as np

print("WHAT EACH SUBLAYER IS FOR:")
print("   attention     mixes information ACROSS positions. it is the only")
print("                 part of the block that does.")
print("   feed-forward  transforms each position INDEPENDENTLY. the same")
print("                 two matrices are applied to every token.")
print()
print("   proof of the second claim -- shuffle the rows and run the")
print("   feed-forward alone:")
perm = np.array([4, 1, 0, 5, 3, 2])
f = layernorm(x)
ff_normal = gelu(f @ W1 + b1) @ W2 + b2
ff_shuffled = gelu(layernorm(x[perm]) @ W1 + b1) @ W2 + b2
print("      identical up to the permutation: %s"
      % np.allclose(ff_normal[perm], ff_shuffled))
print("   so a transformer alternates: mix across tokens, then think about")
print("   each token separately. every layer does both.")
print()

The next block goes on with the left demonstrations.

import numpy as np

print("THE PARAMETER COUNT, and where it goes:")
attn_p = 4 * D * D                       # Q, K, V and the output projection
ff_p = 2 * D * FF + FF + D
ln_p = 2 * 2 * D
total = attn_p + ff_p + ln_p
print("%26s %12s %10s" % ("component", "parameters", "share"))
for name, n in (("attention (Q,K,V,O)", attn_p), ("feed-forward", ff_p),
                ("layer norms", ln_p)):
    print("%26s %12d %9.1f%%" % (name, n, 100 * n / total))
print("%26s %12d" % ("total per block", total))
print()
print("   the feed-forward layer holds two thirds of the parameters, and it")
print("   is the part people think about least. the standard ratio is")
print("   FF = 4 x d_model, which is why:")
for d in (768, 4096):
    print("      d_model %5d -> feed-forward width %5d, %s parameters"
          % (d, 4 * d, "{:,}".format(2 * d * 4 * d)))
print()

print("BUT ATTENTION DOMINATES THE COMPUTE, because its cost grows with")

Building on the same setup, the next block shows the left cases.

print("sequence length and the feed-forward's does not:")
print("%10s %20s %20s %14s" % ("tokens", "attention FLOPs", "feed-forward FLOPs", "attn share"))
for t in (128, 512, 2048, 8192):
    a_f = 4 * t * D * D + 2 * t * t * D          # projections + the T x T part
    f_f = 4 * t * D * FF
    print("%10d %20s %20s %13.0f%%"
          % (t, "{:,}".format(a_f), "{:,}".format(f_f), 100 * a_f / (a_f + f_f)))
print("   at short lengths the feed-forward is the expensive part. as the")
print("   context grows the T x T term takes over, and that crossover is")
print("   why long-context work is all about attention and never about the")
print("   feed-forward.")
print()

print("STACKING. run the same block 6 times and watch the residual stream")
print("grow while every sublayer INPUT stays normalised:")
h = x.copy()
for i in range(6):
    h, _ = block(h)
    print("   after block %d: mean %+8.4f, sd %8.4f, max |value| %8.4f"
          % (i + 1, h.mean(), h.std(), np.abs(h).max()))
print("   the residual stream grows -- each block ADDS to it -- while the")
print("   layer norms keep each sublayer's INPUT on scale. that is the")
print("   arrangement that makes 100-layer models trainable.")

Try it yourself

1. Take the residuals away. Untick Residual Connections and set the Number of Layers slider to 24. Watch the gradient curve decay toward zero and Trainable flip to no. Each sublayer shrinks what passes through it, and twenty-four of those in a row leave nothing to learn from.
2. Put them back. Tick Residual Connections with the depth still at 24. The gradient holds at 1 all the way back to layer 1. This is the single change that makes deep stacks trainable at all.
3. Take normalisation away instead. Untick Layer Normalisation with residuals on and depth at 24. Now the opposite failure: with nothing rescaling between layers the gradient compounds instead of decaying. And a deep stack becomes unstable rather than dead.
4. Find where the parameters are. With Feed-Forward Width at 4, read FFN Share. Roughly two thirds of each block sits in the feed-forward network, not in attention — which surprises most people. Since attention is the part that gets all the attention.
5. Change the width. Set the Feed-Forward Width slider to 1 and watch FFN Share collapse. Then to 8 and watch it dominate. The 4× default is a convention, not a law.
6. Scale it up. Set the Model Dimension slider to its maximum and watch the total. Parameters grow with the square of the dimension but only linearly with depth. This is why widening is so much more costly than deepening.

Why residuals matter so much

Without a residual connection, the gradient reaching layer 1 has to pass back through each layer above it, being times at each step. Times twenty-four numbers smaller than one and you get about zero. The vanishing gradient problem, arriving from a different direction.

The residual gives the gradient a route that skips the sublayer. So it reaches the early layers in effect intact. That is why the same trick appears in ResNets, and why almost each deep design built since 2015 has some version of it. It is not an optimisation detail; without it, "deep" is not available.

Pre-norm and post-norm

The first paper put normalisation after the addition — LayerNorm(x + Sublayer(x)). This is what the diagram above shows. Almost each model since does the opposite, x + Sublayer(LayerNorm(x)), called pre-norm.

The reason is practical: with post-norm the residual path itself gets normalised at each layer. This weakens it, and deep post-norm models need a careful learning-rate warmup to train at all. Pre-norm leaves the residual path clean from input to output and trains far more forgivingly. If you read the 2017 paper and then read a modern implementation, this is the gap you will notice first.

Where this goes wrong

• Thinking the FFN is the small part. It is about two thirds of the block. Most of a transformer's parameters do per-token work, not attention.
• Forgetting the FFN is spot-wise. It has no view of other tokens at all. Each bit of cross-token data arrived through attention.
• Mixing up pre-norm and post-norm. Implementing the paper's post-norm without warmup and wondering why a deep model will not converge is a rite of passage.
• Expecting depth to be free. Layers add parameters linearly but add sequential compute too; width adds parameters quadratically but parallelises well. The trade is why model shapes look the way they do.

In one line

A transformer block is two sublayers, each wrapped in the same add-and-normalise pattern: attention moves data between spots. And a spot-wise feed-forward network processes each token on its own. The residual connection is what makes depth doable. This gives the gradient a path that skips each sublayer. Remove it and a deep stack stops training fully. While layer normalisation holds the trigger scale steady as layers build up. About two thirds of the parameters live in the feed-forward network. Rather than in attention, and parameters grow with the square of the model dimension but only linearly with depth.

Why it replaced recurrent networks

Two properties, both structural.

Parallelism. An RNN must compute spot t before t+1. A transformer computes all spots at once as one matrix product. This is what a GPU is built for. That alone made training on far more data feasible.

Constant path length. Any two spots are one attention step apart. In an RNN the path grows with the distance, and gradients decay along it —. This is why long-range dependencies were the central trouble.

The cost of those wins is quadratic attention: doubling the sequence length quadruples the work. That is the trade the design makes. And it is why long context needed years of engineering (FlashAttention, sparse patterns, KV caching) rather than being free.

Note that recurrent designs have not disappeared — state-space models such as Mamba revisit the idea with linear scaling, and are strong on long sequences.

Scaling, and what changes with size

The same block, repeated more times and made wider:

Nothing about the block changes. What changes is depth, width, vocabulary, context length and the amount of data.

The scaling laws that emerged from studying this. Speed improving predictably with compute, parameters and data together. Are why the field spent a few years mostly making the same design bigger. The Chinchilla result refined it: for a given compute budget, most early large models were under-trained on too little data next to their size.

Practical results of size: Adam's optimiser state is twice the model. So large-scale training needs sharding (ZeRO, FSDP). Inference needs KV caching and quantisation; and fine-tuning uses parameter-fast methods such as LoRA rather than updating each weight.

Questions people ask

Why is the feed-forward layer four times wider? Convention from the first paper, and it works. Some modern models use different ratios, and gated variants (SwiGLU) with a smaller expansion.

What is the gap between BERT and GPT? Bidirectional encoder trained on masked tokens, versus causal decoder trained on next-token prediction. Grasp versus generation.

Do transformers understand language? They model numeric structure in text very well. Whether that constitutes grasp is a genuine open question, not a settled one in either direction.

Why pre-norm rather than post-norm? It leaves the residual path clean. So gradients flow back unimpeded. Post-norm needs careful warm-up beyond about 12 layers.

How long can the context be? Architecturally unlimited; practically bounded by quadratic cost and by how far the positional scheme extends. Current models reach hundreds of thousands of tokens.

Are transformers only for text? No — vision transformers treat image patches as tokens. And the design is used for audio, protein structure, video and time series.

Recap in one screen

• One block: attention to move data between spots, a feed-forward network to process each spot, both wrapped in residual connections and LayerNorm.
• Most parameters are in the feed-forward layers, not in attention.
• Encoder-only for grasp, decoder-only for generation, encoder-decoder for sequence-to-sequence.
• Parallelism and constant path length are why it replaced RNNs; quadratic attention is what it costs.
• Scaling changes depth, width and data — not the block.