Transformer Architecture and Its Variants

Chapter 7 ended with a stable mapping from text to token ids. This chapter follows those ids through a decoder-only transformer: embedding, a stack of residual blocks, a final normalization, and a vocabulary projection. The 2017 Transformer introduced the attention-and-feed-forward block in an encoder-decoder system (Vaswani et al. 2017). Modern causal language models reuse that basic pattern, but differ in normalization, position handling, feed-forward design, attention heads, and cache representation.

Those choices are not cosmetic. They determine which positions can exchange information, how gradients travel through depth, how many parameters sit in each block, and how much state a serving process retains for every active sequence. A useful architecture description therefore needs tensor shapes, mask semantics, and memory units, not only a box labeled “Transformer.”

Scope: a causal decoder

Here a batch contains $B$ sequences of length $S$, and the model width is $d$. The tokenizer emits integer ids $T\in\{0,\ldots,V-1\}^{B\times S}$ for a vocabulary of size $V$. An embedding table maps them to the initial residual stream $X_0\in\mathbb{R}^{B\times S\times d}$. After $L$ blocks and a final normalization, an output matrix produces logits over the vocabulary:

$$\widehat{X}_L=\operatorname{Norm}_{f}(X_L), \qquad Z=\widehat{X}_LW_U, \qquad p(t_{i+1}=v\mid t_{\le i})=\operatorname{softmax}(Z_i)_v.$$

Here $X_L$ is the last block's output, $\operatorname{Norm}_{f}$ is the final normalization, $W_U\in\mathbb{R}^{d\times V}$ is the output or unembedding matrix, $Z_i\in\mathbb{R}^{V}$ contains the logits at position $i$, and $v$ is a candidate next-token id. The causal condition $t_{\le i}$ means that position $i$ may depend only on tokens at positions up to $i$. Encoder-only models use bidirectional masks; encoder-decoder models also add cross-attention. Those architectures share components with this one, but they do not share its exact information-flow contract.

The input embedding and $W_U^\top$ may share parameters, a design called weight tying (Press and Wolf 2017). Tying removes one $Vd$ matrix, but also forces the input and output roles to use the same rows. It is an architecture choice, not a rule for large or small models.

A block updates one residual stream twice

A common pre-normalized decoder block has two residual updates:

$$U_\ell=X_\ell+ \operatorname{Attn}_\ell(\operatorname{Norm}_{\ell,1}(X_\ell)), \qquad X_{\ell+1}=U_\ell+ \operatorname{FFN}_\ell(\operatorname{Norm}_{\ell,2}(U_\ell)).$$

Here $\ell\in\{0,\ldots,L-1\}$ is the block index; $X_\ell$ is the block input; $U_\ell$ is the state after self-attention; and $X_{\ell+1}$ is the block output. Attention mixes information across allowed positions. The feed-forward network (FFN) transforms each position independently. Both return a tensor with the same shape as the residual stream so that addition is defined.

The residual stream is the width-$d$ state carried through the stack. Calling it a stream is helpful, but it is not a separate memory module. It is the sequence of hidden vectors that every sublayer reads and updates.

This diagram specifies pre-normalization. The original Transformer instead normalized after each residual addition (Vaswani et al. 2017):

$$U_\ell=\operatorname{Norm}_{\ell,1} \bigl(X_\ell+\operatorname{Attn}_\ell(X_\ell)\bigr), \qquad X_{\ell+1}=\operatorname{Norm}_{\ell,2} \bigl(U_\ell+\operatorname{FFN}_\ell(U_\ell)\bigr).$$

Here the symbols have the same shapes as in the pre-normalized block, but the normalizers now act after each addition. Xiong et al. showed why pre-normalization has better-behaved gradients at initialization and can reduce reliance on learning-rate warmup (Xiong et al. 2020). That evidence does not make post-normalization invalid: initialization, residual scaling, depth, and optimizer settings change the comparison. Record the exact block equation rather than relying on the labels “pre-norm” or “post-norm.”

