Transformer Architecture Deep Dive

Why an architecture lesson earns its place

You can build with these models without knowing any of this. But the architecture is what explains the things that otherwise look arbitrary: why long context is expensive rather than merely slow, why the first token takes so much longer than the rest, why a fact buried mid-prompt gets missed, and why prompt caching is a prefix operation rather than a semantic one.

Every one of those is a direct consequence of what follows.

Self-attention

The transformer’s contribution was letting every token look at every other token directly, in parallel.

Each token is projected into three vectors:

• Query — what this token is looking for
• Key — what this token offers
• Value — what this token contributes if attended to

Attention is then:

$\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V$

Read it as a lookup. $QK^{\top}$ scores every query against every key. Softmax turns scores into weights summing to one. Multiplying by $V$ produces a weighted blend of what the attended tokens contribute.

Why divide by $\sqrt{d_k}$? For vectors of dimension $d_k$ with unit-variance components, the dot product has variance $d_k$. Large-magnitude scores drive softmax towards a one-hot distribution, where gradients vanish. Dividing by $\sqrt{d_k}$ restores unit variance and keeps the distribution soft enough to train.

The quadratic cost

$QK^{\top}$ is $n \times n$ for a sequence of length $n$. Both compute and memory for attention scale as $O(n^2 d)$.

This is the structural fact about long context. Ten times the input is a hundred times the attention work. Every long-context system — sliding windows, sparse attention, FlashAttention’s tiling, linear-attention variants — exists to attack this term. FlashAttention is worth understanding specifically: it does not change the mathematics, it avoids materialising the full $n \times n$ matrix in high-bandwidth memory, which is a memory-movement win rather than an arithmetic one.

Multi-head attention

Rather than one attention operation of dimension $d$, run $h$ of them at dimension $d/h$ and concatenate.

Different heads specialise, and this is empirically observable: some track syntactic dependencies, some resolve coreference, some attend to delimiters, some attend to the first token as a no-op sink. Same computation, different learned projections.

Modern variants trade quality for cache size. Multi-query attention shares one key/value head across all query heads; grouped-query attention shares across groups. The motivation is entirely practical — the KV cache is the memory bottleneck at inference, and this shrinks it several-fold for a small quality cost. Almost every deployed model now uses one of these.

The KV cache

This is the most operationally important idea in the lesson, and it is usually left out.

Generation is autoregressive: produce a token, append it, produce the next. Naively, each step re-runs attention over the whole sequence — $O(n^2)$ work per token, $O(n^3)$ for a response.

But the keys and values for previous tokens do not change. Causal masking means token $i$ never attends to anything after it, so its K and V projections are fixed the moment it is processed. Cache them.

With the cache, each new token computes its own Q, K, V and attends against the stored ones: $O(n)$ per step.

Three consequences follow directly:

Generation has two distinct phases with different bottlenecks.

This is why time-to-first-token scales with prompt length while inter-token latency is roughly flat. It is also why batching helps decode enormously — you are reading the same weights for many sequences at once — and why output tokens are priced several times higher than input.

Cache memory grows linearly with context. For a long context, the KV cache can exceed the size of the model weights. This is the real constraint on how many concurrent conversations a server holds, and the reason grouped-query attention was adopted so quickly.

Prompt caching is prefix-only, and now you know why. A provider caching your prompt is storing the KV cache for that prefix. Change one byte at position 500 and every subsequent key and value is different — the cached state is invalid from that point on. It is not that providers were lazy about semantic matching; the structure genuinely does not permit it.

Position

Attention is permutation-invariant. Without positional information, “dog bites man” and “man bites dog” are identical inputs.

The original paper added sinusoidal encodings to the embeddings:

$PE_{(pos,2i)} = \sin\!\left(\frac{pos}{10000^{2i/d}}\right), \qquad PE_{(pos,2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d}}\right)$

Modern models overwhelmingly use rotary position embeddings instead. RoPE rotates the query and key vectors by an angle proportional to position. Because the dot product of two rotated vectors depends on the difference of their angles, attention scores become a function of relative position — which is what actually matters linguistically, and which extrapolates to longer sequences far better than an added absolute signal.

Extending context length in practice usually means interpolating or rescaling RoPE frequencies, then fine-tuning. This is why “we extended the context window” is a real engineering effort rather than a configuration change — and why quality at the top of an extended window is often worse than quality in the range the model was actually trained on.

The block

A modern decoder block, in the arrangement almost everyone converged on:

x = x + Attention(RMSNorm(x))       # pre-norm residual
x = x + FeedForward(RMSNorm(x))

Four components, each earning its place:

• Pre-norm. Normalising before the sublayer rather than after gives a clean residual path and makes deep stacks trainable without warmup tricks.
• RMSNorm over LayerNorm — drops the mean-centring term, marginally cheaper, no measured quality cost.
• Residual connections. The gradient highway. Without them, depth beyond a dozen layers does not train.
• Feed-forward network. Typically 4× the model dimension, now usually with a SwiGLU activation. This is where most parameters live, and the evidence suggests it is where most factual knowledge is stored.

Stack 32 to 128 of these.

Mixture of experts

Most frontier models now replace the dense feed-forward layer with a sparse mixture of experts: many parallel FFNs, and a router that sends each token to only a few.

The economics are the point. A model can hold a very large total parameter count while activating only a fraction per token, so capacity is decoupled from inference cost. Total parameters determine memory footprint; active parameters determine compute.

The awkward parts are real: load balancing across experts needs an auxiliary loss or some experts starve; the router is a discrete decision that complicates training; and memory requirements stay high because every expert must be resident even though most are idle.

Decoder-only, and why

Decoder-only with causal masking won because it makes the training objective uniform — predict the next token — over any text at all, which lets you train on the entire internet without paired data. Every task then becomes a text-continuation task. The generality was worth more than the architectural specialisation.

Encoder models have not disappeared; they remain the right tool for embeddings, which is why your retrieval pipeline probably uses one.

Sampling

The model emits logits over the vocabulary. Turning those into a token is a separate, controllable step:

• Temperature $T$ divides the logits before softmax. $T \to 0$ approaches argmax; $T > 1$ flattens the distribution.
• Top-k restricts to the $k$ highest-probability tokens.
• Top-p (nucleus) restricts to the smallest set whose cumulative probability exceeds $p$ — adapts to how peaked the distribution is, which is why it generally beats top-k.

Worth stating precisely: temperature 0 is not “accurate mode”. It is deterministic mode. A model can be reproducibly wrong, and often is.

What this buys you

• Long context is quadratic in attention and linear in cache memory — a cost, not just a limit.
• Prefill and decode are different workloads, which is why streaming, batching and prompt caching are the levers they are.
• Prompt caching must be prefix-exact, because it is a KV cache.
• Position is relative and extrapolated, which is why quality degrades near the top of an extended window.
• Mixture-of-experts separates capacity from per-token cost.

Try this: Send a 200-token prompt and a 20,000-token prompt to the same model and measure time-to-first-token for each, then inter-token latency for each. The first will scale roughly with prompt length; the second will barely move. You have just measured prefill against decode.