The Transformer Architecture

The block

Every 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 remember as "sublayer, add, norm" repeated twice than as a diagram.

What each piece is for

• Attention moves information between positions. It is the only part of the block that lets one token see another.
• The feed-forward network processes each position independently — the same small two-layer MLP applied to every token separately. It is where the model does per-token computation on whatever attention just gathered.
• The residual connection gives the gradient a path that skips the sublayer entirely.
• Layer normalisation keeps the activations 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 information between positions. Each token gathers a content-weighted mix of the others.

The feed-forward network processes each position independently. 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 activations stable, and placing it before each sublayer rather than after is what makes 80-layer stacks trainable.

Everything 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 every token's representation is informed by both directions. That makes them excellent at understanding and useless at generating — there is no notion of "next".

Decoder-only models mask future positions, so each token depends only on what precedes it. That constraint is what makes generation possible, 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

Following 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 information — sinusoidal, learned, or rotary applied inside attention.
4. N transformer blocks — each refining the representation.
5. Final LayerNorm.
6. Output projection to vocabulary size, giving one logit per possible next token.
7. Softmax (inside the loss during training; explicitly at generation time).

Step 6 is often tied to the embedding table — the same matrix used in reverse. That saves a large number of parameters and usually improves quality slightly.

The shape of the tensors flowing through: (batch, sequence, model_dim) from step 2 onwards, unchanged by every 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 actually 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()

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()

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()

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")
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 every 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, which is why widening is so much more expensive than deepening.

Why residuals matter so much

Without a residual connection, the gradient reaching layer 1 has to pass back through every layer above it, being multiplied at each step. Multiply twenty-four numbers smaller than one and you get approximately 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 essentially intact. That is why the same trick appears in ResNets, and why almost every deep architecture 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 original paper put normalisation after the addition — LayerNorm(x + Sublayer(x)), which is what the diagram above shows. Almost every 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 every layer, which 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 difference 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 computation, not attention.
• Forgetting the FFN is position-wise. It has no view of other tokens at all. Every bit of cross-token information 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 information between positions, and a position-wise feed-forward network processes each token on its own. The residual connection is what makes depth possible, giving the gradient a path that skips each sublayer — remove it and a deep stack stops training entirely — while layer normalisation holds the activation scale steady as layers accumulate. 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 position t before t+1. A transformer computes all positions simultaneously as one matrix multiplication, which is what a GPU is built for. That alone made training on far more data feasible.

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

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

Note that recurrent architectures have not disappeared — state-space models such as Mamba revisit the idea with linear scaling, and are competitive 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 — performance improving predictably with compute, parameters and data together — are why the field spent several years mostly making the same architecture bigger. The Chinchilla result refined it: for a given compute budget, most early large models were under-trained on too little data relative to their size.

Practical consequences 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-efficient methods such as LoRA rather than updating every weight.

Questions people ask

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

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

Do transformers understand language? They model statistical structure in text extremely well. Whether that constitutes understanding 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 generalises. Current models reach hundreds of thousands of tokens.

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

Recap in one screen

• One block: attention to move information between positions, a feed-forward network to process each position, both wrapped in residual connections and LayerNorm.
• Most parameters are in the feed-forward layers, not in attention.
• Encoder-only for understanding, 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.