Normalization controls scale

LayerNorm centers a hidden vector and divides by its standard deviation (Ba et al. 2016). root mean square layer normalization (RMSNorm), introduced by Zhang and Sennrich in 2019, omits mean subtraction and normalizes by root mean square (Zhang and Sennrich 2019):

$$\operatorname{RMSNorm}(x) =\gamma\odot \frac{x}{\sqrt{\frac{1}{d}\sum_{j=1}^{d}x_j^2+\varepsilon}}.$$

Here $x\in\mathbb{R}^{d}$ is one token's hidden vector; $x_j$ is its $j$-th component; $\gamma\in\mathbb{R}^{d}$ is a learned element-wise scale; $\varepsilon>0$ prevents division by zero; and $\odot$ denotes element-wise multiplication. Implementations must pin $\varepsilon$, accumulation precision, and whether a learned bias is present. These details can affect checkpoint compatibility and low-precision stability.

Some architectures also normalize queries and keys before their dot product. That QK normalization solves a different problem from normalizing the residual stream: it bounds attention-score scale. A model manifest should state both locations independently.

The feed-forward network supplies per-position capacity

The original block used two matrices with a ReLU between them. A common gated alternative is Swish-gated linear unit (SwiGLU) (Shazeer 2020):

$$\operatorname{SwiGLU}(x) =\bigl(\operatorname{SiLU}(xW_g)\odot(xW_u)\bigr)W_d.$$

Here $x\in\mathbb{R}^{d}$ is one residual vector; $W_g,W_u\in\mathbb{R}^{d\times d_f}$ are the gate and value projections; $W_d\in\mathbb{R}^{d_f\times d}$ projects back to model width; $d_f$ is the FFN hidden width; and $\operatorname{SiLU}(z)=z/(1+e^{-z})$ is applied element-wise. Ignoring biases, this gated FFN has $3dd_f$ parameters. A plain two-matrix FFN has $2dd_f$.

That count explains the often-quoted “two-thirds” adjustment. A plain FFN with hidden width $4d$ has $8d^2$ parameters. Matching it with a three-matrix gated FFN requires $d_f=8d/3$, before rounding to a hardware-friendly multiple. The ratio is a budgeting rule, not an intrinsic property of SwiGLU. Actual models choose different expansion ratios, and equal matrix counts do not guarantee equal wall-clock time across kernels or hardware.

Attention and the FFN play different roles, but claims that one stores “knowledge” while the other only “routes” it are interpretations, not tensor invariants. Ablations can localize behavior; the forward equations alone do not assign semantic ownership.

Causal self-attention with explicit shapes

Here, for one sequence, let $X\in\mathbb{R}^{S\times d}$. Assume equal query, key, and value head widths. Let $H_q$ be the number of query heads, $H_{kv}$ the number of distinct key-value heads, and $d_h$ the width of one head. Learned projections produce

$$Q\in\mathbb{R}^{S\times H_q\times d_h}, \qquad K,V\in\mathbb{R}^{S\times H_{kv}\times d_h}.$$

Each query head $a$ is assigned to a key-value head $g(a)$. For query position $i$, its output is

$$\alpha_{ij}^{(a)} =\operatorname{softmax}_{j} \left( \frac{q_i^{(a)\top}k_j^{(g(a))}}{\sqrt{d_h}}+M_{ij} \right), \qquad z_i^{(a)}=\sum_{j=1}^{S}\alpha_{ij}^{(a)}v_j^{(g(a))}.$$

Here $q_i^{(a)}$ is the query for head $a$ at position $i$; $k_j^{(g(a))}$ and $v_j^{(g(a))}$ are the selected key and value at position $j$; $\alpha_{ij}^{(a)}$ is their normalized attention weight; and $M$ is the mask. For ordinary causal attention, $M_{ij}=0$ when $j\le i$ and $M_{ij}=-\infty$ when $j>i$. The $-\infty$ entries become zero probability after softmax. Padding, packed documents, sliding windows, and prefix-language models require different masks.

