17  The Transformer

TipTL;DR

The Transformer wraps self-attention in a repeatable block: attention to mix information across positions, a position-wise feed-forward network to process each position, with residual connections and layer normalisation around both. Stacking \(N\) such blocks gives an encoder; adding a second attention that queries the encoder gives a decoder. A block’s update is \[x \leftarrow x + \text{Sublayer}\big(\text{LayerNorm}(x)\big),\] applied first to attention, then to the FFN. Vaswani et al. (2017) showed this stack — no recurrence, no convolution — sets a new state of the art on translation while training in a fraction of the time.

Depends on: Self-Attention & Multi-Head Attention

17.1 Why this matters

The previous chapter gave us one operation: multi-head self-attention. But one attention layer is not a model. Three problems remain before it becomes a trainable architecture:

  1. Attention only mixes; it does not transform. Its output is a weighted average of value vectors — a convex combination. On its own it cannot apply a non-linear per-token computation.
  2. Deep stacks do not train. Stacking attention layers naively reproduces the vanishing/exploding-gradient problem from the RNN era.
  3. Attention is permutation-equivariant. It has no built-in sense of order (handled in the next chapter).

The Transformer block solves (1) with a feed-forward network and (2) with residuals + normalisation. The result is the most copied architecture in machine learning: the same block, restacked and rescaled, is GPT, BERT, T5, ViT, and every LLM since.

17.2 The mechanism

17.2.1 The block: four ingredients

Each Transformer layer composes two sublayers, each wrapped identically:

\[ \text{Sublayer output} = \text{LayerNorm}\big(x + \text{Sublayer}(x)\big). \]

  • Multi-head self-attention — mixes information across positions.
  • Position-wise feed-forward network (FFN) — transforms each position independently and identically: \[\text{FFN}(x) = \max(0,\, xW_1 + b_1)\,W_2 + b_2,\] typically expanding to an inner dimension \(d_{\text{ff}} = 4\,d_{\text{model}}\) then projecting back. This is where most of a Transformer’s parameters live. The original uses ReLU; modern stacks swap in GELU and the gated SwiGLU — see Modern Activations.
  • Residual connection — the \(x +\) term gives gradients a direct path to every layer, the same identity-shortcut idea that let ResNets (He et al., 2016) go deep.
  • Layer normalisation — normalises across the feature dimension per token, stabilising the scale of activations regardless of sequence content.
NotePost-LN vs Pre-LN

The original paper applies LayerNorm after the residual add (post-LN): \(\text{LayerNorm}(x + \text{Sublayer}(x))\). Modern LLMs almost universally use pre-LN\(x + \text{Sublayer}(\text{LayerNorm}(x))\) — which keeps a clean residual highway and trains stably without learning-rate warmup. The TL;DR formula above is the pre-LN form; the difference is purely where the norm sits.

17.2.2 Encoder and decoder

Encoder ×N Multi-head self-attn Feed-forward + residual & LayerNorm around each sublayer Decoder ×N Masked self-attn (causal) Cross-attn (Q=dec, K,V=enc) Feed-forward K, V Linear → softmax Encoder reads the source in parallel; decoder generates left-to-right, attending to itself and the encoder.
  • Encoder\(N\) identical blocks (self-attention + FFN). Sees the whole input at once, no masking. Produces a contextual representation per token.
  • Decoder\(N\) blocks with three sublayers: masked self-attention (causal, so generation stays autoregressive), cross-attention (queries from the decoder, keys/values from the encoder output — the same encoder-decoder bridge inherited from RNN attention), then the FFN.

The three later families simply keep one half: encoder-only (BERT), decoder-only (GPT), or both (T5). All covered in BERT & GPT.

17.2.3 The full forward pass

Tokens → embeddings → add positional encoding → \(N\) encoder blocks. On the decoder side: shifted-right target embeddings → \(N\) decoder blocks → a final linear projection to vocabulary logits → softmax. The output embedding matrix is often tied to the input embedding to save parameters.

WarningPitfall: the residual sets the width

Every sublayer’s input and output must share dimension \(d_{\text{model}}\), because the residual adds them. This is why attention concatenates its heads back to \(d_{\text{model}}\) and the FFN expands then contracts. The residual stream is a fixed-width “bus” running the depth of the network; every block reads from and writes to it.

TipGoing deeper
  • Self-attention has no order; the embeddings step needs a position signal — Positional Encoding.
  • The encoder-only / decoder-only split and pre-training objectives — BERT & GPT.

17.2.4 Minimal sketch

A single pre-LN encoder block in NumPy — attention placeholder plus FFN, residuals, and layer norm — to show the wiring.

import numpy as np

rng = np.random.default_rng(0)
n, d_model, d_ff = 4, 8, 32

def layer_norm(x, eps=1e-5):
    mu = x.mean(-1, keepdims=True)
    var = x.var(-1, keepdims=True)
    return (x - mu) / np.sqrt(var + eps)

def ffn(x, W1, b1, W2, b2):
    return np.maximum(0, x @ W1 + b1) @ W2 + b2

W1 = rng.normal(scale=0.1, size=(d_model, d_ff)); b1 = np.zeros(d_ff)
W2 = rng.normal(scale=0.1, size=(d_ff, d_model)); b2 = np.zeros(d_model)

def attn(x):                      # stand-in for multi-head self-attention
    return x

def encoder_block(x):
    x = x + attn(layer_norm(x))               # sublayer 1: attention
    x = x + ffn(layer_norm(x), W1, b1, W2, b2) # sublayer 2: position-wise FFN
    return x

X = rng.normal(size=(n, d_model))
print("block output shape:", encoder_block(X).shape)  # (n, d_model) — width preserved
block output shape: (4, 8)

17.3 Application & impact

Component Role Lives on as
Self-attention sublayer Cross-position mixing The defining op of every Transformer
Position-wise FFN Per-token non-linear transform ~2/3 of LLM parameters; the MoE target
Residual stream Gradient highway, fixed-width bus Interpretability’s “residual stream” view
LayerNorm (→ pre-LN, RMSNorm) Activation scaling Every modern LLM (RMSNorm in LLaMA (Touvron et al., 2023))
Encoder / decoder split Bidirectional vs autoregressive BERT / GPT / T5 families
  • The result that ended the RNN era: Vaswani et al. (Vaswani et al., 2017) reached a new SOTA on WMT translation, training in ~12 hours on 8 GPUs where comparable recurrent models took days — parallelism over sequence length was the unlock.
  • The block is the unit of scaling. “Make the model bigger” almost always means more blocks and wider blocks of exactly this design — the premise of the scaling laws.
  • The FFN became the efficiency frontier: because it holds most parameters and runs per-token, it is what Mixture of Experts makes sparse.
NoteKey takeaway

The Transformer is not one idea but a recipe: attention to mix, FFN to transform, residual + norm to make it trainable at depth, stacked \(N\) times. Self-attention is the engine; this chapter is the chassis. Every model in the rest of this era is this block — kept whole, split in half, scaled up, or made sparse.

→ Next: Positional Encoding