The Transformer
TL;DR — A Transformer is a stack of identical blocks, each of which is just two things: a mixing step (attention, which moves information between positions) and a computing step (a position-wise MLP, which processes each position independently). Around every one of those sits a residual connection, so the block’s input is added to its output and the stream flowing up the stack is never overwritten — only read from and added to. That residual stream is the mental model that makes the rest of the architecture obvious: norm layers keep its scale sane, the pre-norm placement keeps its gradient path clean, and depth is what lets the model compose operations. Attention is the only part that is ; the MLP holds most of the parameters and most of the FLOPs at realistic context windows.
1. Why this exists
The 2017 paper’s own framing is the right one: this was not a better layer, it was a better graph.
A recurrent network computes . Training it on a 1000-token sequence means 1000 dependent steps; the GPU’s parallel width is useless along the time axis, and the gradient from token 1000 to token 1 must survive 999 multiplies. A convolutional stack removes the sequential dependency but pays in path length: to relate two positions far apart you need layers of contiguous kernels, or with dilations. Both designs force you to choose between parallelism and reach.
The Transformer’s claim in Table 1 of the paper is narrower and more useful than “it’s better”: self-attention is per layer with sequential operations and maximum path length between any two positions. You pay a quadratic bill in exchange for removing the architectural barrier to parallelism and long-range mixing. Everything else — the hardware that got built, the scaling laws that fit, the fact that GPT-4-class models are decoder-only Transformer stacks — follows from that trade.
Here is the concrete failure the original design fixes. Take a task where the answer depends on two tokens 800 apart, e.g. “the error code in section 3 is which one, given the fix described in section 40?” An LSTM must route that dependency through 800 cell-state updates; in practice the signal is gone and the model attends to recency instead. Symptom in the product: the model answers using the last thing it read rather than the right thing, and retrieval quality is blamed when the architecture is the culprit.
What this chapter is and is not. Attention derives the mixing operation itself. This chapter is the block that wraps it, the stream that carries it, and the reason a 32-deep stack trains at all. If you can draw the block from memory and justify every arrow, this chapter did its job.
2. How it works
2.1 The one idea: the residual stream
Write for the vector at position after layer . Every modern block has the shape
and the whole point is the leading .
- — the residual stream. It is a -dimensional vector per position that is created once by the embedding and then never replaced. Every sublayer reads from it and writes an increment back into it. Without this identity path, each block would be a function that destroys its input: the only way information from layer 1 could reach layer 32 would be through 31 nonlinear transforms, and the gradient would have to survive the same journey backwards.
- — the increment: either the attention output or the MLP output. It is not the state; it is a correction to the state. Because the residual is additive, sublayers compose by summation, which is what makes the “write to a shared bus” interpretation (Elhage et al., A Mathematical Framework for Transformer Circuits, 2021, https://transformer-circuits.pub/2021/framework/index.html) legitimate rather than merely poetic.
- The absence of a scale on the residual — no with — is deliberate. Any per-layer decay would bound the effective path length between layer 1 and layer , reintroducing exactly the vanishing-gradient problem the design removed.
If you delete the residual connection, a 12-layer Transformer still trains but noticeably worse and becomes far more sensitive to learning rate; at 32 layers it typically does not converge. It is the cheapest ablation in the field.
2.2 The block, end to end
flowchart TB
X["x (n × d)"] --> R1["+"]
X --> N1["Norm"]
N1 --> ATT["Self-attention<br/>(mixes across positions)"]
ATT --> R1
R1 --> R2["+"]
R1 --> N2["Norm"]
N2 --> MLP["Feed-forward<br/>(each position alone)"]
MLP --> R2
R2 --> OUT["x (n × d) → next block"]
style R1 fill:#2E7A5A,stroke:#1E6A4A,color:#fff
style R2 fill:#2E7A5A,stroke:#1E6A4A,color:#fff
Formally, with the sequence length and the model width:
- — the current block’s input: one row per position, all rows carried on the same residual stream.
- — the only operation in the block that lets row see row . Its internals are the subject of Attention; here it matters only that it is position-aware (via positional encoding) and sequence-mixing.
- — applied to each row independently and with the same weights, so it is a per-position transform. It cannot move information between positions; it can only reshape what is already there.
- Two norms, not one, and one per sublayer. The norm belongs to the sublayer it feeds, not to the block.
The sequencing is the load-bearing part: the block alternates mix → think. Attention lets positions average each other; the MLP then applies a learned nonlinearity to the result. Neither is sufficient alone — attention with no MLP can only re-weight existing embeddings, and an MLP with no attention is a per-token function, no better than a bag of words for anything order-dependent.
2.3 The feed-forward network is not “just an MLP”
- — the up-projection. It lifts the residual stream into a wider space ( against in the 2017 base model; is the classic ratio). The expansion is what creates room for the activation to act per-feature; without it, the activation is applied in the same basis as the residual reads and the layer has far less representational room.
- — the elementwise nonlinearity. The 2017 paper uses (ReLU); modern models largely use gated variants (SwiGLU, Shazeer 2020, https://arxiv.org/abs/2002.05202). Without some nonlinearity the two linears collapse into one and the block becomes linear — attention included, since the whole thing is then a matrix product.
- — the down-projection back to the residual width, so the increment can be added to the stream.
- — biases. The 2017 paper has them; most modern LLMs drop them from the MLP and attention projections (LLaMA, PaLM, Mistral, Qwen). The reason is partly that a norm sits on the sublayer’s input, so a constant shift can be partly absorbed by of the preceding norm, and partly that removing them saves parameters and a kernel launch for no measured loss.
Two numbers make the MLP matter. At it holds parameters while the attention projections hold — so of the block’s total, or two-thirds of every block’s weights, live in the MLP. That is why “the model’s knowledge is in the MLP” is a reasonable first guess when you are deciding what to freeze, quantize, or ablate.
2.4 Pre-norm vs post-norm: where the norm goes
There are two placements, and the difference is not cosmetic.
Post-norm (the 2017 paper).
Here the norm is on the residual path: the identity you just added gets renormalised before it continues. Xiong et al. 2020 (On Layer Normalization in the Transformer Architecture, ICML 2020, https://arxiv.org/abs/2002.04745) prove with mean-field analysis that at initialization the expected gradients of parameters near the output layer of a post-norm stack are large, and that this is precisely why a learning-rate warmup is mandatory. The practical consequences:
- You need a warmup schedule or the run diverges.
- The required warmup grows with depth, so scaling the stack means retuning optimization, not just the model.
Pre-norm (most modern LLMs).
Now the norm is inside the increment, and the residual path is a pure identity. Xiong et al. show the gradients are well-behaved at initialization and the warmup can be dropped. That is why the pre-2020 encoder-decoder recipes treated warmup as mandatory and why modern recipes treat it as an optimization detail: warmup is still almost always used — it rarely hurts — but a pre-norm run without it no longer diverges.
The cost of pre-norm: the residual stream’s norm grows monotonically with depth (each block adds an increment), so the final norm before the output head is doing real work — it is not decoration. A side effect worth knowing: because the stream grows, per-layer increments become relatively smaller as you go deeper, which is one motivation for scaled residual-init schemes.
2.5 Layer norm: what each symbol does
- — one position’s vector, on the residual stream. The norm is applied per position, normalizing across the features; it never mixes across the sequence dimension.
- — the mean across features. Subtracting it re-centers the vector at the origin. Removing this term gives RMSNorm (below).
- — the variance across features. Dividing by rescales the vector to unit scale. This is the term that actually keeps training stable: it makes the sublayer’s input statistics independent of how large the residual stream has grown.
- — a small constant (commonly to ) that keeps the denominator away from zero when the vector is (near-)constant. Without it, a saturated or collapsed activation produces a division by ~0 and NaNs.
- — the learned scale. It restores the model’s ability to represent features of different magnitudes after normalization; without it, every sublayer would be forced to read unit-variance inputs.
- — the learned shift. It restores a per-feature offset. This is the parameter RMSNorm drops, on the empirical finding (Zhang & Sennrich 2019, https://arxiv.org/abs/1910.07467) that the re-centering is not what is doing the work — the re-scaling is. RMSNorm is therefore cheaper and has fewer parameters, and it is what LLaMA, PaLM, Gemma, Mistral, Qwen and DeepSeek use.
The interviewer’s version of this question. “Why layer norm and not batch norm?” Strong answer: batch norm computes statistics across the batch, so at inference you must freeze running estimates, and the statistics depend on sequence padding and on batch composition — fatal for variable-length text and for the autoregressive decode loop, where batch size is 1. Layer norm normalizes each position using only that position’s own features, so train and inference behave identically and batching does not change the result.
2.6 Encoder, decoder, and cross-attention
The original architecture is a stack of identical encoder layers and a stack of identical decoder layers, plus one extra sublayer in the decoder. The three attention layouts are the whole difference:
| Layout | from | from | Mask | Used by |
|---|---|---|---|---|
| Encoder self-attention | input | input | none (bidirectional) | BERT, T5 encoder |
| Decoder self-attention | output so far | output so far | causal | GPT, T5 decoder |
| Cross-attention | decoder | encoder | none on the encoder axis | T5/BART decoder, Whisper |
- Cross-attention’s has rows while have rows. The score matrix is therefore — rectangular, which is fine, and is the reason a seq2seq model can read a source of a different length than its target.
- There is no causal mask on the encoder axis because the whole source is available at once; masking it would hide input the decoder is entitled to.
- Encoder self-attention is bidirectional: masking it costs you exactly the right-context signal that makes BERT embeddings good for retrieval.
Decoder-only models (GPT and everything since) keep only the causal self-attention column. BERT vs GPT vs Encoder-Decoder does the objective-by-objective comparison; Modern Architecture Deltas covers what changed in the block since 2017 (pre-norm, RMSNorm, SwiGLU, GQA, FlashAttention, MoE).
2.7 Two 2017 details: the scale and weight tying
The paper multiplies the token embedding by before adding the positional encoding. With a norm immediately downstream, the pre-norm value does not affect forward outputs at all — the norm divides it out — and it matters only for initialization calibration: the embedding rows start with small variance, and puts them in the regime the norm’s (initialized to 1) expects. Delete it and the model still trains; the symptom is a slightly worse early loss curve, not a break.
Weight tying is the second: the input embedding, the output embedding, and the pre-softmax projection are the same matrix. That is parameters saved (16.4M of the base model’s ~60M, 27%), and it encodes the prior that a token’s input representation and its output logit direction should be related.
3. Implementation
Runnable block and stack, torch==2.4.0 (also correct on 2.1+), CPU-only:
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class TransformerBlock(nn.Module):
"""Pre-norm block: x = x + attn(norm(x)); x = x + mlp(norm(x))."""
def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.0):
super().__init__()
assert d_model % n_heads == 0
self.n_heads, self.d_k = n_heads, d_model // n_heads
# QKV in one matmul: 3x fewer kernel launches than three Linears.
self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
self.norm1 = nn.LayerNorm(d_model) # owns the attention sublayer
self.norm2 = nn.LayerNorm(d_model) # owns the MLP sublayer
self.mlp = nn.Sequential(
nn.Linear(d_model, d_ff, bias=False),
nn.GELU(),
nn.Linear(d_ff, d_model, bias=False),
)
self.dropout = nn.Dropout(dropout)
def forward(self, x, causal: bool = True):
B, T, _ = x.shape
# --- sublayer 1: attention, operating on the NORMED copy ---
h = self.norm1(x)
q, k, v = self.qkv(h).view(B, T, 3, self.n_heads, self.d_k) \
.permute(2, 0, 3, 1, 4) # each (B, H, T, d_k)
attn = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
attn = attn.transpose(1, 2).reshape(B, T, -1)
x = x + self.dropout(self.out_proj(attn)) # <-- the residual
# --- sublayer 2: MLP, position-wise, again on the NORMED copy ---
x = x + self.dropout(self.mlp(self.norm2(x))) # <-- the residual
return x
class TinyTransformer(nn.Module):
"""Decoder-only stack, pre-norm, final norm before the head."""
def __init__(self, vocab: int, d_model: int, n_heads: int, d_ff: int,
n_layers: int, max_len: int = 512):
super().__init__()
self.tok_emb = nn.Embedding(vocab, d_model)
self.pos_emb = nn.Embedding(max_len, d_model) # see positional-encoding.md
self.norm_f = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, vocab, bias=False)
# Tie AFTER the head exists: same tensor, and it must be set before the
# optimizer is constructed.
self.head.weight = self.tok_emb.weight
self.blocks = nn.ModuleList(
TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)
)
self.d_model = d_model
def forward(self, idx):
B, T = idx.shape
pos = torch.arange(T, device=idx.device)
x = self.tok_emb(idx) * math.sqrt(self.d_model) + self.pos_emb(pos)
for block in self.blocks:
x = block(x)
return self.head(self.norm_f(x)) # (B, T, vocab) logits
if __name__ == "__main__":
torch.manual_seed(0)
m = TinyTransformer(vocab=1000, d_model=512, n_heads=8, d_ff=2048, n_layers=6)
print(m(torch.randint(0, 1000, (2, 64))).shape) # torch.Size([2, 64, 1000])
print(sum(p.numel() for p in m.parameters())) # 19_661_824
That last number is checkable by hand: tied tok_emb ()
512,000, plus pos_emb () 262,144, plus 6 blocks at
each = 18,886,656, plus norm_f 1,024. Sum:
19,661,824. If your count is off by ~512,000 the head is not actually tied;
if it is off by a factor of two, you changed d_ff.
The four knobs that matter, with the ranges you actually see:
| Knob | Range in practice | What moving it does |
|---|---|---|
(d_model) | 512 (2017 base) → 8192 (Llama-3-70B) | The master dial. Compute and parameter count scale ~; the residual stream’s width, and therefore the model’s working memory, scales . |
(n_layers) | 6 (2017 base/encoder) → 80–126 in frontier models | Buys composition depth: each block applies one more mix-then-think step. Linear in parameters and compute. |
| classically; Llama-3-8B uses 14336 for () | The MLP’s inner width. Together with this dominates parameter count. | |
(n_heads) | 8 (2017) → 32–128 | Splits into subspaces. is kept at 64–128; heads are free in parameter count but not in . |
Notes on the code that interviewers probe:
out_projis applied inside the residual add, not outside it. Adding before projecting changes the block’s function, not just its efficiency.self.head.weight = self.tok_emb.weightrebinds both to one tensor and must run before the optimizer is built. Two parameter groups with different LRs will then apply conflicting updates to the same weights.norm_fbefore the head is not optional in a pre-norm stack: the residual stream’s norm grows with depth, so without it the logit scale — and hence the effective softmax temperature — varies with depth.is_causal=Truemust match the training-time masking exactly. A model trained bidirectionally and served causally produces garbage from token one.
4. Numbers and trade-offs
Parameter count, derived
Per block, with and no biases:
| Component | Parameters | For |
|---|---|---|
| Attention (, each ) | 1,048,576 | |
| MLP (, ) | 2,097,152 | |
| Two layer norms () | 2,048 | |
| Per block | 3,147,776 |
Six encoder layers: . Word embedding, tied
to the output projection, at a 32k-token vocabulary:
. The encoder alone is therefore
M. The decoder adds , giving a full encoder–decoder base model of roughly
M counting weights only. The familiar headline is ~65M, and the
bulk of that gap is vocabulary, not biases: the paper’s BPE vocabulary is ~37k
tokens rather than the 32k assumed here, and the extra M
of embedding takes the derived total to ~63M. The biases this derivation omits
and the fixed sinusoidal encoding (if you materialise it as a buffer) add only
~0.33M. Two observations worth carrying: the MLP is
two-thirds of every block, and the embedding is nearly half of a small
encoder. That share falls as grows, but not to
zero: Llama-3-8B uses a 128,256-token vocabulary and untied input and output
embeddings (tie_word_embeddings=False, as in Llama-2), so its single embedding
matrix is M parameters — about 6.5% of an 8B model —
and the input+output pair is B, ~13% of 8.03B. Those two are
by far its largest single matrices.
Compute per token, derived
Count multiply-accumulates as 2 FLOPs and ignore the softmax (it is with a tiny constant). Per layer, per token:
- Attention: QKV projections , output projection , plus the score and value matmuls at . Total .
- MLP: .
So the per-token budget is , and the ratio of attention to MLP work is . The crossover is . For the 2017 base model () that is 2048 tokens — at or above the length of most WMT sentence pairs — so 2017 was an MLP-dominated computation. For Llama-3-8B () the crossover is 16384 tokens: at 8k context, attention is
of the FLOPs per token, so the MLP still dominates. This is the honest answer to “why is long context expensive?” — not because attention is inherently the compute cost, but because it is the only term that grows with , and it grows quadratically.
Memory at inference: the KV cache, in context
The cache size is the formula derived in Attention §4: , where is the per-head width and is the number of KV heads (equal to the head count without GQA). Plug in the 2017 base model — 6 encoder plus 6 decoder layers, 8 heads, , fp32 — and it is bytes/token KiB/token (binary units: 1 KiB = 1024 B), about 192 MiB (201 MB) at 4096 tokens: trivial next to a modern model’s weights. The regimes where that stops being true, and the serving consequences, are in Attention §4 and vLLM and PagedAttention.
Which axis you are optimizing
| Decision | Buys you | Costs you |
|---|---|---|
| Deeper stack () | More composition; better sample efficiency at fixed params | Sequential layers, startup/warmup issues, no parallelism win |
| Wider stack () | Bigger residual stream, more per-token memory | Compute ~; needs more data to train |
| Post-norm | Slightly better final quality if you tune warmup | A warmup hyperparameter that grows with depth; divergence without it |
| Pre-norm | Trains without warmup; gradient path is clean | Unbounded residual growth; a mandatory final norm |
| SwiGLU MLP | Better loss at equal parameter count | Three matrices instead of two; more memory traffic |
5. Pitfalls
Post-norm with a pre-norm learning rate. Symptom: the loss spikes in the first few hundred steps and either recovers only with a warmup restart or diverges outright. The gradient at initialization near the output layer is simply too large for the LR you chose — Xiong et al. 2020 derives exactly this.
Forgetting the final norm in a pre-norm stack. Symptom: training is fine but the logits’ scale drifts with depth; you see the loss plateau at a value that looks like a calibration problem rather than a capacity problem, and sampling temperature has to be retuned per checkpoint.
Applying the residual add to the wrong tensor. x = x + mlp(norm2(x)) is
correct; x = norm2(x + mlp(x)) is post-norm; x = x + mlp(x) (no norm) usually
still trains at small scale and hides the bug. Symptom: no crash, a loss curve
that is merely worse, and a model that is unusually sensitive to initialization
scale.
Assuming the MLP is cheap. Symptom: profile shows the “attention” budget blown at 32k context and you optimize attention, while 80%+ of the FLOPs are in / GEMMs. At the MLP is the thing to optimize — quantization and MoE (see Modern Architecture Deltas) attack it directly.
Tying weights with two optimizer groups. Symptom: the loss increases for a few hundred steps and then recovers, or embedding norms diverge from head norms. The two parameter groups are applying conflicting updates to one tensor.
Assuming a causal model can be evaluated bidirectionally. Symptom: an elegant offline evaluation using masked-token prediction scores far above the model’s real generation quality, because the model has never seen a right context during training.
Treating as a rounding detail. Symptom: you port a checkpoint, change to a “nicer” number, and get NaN-free garbage — the shapes match and the weights do not.
6. Using it in a project
You almost never build a Transformer. What you actually control:
The checkpoint’s shape, not the architecture. Vocabulary size, , , , and head count are fixed when you download a model. Your only levers are which checkpoint, how many bits, and how it is served.
When NOT to use a Transformer block. Reach for something else when:
- is small and the structure is tabular. Gradient-boosted trees beat a Transformer on a 10k-row, 50-column tabular problem, they need no tokenization, and they train in seconds. The Transformer’s inductive bias buys nothing when there is no meaningful long-range order to model.
- Latency budget is single-digit milliseconds. A Transformer’s cost is dominated by sequential block evaluations; a distilled CNN or a n-gram/feature-hashing model will fit in a budget the Transformer cannot.
- The task is a per-token classification with abundant labels and short inputs. A frozen encoder plus a linear head is usually the right call; a decoder-only LLM is a much more expensive way to get the same label.
- You need a hard guarantee of determinism and auditability. Transformer generations are stochastic functions of weights and sampling; if the domain demands a rule engine, no amount of prompting fixes it.
- The data is a few hundred labeled examples. Fine-tuning the block structure will overfit; few-shot prompting plus retrieval, or a linear probe on a frozen encoder, is the better first move.
When a Transformer is the right tool but the block is not the problem. If your RAG system gives bad answers, the block is nearly never at fault — retrieval recall, prompt construction, and output parsing are. See RAG. Optimizing the block before measuring those is the classic premature optimization.
Two block-level decisions you do get to make. Pre-norm vs post-norm when you
train from scratch (choose pre-norm), and whether to use the fused attention
kernel (choose it; F.scaled_dot_product_attention is a drop-in with the same
math). Everything else is already decided by whoever trained the checkpoint.
7. How you would prove it worked
A Transformer has no product metric of its own. What you prove is that a component earns its place. Three experiments, in increasing cost:
The residual ablation. Train two identical small models, one with and one with only. (example) Expect the no-residual run to plateau at a substantially higher loss and to be far more LR-sensitive. This is the one experiment in this chapter that is nearly always worth running, because it makes the central claim of the architecture measurable in an afternoon.
Pre-norm vs post-norm with warmup removed. Train both, both without warmup. (example) Post-norm shows an early loss spike and often diverges; pre-norm descends smoothly. Then add warmup back to post-norm and show it recovers — which is the demonstration that the norm placement was the variable, not the optimizer.
The mixing/thinking split. Ablate attention only, then the MLP only, at matched parameter count. (example) Removing attention hurts order-sensitive and retrieval-sensitive tasks much more than it hurts token-level tasks; removing the MLP hurts nearly everything. If your two curves look identical, you are evaluating on a task that does not depend on sequence order, and the ablation proves nothing. The cheap version of this experiment uses the permutation test from Positional Encoding §7 as the order-sensitive probe.
Reporting template (example — replace with your own numbers). “We compared pre-norm and post-norm 12-layer blocks on the same data at the same LR. Post-norm without warmup diverged at step ~400; with a 2000-step warmup it reached 2.61 validation loss. Pre-norm without warmup reached 2.58 and was stable, removing one hyperparameter and 2000 steps of schedule.”
8. Interview drill
“Explain a Transformer block in 60 seconds.” Strong answer: two sublayers — attention, which mixes information across positions, and a position-wise MLP, which transforms each position independently — each wrapped in a residual connection with a layer norm on the sublayer’s input. The residual stream is the state; the sublayers only add to it. Pre-norm placement is what makes deep stacks trainable without warmup. Common wrong answer: reciting the 2017 figure as “attention then layer norm then feed-forward then layer norm” without saying what the residual path is for. That is a diagram, not an explanation.
“Why is the residual connection load-bearing?” Strong answer: it makes the identity path between the input embedding and the output head unbroken, so gradients reach early layers unimpeded and information is added rather than overwritten — the “residual stream” reading. Removing it reintroduces depth-dependent gradient decay. Follow-up: “why does the norm then have to move inside the residual?” → because a norm on the residual path renormalises the identity path, which is exactly the post-norm instability Xiong et al. 2020 analyses.
“Why did the Transformer replace the RNN? Be specific.” Strong answer: self-attention has sequential operations per layer and maximum path length between positions, versus and for a recurrence. That removes the parallelism barrier on the time axis and the long-range gradient-decay problem at once, at the cost of compute. Common wrong answer: “it has better accuracy.” Accuracy is an outcome; the architectural claim is about parallelism and path length, and both are exactly what made training on web-scale data feasible.
“What does the feed-forward network do that attention can’t?” Strong answer: it is the only nonlinear, per-position computation in the block. Attention only re-weights and averages existing vectors — it is a convex combination, so with softmax weights the output is bounded by the inputs’ convex hull. The MLP creates new features (an expanded hidden space, a nonlinearity, a projection back), and it holds roughly two-thirds of the block’s parameters. Common wrong answer: “it’s just a classifier” or “it adds depth.” It is a key-value memory: the up-projection’s rows act as learned feature detectors and the down-projection turns each on into a residual-stream direction, which is why the MLP — not attention — is the target of MoE and the layer that quantization damages most.
“Which is more expensive, attention or the MLP?” Strong answer: it depends on relative to . They cross at for : below it the MLP dominates, above it attention does. For a 2017 base model at 512 tokens attention is ~14% of the FLOPs (); at 32k tokens it is ~91%. “Attention is the expensive part” is only true at long context. Common wrong answer: “attention, because it’s ” — dropping the constant is fine for asymptotics and wrong for the regime you are actually in.
“Why layer norm and not batch norm?” Strong answer: layer norm computes statistics per position across features, so its behaviour is identical at training and inference and independent of batch size and padding. Batch norm couples positions in a batch, which is degenerate in autoregressive decode where the effective batch is 1, and it makes sequence padding statistically meaningful. Follow-up: “why does RMSNorm work?” → the re-centering term matters less than the re-scaling term empirically (Zhang & Sennrich 2019); RMSNorm keeps the rescaling and drops the mean and bias, saving parameters and a reduction.
“The 2017 paper scaled embeddings by — why?” Strong answer: a scale calibration for the embedding’s initial statistics relative to the norm’s learned ; with a norm immediately after, the forward output is invariant to it, so removing it costs a slightly worse early loss curve and nothing structural. Common wrong answer: “to prevent the positional encoding from dominating,” or “to keep gradients from vanishing.” Both are plausible-sounding and neither is what the paper’s own code does.
“What would you change in the block for a 2026 model?” Strong answer: pre-norm with RMSNorm, SwiGLU MLP with (rounded for hardware), no biases, RoPE instead of additive positional encoding, GQA to shrink the KV cache, and a fused FlashAttention kernel. Each of those is a response to a specific cost — see Modern Architecture Deltas.
Next: Attention — the mixing half of the block, derived rather than asserted: why the query–key dot product, why , and why multiple heads.