The $\sqrt{d_h}$ divisor keeps score variance from growing with head width under common initialization assumptions (Vaswani et al. 2017). Concatenating the $H_q$ head outputs and applying an output projection returns a width-$d$ update for the residual addition.

When $H_qd_h=d$ and biases are omitted, the query and output projections hold $2d^2$ parameters. The key and value projections hold $2dH_{kv}d_h$. Standard MHA with $H_{kv}=H_q$ therefore uses $4d^2$ projection parameters. GQA and MQA reduce the K/V term as well as the cache, while dense query-key interaction still scales with $S^2H_qd_h$.

A dense attention layer forms $S^2$ query-key relationships across a full sequence. The arithmetic is therefore quadratic in $S$. A naive implementation also materializes an $S\times S$ score matrix, but that temporary storage is not required by the mathematical definition.

Position enters through queries and keys

Unmasked content-only self-attention is permutation-equivariant: permuting the input positions permutes its outputs. A causal mask already imposes a direction, but it does not give equal content vectors an explicit distance representation. rotary position embedding (RoPE) applies a position-dependent rotation to pairs of query and key coordinates (Su et al. 2024). For an even rotated width $d_R$, define

$$\omega_r=\theta^{-2r/d_R}, \qquad R_p^{(r)}= \begin{bmatrix} \cos(p\omega_r)&-\sin(p\omega_r)\\ \sin(p\omega_r)&\cos(p\omega_r) \end{bmatrix}, \quad 0\le r<d_R/2.$$

Here $\theta>0$ is the frequency base; $r$ indexes a pair of rotated coordinates; $\omega_r$ is that pair's angular frequency; and $p$ is token position. Let $R_p$ be the block-diagonal matrix made from all $R_p^{(r)}$. Then

$$q_i'=R_iq_i, \qquad k_j'=R_jk_j, \qquad q_i'{}^\top k_j'=q_i^\top R_{j-i}k_j.$$

Here $q_i$ and $k_j$ are unrotated query and key vectors; $R_i$ and $R_j$ use the configured frequency schedule; and the last identity shows that the positional part of their dot product depends on relative offset $j-i$. RoPE does not make the complete attention score a function of distance alone; the content vectors still matter.

The frequency base, the fraction of dimensions rotated, and any scaling used for context extension belong in the checkpoint contract. Changing them after training is an adaptation experiment, not a metadata edit. RoPE is common, but it is not the only valid scheme. ALiBi instead adds a head-specific linear distance penalty to each causally valid score (Press et al. 2022):

$$s_{ij}^{(a)}= \frac{q_i^{(a)\top}k_j^{(a)}}{\sqrt{d_h}} -m_a(i-j)+M_{ij}.$$

Here $m_a>0$ is a fixed slope for head $a$; $i-j$ is non-negative for an allowed past position; and $M_{ij}$ is the causal mask defined earlier. ALiBi introduces no learned position vectors. Long-context behavior under either method depends on training length, frequency or slope configuration, later adaptation, and the evaluation distribution.

Prefill creates the cache; decode reuses it

Autoregressive serving exposes a distinction that a full-sequence training diagram hides. The first pass is prefill: it reads the entire prompt under the causal mask, produces logits at its positions, and populates keys and values for the prompt. The logits at the last prompt position predict the first generated token. The positions can run in parallel, although serving systems may split a long prefill into chunks. The second phase is decode: it processes one new position per sequence, attends from that new query to stored past keys and values, selects or samples a token, appends that position's keys and values, and repeats.

In standard causal self-attention with a fixed prefix, a past position's key and value at a given layer do not change during later decode steps. Reusing them avoids projecting the whole prefix again. A dense decode step still reads the retained cache, making it a bandwidth cost as well as a capacity cost. For a ragged batch, the persistent key-value (KV) tensor payload is

