Transformer Architecture
Self-attention, positional encoding, encoder vs decoder, pre-norm vs post-norm
Attention has a surprising blind spot. Because it looks at all the words at once and just computes weighted averages, it does not inherently know their *order* — "the dog bit the man" and "the man bit the dog" contain the exact same words, and to raw attention they look identical. So the first thing a Transformer has to add is a sense of *position*.
Positional encoding: telling the model where each word sits
The fix is to stamp each word's representation with a position signal before the first layer — a little pattern that says "I am word 1," "I am word 2," and so on. The original Transformer used sine and cosine waves of different frequencies for this (the original paper only hedged that this *may* let the model extrapolate to unseen lengths — in practice that extrapolation is weak); modern models like LLaMA use RoPE, which bakes *relative* position straight into the attention comparison. RoPE's relative encoding is friendlier to extrapolation than absolute position encodings, but base RoPE still degrades past the trained context length without added scaling tricks (position interpolation, NTK-aware/YaRN scaling). Either way, once positions are added, the model can finally tell word order apart.
The Transformer block
Stack the pieces and you get the repeating block every Transformer is built from: normalise the inputs, run multi-head attention (words look at each other), add the result back through a residual shortcut, normalise again, run a small two-layer feed-forward network (FFN) on each word, and add that back too. Two details matter. The FFN is deliberately *wide* — usually 4× the model's width in the middle — because it acts as the model's *memory*, where a lot of its factual knowledge is stored; shrinking it is reported in interpretability and scaling studies to cost the model retrievable facts (this module doesn't walk through a worked number for that drop, unlike the √d_k computation in the attention module — treat the direction as right, not the exact magnitude). And the residual shortcuts are not decoration: they give the gradient a direct path back to every layer — necessary, but not sufficient, for stacking dozens of these blocks without the signal dying (exactly the vanishing-gradient fix from earlier). Residuals give the gradient a path; whether that path stays well-scaled as depth grows still depends on where LayerNorm sits relative to it, which is why Pre-LN vs Post-LN (below) matters.
Two flavours: encoder and decoder
The same block comes in two modes, set by *who is allowed to look at whom*. Encoder-only (BERT-style) lets every word see every other word, in both directions — great for *understanding* tasks like classification, where you want the fullest possible context. Decoder-only (GPT-style) masks the future, so each word can only see the words *before* it — which is exactly what you need to *generate* text one token at a time. Decoder-only models also get a training bonus: *every* token in a sequence is a prediction target at once, giving them far more learning signal per pass than BERT's "predict just the 15% we masked," which is a big part of why decoder-only models dominate at scale.
Key points
- Encoder-only for understanding (classification, NER, QA); decoder-only for generation (language models, chat); encoder-decoder for sequence-to-sequence tasks (translation, summarisation). The architecture determines the training objective, which determines what the model can do. A decoder-only model cannot be directly used for tasks requiring bidirectional context (e.g., masked span filling) without changing its attention pattern. Use Pre-LN (LayerNorm before attention and FFN sub-layers) — every modern large model (LLaMA, GPT-3, PaLM) uses Pre-LN because Post-LN (original Transformer) causes gradient norms to blow up in early training and requires very careful warmup schedules to stabilise.
- The production trap: reducing d_ff below 4×d_model to cut compute. The FFN stores factual knowledge in its weight matrices. Reducing d_ff from 4× to 2× on a large language model is reported in interpretability and scaling studies to cost downstream knowledge-intensive task performance — not just a small accuracy delta but a loss of specific factual associations the model can no longer store. (No worked number for this is shown in this module; treat the direction as reliable, not the exact size.) Profile d_ff reduction on your specific tasks before accepting this tradeoff. The attention heads are frequently a safer target for compute reduction (fewer heads, or smaller d_k per head) without losing as much stored knowledge.
- The diagnostic: check training loss curve shape in the first 1000 steps. A healthy Transformer training run shows a rapid initial drop followed by smooth decay. A loss spike in the first 500 steps (then recovery) is the signature of insufficient learning rate warmup — Adam's moment estimates were too noisy for the initial learning rate. A flat loss that does not decrease at all is the signature of a missing or inverted causal mask in a decoder model — the model is attending to future tokens and the prediction task is trivially solved (training loss looks low, generation is random). Both are diagnosable before committing GPU hours.
The Transformer's power rests on three mutually dependent components: direct all-to-all attention for O(1) path length, residual connections that route gradients to every layer simultaneously, and a 4×-expanded FFN that stores and retrieves factual knowledge — each degrades measurably without the others.
Recap
- Attention's blind spot — no sense of order: raw self-attention is permutation-equivariant, so "dog bit man" and "man bit dog" produce the same (reordered) outputs. It sees a *set* of words, not a sequence.
- Positional encoding fixes it: add a position-dependent vector to each token's embedding. Sinusoidal encodings (sin/cos at many frequencies) have the property that PE(pos+k) is a linear function of PE(pos), which the original paper hedged *may* help extrapolation to unseen lengths — in practice that extrapolation is weak; modern models use RoPE, which encodes *relative* position directly in the attention scores. RoPE is friendlier to extrapolation than absolute encodings, but base RoPE still degrades past the trained context length without added scaling tricks (position interpolation, NTK-aware/YaRN scaling).
- The block, in order: norm → multi-head attention → residual add → norm → feed-forward network (FFN) → residual add. Pre-norm (norm inside the residual branch) is the stable modern default.
- The FFN is deliberately ~4× wide: it's the model's *memory* where factual knowledge is stored and retrieved — shrink it and the model measurably forgets facts, even though it's often overlooked next to attention.
- Residual shortcuts route gradients to every layer at once — ∂L/∂x keeps a direct identity path back regardless of the block's transform. That's necessary but not sufficient to stack dozens of Transformer blocks and train them: LayerNorm placement still determines whether that path stays well-scaled (see Pre-LN vs Post-LN below) — residuals alone don't stop Post-LN's early-training blow-ups.
- Two flavours: encoder-only (BERT — bidirectional, both-sides context, best for *understanding* tasks) and decoder-only (GPT — masks future tokens so it can only look left, best for *generation*, and since every token is a training target it gets denser signal per pass).
- Pre-LN over Post-LN: every modern large model uses Pre-LN because Post-LN blows up gradient norms early in training and needs careful warmup to survive — a stability property, not a preference.
Check your understanding
Q1. Why does the transformer use positional encodings, and why does standard sinusoidal encoding have a theoretical property that could support extrapolation to longer sequences than seen in training — even though, per the original paper and in practice, that extrapolation benefit is weak?
- A) Self-attention operates on the full sequence at once, so it needs explicit position markers. Sinusoidal encoding doesn't beat learned embeddings at generalising — both fail past the training max — but it saves the parameter cost of storing a learned table.
- B) Positional encodings exist because fixed-size attention matrices require every position to share the same embedding dimension. Sinusoidal encoding generalises because its high-frequency components repeat at short, regular wavelengths, giving familiar sub-patterns even at unseen positions.
- C) Self-attention is permutation-equivariant, so it needs an added position-dependent vector per token. Sinusoidal encoding has a theoretical edge because PE(pos+k) is a linear function of PE(pos), unlike learned embeddings with no representation for unseen positions — though the original paper only hedged this *may* help, and in practice the extrapolation is weak.
- D) Positional encodings compensate for the lack of recurrence, since RNNs encode order through sequential hidden-state updates. Sinusoidal encoding generalises because its frequency bands are tuned to match the natural frequency spectrum of language, independent of sequence length.
Q2. The feed-forward sublayer in a transformer block has two linear layers with a nonlinearity in between: FFN(x) = W₂·ReLU(W₁x + b₁) + b₂. Select the TWO correct statements about its role and 4× width.
- A) Attention is fundamentally a linear combination of value vectors, so the FFN supplies the essential nonlinearity; it's applied position-wise (same W₁, W₂ at every position), and 4× expansion (2048 for d_model=512) is empirically robust from BERT to GPT-3.
- B) The 4× expansion functions like a key-value memory: the up-projection selects which sparse "memories" activate in the high-dimensional space, and the down-projection reads their values back out — the source of the FFN's role in storing factual associations.
- C) The FFN provides cross-position communication that attention itself cannot achieve, and 4× width exists specifically to avoid an information bottleneck in mixing tokens across positions — below 4× the model must rely on attention alone for cross-position mixing.
- D) The FFN's real job is normalising attention output magnitude, which would otherwise grow unboundedly with sequence length from summing more values; the 4× width gives the compression back to d_model room to be nonlinear.
Q3. Why does training a transformer require a learning rate warmup, and what happens without it?
- A) Transformer weights start near zero under Xavier/He init; a full learning rate from step 1 overshoots that careful initialisation before the model has seen enough data to correct it, so warmup takes small steps to preserve trainability at depth.
- B) Residual connections create a gradient imbalance at init — layers near the loss get full gradient, early layers get almost none — so full LR makes late layers diverge while early layers barely move; warmup lets early layers accumulate gradient history first.
- C) At step 1, positional encodings are the largest signal in a first-layer input of near-zero random weights, so a full LR immediately overwrites the embedding layer's use of position; warmup lets the model learn to use positional encodings before large updates land.
- D) Adam's second-moment estimates are noisy for the first few batches, so a full LR takes confident steps in unreliable directions. Warmup ramps α from near-zero over 2000-4000 steps so moments settle — without it, transformers often diverge early.
Try it interactively
ML Systems Lab is a free interview-prep platform for ML engineers — work through the full interactive module, quizzes, and drills.
Open ML Systems Lab →