$$M_{\mathrm{KV}} =b\sum_{r=1}^{B}S_r \sum_{\ell=1}^{L}H_{kv,\ell}(d_{k,\ell}+d_{v,\ell}) \quad\text{bytes}.$$

Here $B$ is the number of resident sequences; $S_r$ is the cached length of sequence $r$; $L$ is the number of attention layers; $H_{kv,\ell}$ is the number of distinct KV heads in layer $\ell$; $d_{k,\ell}$ and $d_{v,\ell}$ are the key and value elements per head; and $b$ is bytes per cached element. For uniform length $S$ and equal $d_k=d_v=d_h$, this reduces to $2LBSH_{kv}d_hb$. Real servers also have allocation blocks, metadata, padding or fragmentation, quantization scales, speculative or beam state, prefix sharing, and possible offload. The formula is the dense tensor payload, not total process memory.

With cache reuse, dense attention for one new token is linear in retained length $S$ per layer, omitting head and dimension factors. Generating $T$ new tokens after a prompt of length $P$ therefore performs attention work proportional to $TP+T^2$ across decode. Cache reuse prevents repeated prefix projection; it does not make long decode constant-time.

The cache scales linearly with resident tokens. Whether it exceeds model weights depends on all the quantities above, plus weight and cache precision. It is therefore incorrect to present one crossover as a property of Transformers in general.

Figure 8.4 uses explicit assumptions: 32 layers, 128-wide heads, two bytes per KV element, batch size one, and 7 billion two-byte weight parameters. The plotted MHA, GQA, and MQA lines differ only in $H_{kv}$, so every value follows directly from the cache equation.

The runnable cell prints the cache payload per token and the context length at which that payload equals the stated weight bytes. Change batch size, precision, or KV-head count to test a different deployment.

layers = 32
head_width = 128
bytes_per_kv_element = 2
batch_size = 1
parameter_count = 7_000_000_000
bytes_per_weight = 2
weight_bytes = parameter_count * bytes_per_weight

def kv_payload_bytes(sequence_length, kv_heads):
    return (
        2
        * layers
        * batch_size
        * sequence_length
        * kv_heads
        * head_width
        * bytes_per_kv_element
    )

for name, kv_heads in [("MHA", 32), ("GQA", 8), ("MQA", 1)]:
    bytes_per_token = kv_payload_bytes(1, kv_heads)
    crossover_tokens = weight_bytes / bytes_per_token
    cache_at_128k_gib = kv_payload_bytes(128_000, kv_heads) / 2**30
    print(
        f"{name}: {bytes_per_token / 1024:.0f} KiB/token; "
        f"128k cache={cache_at_128k_gib:.2f} GiB; "
        f"equals weight bytes at {crossover_tokens:,.0f} tokens"
    )

MHA, GQA, and MQA change KV-head sharing

Multi-head attention (MHA) uses $H_{kv}=H_q$: every query head has its own key and value head (Vaswani et al. 2017). Multi-query attention (MQA) uses $H_{kv}=1$, so all query heads share one key and value head (Shazeer 2019). Grouped-query attention (GQA) chooses an intermediate $1<H_{kv}<H_q$ and assigns several query heads to each KV head (Ainslie et al. 2023).

Under the dense-payload formula, changing MHA to GQA reduces cache bytes by $H_q/H_{kv}$ when all other dimensions stay fixed. Changing MHA to MQA reduces them by $H_q$. These are exact storage ratios, not quality claims. Ainslie et al. reported that an MHA checkpoint could be uptrained into GQA using 5% of its original pretraining compute, with quality close to MHA and speed comparable to MQA in their experiments (Ainslie et al. 2023). The result supports GQA as a useful design point; it does not prove one KV-head count is optimal for every model and context length.

MLA changes the cached representation

Multi-head latent attention (MLA), introduced in DeepSeek-V2, does not obtain its main saving by sharing a conventional KV head (DeepSeek-AI 2024). It projects each residual vector into a low-dimensional latent $c_i^{KV}\in\mathbb{R}^{d_c}$, from which content keys and values are derived. A separate small RoPE key component is retained because applying a position-dependent rotation prevents the full key projection from being absorbed into later matrix multiplications.

In the DeepSeek-V2 construction, the per-position payload is

$$M_{\mathrm{MLA}}=LBS(d_c+d_h^R)b\quad\text{bytes}.$$

Here $d_c$ is the joint KV latent width; $d_h^R$ is the decoupled RoPE-key width; and $L$, $B$, $S$, and $b$ have the meanings defined for the dense cache. DeepSeek-V2 sets $d_c=4d_h$ and $d_h^R=d_h/2$, or $4.5d_h$ cached elements per position and layer. A conventional equal-width MQA cache stores $2d_h$ elements, so MLA is not generically “smaller than MQA.” During optimized decode, the up-projection matrices can be absorbed into the query and output projections; full per-head K and V need not be explicitly reconstructed.

DeepSeek reported lower cache requirements than its MHA baseline and better results on most listed metrics in its controlled ablations. That is evidence for the reported architecture and training recipe, not a general theorem that latent attention dominates GQA.

Separate arithmetic, temporary memory, and persistent state

Several techniques called “efficient attention” act on different resources:

FlashAttention tiles exact attention so the full score matrix need not be written to high-bandwidth memory (Dao et al. 2022). It reduces IO and temporary activation memory during training and prefill. It does not by itself change the number of past key/value elements retained for decode.

Sparse attention changes the mathematical connectivity. Native sparse attention jointly trains three gated branches: compressed coarse-grained tokens, selected fine-grained blocks, and a local sliding window (Yuan et al. 2025). DeepSeek Sparse Attention (DSA) is a distinct later design. Its low-dimensional lightning indexer ranks past tokens, after which the main MLA path attends to a selected top-$k$ subset (DeepSeek-AI 2025). The main attention work moves from all $S$ past positions toward $k\ll S$, but the indexer reported for DeepSeek-V3.2 still scores all query-key pairs and remains quadratic at a smaller dimension. Selector error, continued-training cost, and retained cache state also count. Sparse selection does not automatically delete the full cache.

Record the architecture as a contract

A reproducible decoder specification should name at least:

• vocabulary size, embedding width, output projection shape, and weight tying;
• block count, model width, FFN width, activation, and all bias settings;
• residual ordering, normalization locations, type, epsilon, and precision;
• query-head count, KV-head count, head width, and query-to-KV mapping;
• attention mask, document-boundary behavior, and any local or sparse pattern;
• position mechanism, RoPE base and rotated dimensions, plus scaling rules;
• cache representation, element type, quantization scale granularity, and payload bytes per token;
• final normalization and any logit scaling or clipping.

Names such as “LLaMA-like” or “GQA” leave several of these fields undefined. The serialized model configuration, checkpoint tensors, training code, and serving engine must agree on the same contract.

Validate before the full training run

1. Check shapes and parameter counts. Derive every projection shape and reconcile the total with serialized tensors, including tied parameters.
2. Test causality. Perturb future tokens and confirm earlier logits do not change. Include padding, packed documents, and prefix masks where used.
3. Compare a reference attention path. On small tensors, compare the optimized kernel with a straightforward high-precision implementation and declare numerical tolerances.
4. Test cached equivalence. Confirm token-by-token decode logits match a full-prefix forward pass within tolerance at every position.
5. Measure cache bytes. Compare allocated and payload memory against the declared formula across batch sizes and sequence lengths.
6. Exercise position boundaries. Test the trained length, the deployed maximum, and offsets used by prefix caching or sequence packing.
7. Ablate architecture choices. Compare normalization, FFN, head sharing, and position candidates under matched parameter, token, and compute budgets.
8. Benchmark end to end. Report training throughput, prefill latency, decode latency, maximum resident batch, and quality on the target workload.

The block is compact enough to draw on one page, but its contract reaches from token ids to serving memory. Precision comes from keeping the mathematical operation, its tensor layout, and its runtime representation distinct.

Further reading