Transformer Architecture

Why This Matters

The transformer is the architecture behind every modern large language model: GPT-4, Claude, Gemini, Llama. It replaced recurrent and convolutional architectures for sequence modeling because of one key property: self-attention allows every token to attend to every other token in parallel, enabling the model to learn long-range dependencies without the vanishing gradient problem of RNNs.

Understanding the transformer mathematically (not just as a diagram, but as a sequence of matrix operations with specific dimensions, costs, and properties) is essential for understanding everything built on top of it: RLHF, mechanistic interpretability, scaling laws, and efficiency research.

Mental Model

A transformer processes a sequence of tokens by passing them through a stack of identical blocks. Each block has two sub-layers: a self-attention layer (which lets tokens communicate with each other) and a feed-forward network (which processes each token independently). Residual connections and layer normalization stabilize training.

Self-attention is the key innovation. Each token creates a query ("what am I looking for?"), a key ("what do I contain?"), and a value ("what do I contribute?"). Tokens attend to each other based on query-key similarity, and the output is a weighted sum of values.

Formal Setup and Notation

Let the input sequence have $n$ tokens, each represented as a $d$-dimensional vector. The input is a matrix $X \in \mathbb{R}^{n \times d}$.

Self-Attention

Definition

Scaled Dot-Product Attention

Given an input $X \in \mathbb{R}^{n \times d}$, compute queries, keys, and values:

$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$

where $W_Q, W_K \in \mathbb{R}^{d \times d_k}$ and $W_V \in \mathbb{R}^{d \times d_v}$ are learned weight matrices.

The attention output is:

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

where the softmax is applied row-wise (each row sums to 1).

Proposition

Attention Dimensions

Statement

For input $X \in \mathbb{R}^{n \times d}$:

• $Q \in \mathbb{R}^{n \times d_k}$, $K \in \mathbb{R}^{n \times d_k}$, $V \in \mathbb{R}^{n \times d_v}$
• $QK^\top \in \mathbb{R}^{n \times n}$: the attention matrix
• $\text{softmax}(QK^\top / \sqrt{d_k}) \in \mathbb{R}^{n \times n}$: each row is a probability distribution
• $\text{Attention}(Q, K, V) \in \mathbb{R}^{n \times d_v}$: the output

The output of attention for token $i$ is a weighted average of value vectors:

$\text{output}_i = \sum_{j=1}^{n} \alpha_{ij} \, v_j \quad \text{where } \alpha_{ij} = \frac{\exp(q_i \cdot k_j / \sqrt{d_k})}{\sum_{\ell} \exp(q_i \cdot k_\ell / \sqrt{d_k})}$

Intuition

Each token $i$ computes a query $q_i$ and compares it against all keys $k_1, \ldots, k_n$ via dot products. The softmax converts these similarity scores into attention weights $\alpha_{ij}$. The output is a weighted sum of value vectors, where tokens with similar query-key pairs get higher weight. The $\sqrt{d_k}$ scaling prevents the dot products from becoming too large (which would cause the softmax to saturate).

Why It Matters

Tracking dimensions through the transformer is the single most useful exercise for understanding the architecture. Every research paper assumes you can do this fluently. The $n \times n$ attention matrix is both the source of the transformer's power (global context) and its main computational bottleneck.

Why scale by $\sqrt{d_k}$? If the entries of $q$ and $k$ are independent with zero mean and unit variance, then $q \cdot k$ has variance $d_k$. Without scaling, large $d_k$ causes the dot products to have large magnitude, pushing the softmax into regions with near-zero gradients. Dividing by $\sqrt{d_k}$ normalizes the variance to 1, keeping the softmax in a useful range.

Multi-Head Attention

Definition

Multi-Head Attention

Instead of computing a single attention function, use $h$ heads in parallel:

$\text{head}_i = \text{Attention}(XW_Q^{(i)}, XW_K^{(i)}, XW_V^{(i)})$

where $W_Q^{(i)}, W_K^{(i)} \in \mathbb{R}^{d \times d_k}$ and $W_V^{(i)} \in \mathbb{R}^{d \times d_v}$ with $d_k = d_v = d/h$.

Concatenate the heads and project:

$\text{MHA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W_O$

where $W_O \in \mathbb{R}^{d \times d}$.

Why multiple heads? Each head can attend to different aspects of the input: one head might focus on syntactic relationships, another on semantic similarity, another on positional proximity. Multi-head attention allows the model to jointly attend to information from different representation subspaces. Mechanistic interpretability work gives concrete examples of specialization: previous-token heads that copy information from position $t-1$ to position $t$ (Elhage et al., 2021), induction heads that implement in-context pattern completion (Olsson et al., 2022), and name mover heads that move subject tokens to the final position in the indirect object identification circuit (Wang et al., 2022).

Parameter count for MHA: Each head has $W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}$ each of size $d \times (d/h)$. With $h$ heads, total QKV parameters are $3 \times h \times d \times (d/h) = 3d^2$. The output projection $W_O$ adds $d^2$. Total: $4d^2$ parameters (ignoring biases).

Residual Connections and Layer Normalization

Definition

Transformer Sub-Layer with Residual Connection

Each sub-layer (attention or FFN) is wrapped with a residual connection:

$\text{output} = x + \text{SubLayer}(x)$

This allows gradients to flow directly through the network and enables training of deep transformers.

Definition

Layer Normalization

For a vector $x \in \mathbb{R}^d$, layer normalization computes:

$\text{LayerNorm}(x) = \frac{x - \mu}{\sigma} \odot \gamma + \beta$

where $\mu = \frac{1}{d}\sum_i x_i$, $\sigma = \sqrt{\frac{1}{d}\sum_i (x_i - \mu)^2 + \epsilon}$, and $\gamma, \beta \in \mathbb{R}^d$ are learned scale and shift parameters.

Pre-norm vs. post-norm. The original transformer (Vaswani et al., 2017) uses post-norm: $\text{LayerNorm}(x + \text{SubLayer}(x))$. Most modern LLMs use pre-norm: $x + \text{SubLayer}(\text{LayerNorm}(x))$. Pre-norm is more stable for training deep networks because the residual path is unobstructed.

Feed-Forward Network

Definition

Position-wise Feed-Forward Network

The original 2017 transformer FFN applies two linear transformations with a nonlinearity:

$\text{FFN}(x) = W_2 \, \sigma(W_1 x + b_1) + b_2$

where $W_1 \in \mathbb{R}^{d_{\text{ff}} \times d}$, $W_2 \in \mathbb{R}^{d \times d_{\text{ff}}}$, and $\sigma$ is ReLU (Vaswani et al., 2017) or GeLU in later variants. The standard choice is $d_{\text{ff}} = 4d$.

Parameter count for the 2017 FFN. $W_1$ has $d \cdot d_{\text{ff}}$ parameters, $W_2$ has $d_{\text{ff}} \cdot d$ parameters. With $d_{\text{ff}} = 4d$: total is $8d^2$ parameters (ignoring biases).

Definition

Gated FFN (SwiGLU / GeGLU)

Since 2023, essentially every frontier LLM (Llama 2, Llama 3, Mistral, Mixtral, DeepSeek, Qwen, Gemma) replaces the two-matrix FFN with a gated variant that has three projection matrices $W_1, W_2, W_3$:

$\text{FFN}_{\text{SwiGLU}}(x) = W_2 \left( \sigma(W_1 x) \odot W_3 x \right)$

where $\odot$ is elementwise product, $W_1, W_3 \in \mathbb{R}^{d_{\text{ff}} \times d}$, $W_2 \in \mathbb{R}^{d \times d_{\text{ff}}}$, and $\sigma$ is SiLU (SwiGLU) or GeLU (GeGLU). The gate $W_3 x$ modulates the activation elementwise.

Parameter count for gated FFN. Three projections give $3 \cdot d \cdot d_{\text{ff}}$ parameters. To keep the parameter budget comparable to the 2017 FFN, Shazeer (2020) recommends $d_{\text{ff}} = \frac{2}{3} \cdot 4d = \frac{8d}{3}$, which yields $3 \cdot d \cdot \frac{8d}{3} = 8d^2$ parameters. This matches the legacy $12d^2$ per-layer formula used below. Llama 2 7B ($d = 4096$) picks $d_{\text{ff}} = 11008 \approx 8 \cdot 4096 / 3 = 10922$, rounded up to a multiple of 256 for kernel alignment.

The role of the FFN. Geva et al. (2021) proposed that FFN layers act as key-value memories: $W_1$ maps inputs to a high-dimensional space where patterns are detected, and $W_2$ maps back to the residual stream with the associated information. Under this view, the FFN is where factual knowledge is primarily stored. This is an empirical interpretation from mechanistic interpretability, not a settled fact. Hase et al. (2023) show that the location where Causal Tracing localizes a fact is not a reliable predictor of which layer is best to edit, which complicates the simple "localization equals storage" reading.

Positional Encoding

Self-attention is permutation-equivariant: shuffling the input tokens shuffles the output tokens identically. Without positional information, the model cannot distinguish "the dog bit the man" from "the man bit the dog."

Definition

Sinusoidal Positional Encoding

The original transformer uses fixed sinusoidal encodings added to the input embeddings:

$\text{PE}(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d}}\right), \quad \text{PE}(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d}}\right)$

for position $pos$ and dimension $i$. This allows the model to attend to relative positions because $\text{PE}(pos + k)$ is a linear function of $\text{PE}(pos)$.

Definition

Rotary Position Embedding (RoPE)

RoPE (Su et al., 2021) encodes position by rotating the query and key vectors in $d/2$ independent 2D subspaces:

$\tilde{q}_m = R_m q, \quad \tilde{k}_n = R_n k$

where $R_m$ is a block-diagonal matrix of $d/2$ planar rotations, with the $i$-th block a 2x2 rotation by angle $m \theta_i$ using the base frequencies

$\theta_i = 10000^{-2i/d}, \quad i = 0, 1, \ldots, d/2 - 1.$

Rotations that act on the same 2D subspace commute and satisfy $R_m^\top R_n = R_{n-m}$ block-by-block. Since $R_m$ is block-diagonal, the full matrix product also satisfies $R_m^\top R_n = R_{n-m}$, so the attention score becomes:

$\tilde{q}_m^\top \tilde{k}_n = q^\top R_m^\top R_n k = q^\top R_{n-m} k.$

This depends only on the relative position $n - m$, giving the model translation-invariant attention.

Why RoPE dominates. RoPE naturally encodes relative positions (not absolute), extrapolates better to longer sequences than seen during training, and does not add parameters. It is used in Llama, Mistral, and most modern open-source LLMs.

Computational Complexity

Proposition

Attention is Quadratic in Sequence Length

Statement

The computational cost of self-attention is:

$O(n^2 d)$

The $n^2$ factor comes from computing the $n \times n$ attention matrix $QK^\top$. The memory cost for storing attention weights is $O(n^2)$ per head.

For a full transformer with $L$ layers and $h$ heads:

• Attention cost per layer: $O(n^2 d)$
• FFN cost per layer: $O(n d \cdot d_{\text{ff}}) = O(nd^2)$ with $d_{\text{ff}} = 4d$
• Total cost: $O(L(n^2 d + n d^2))$

Intuition

Every token must attend to every other token, producing an $n \times n$ matrix. For short sequences ($n \ll d$), the FFN dominates. For long sequences ($n \gg d$), attention dominates. This is why extending context length is hard: doubling $n$ quadruples the attention cost.

Why It Matters

The quadratic cost is the fundamental bottleneck for long-context models. A model processing 100K tokens needs attention matrices with $10^{10}$ entries per layer. This has motivated extensive research into efficient attention: sparse attention, linear attention, FlashAttention (which reduces memory but not FLOPs), and sub-quadratic architectures like Mamba.

Failure Mode

The $O(n^2)$ scaling is for standard dense attention. Methods like FlashAttention reduce the memory cost from $O(n^2)$ to $O(n)$ by computing attention in tiles, but the compute cost remains $O(n^2 d)$. True sub-quadratic compute requires architectural changes (sparse or linear attention), which can reduce model quality.

Parameter Counting

Proposition

Transformer Parameter Count

Statement

A decoder-only transformer with $L$ layers has approximately:

$\text{Total params} \approx V \cdot d + L \cdot 12d^2 + V \cdot d$

Breaking this down:

• Token embedding: $V \times d$ parameters
• Per layer:
  • Multi-head attention (QKV + output): $4d^2$
  • FFN (two linear layers): $8d^2$
  • Layer norm (2 per layer): $4d$ (negligible)
  • Subtotal: $\approx 12d^2$ per layer
• Output projection (often tied with embedding): $V \times d$

For GPT-3 scale ($L = 96$, $d = 12288$, $V = 50257$): approximately $12288 \times 50257 + 96 \times 12 \times 12288^2 \approx 175$B parameters.

This is an order-of-magnitude approximation. It ignores biases, layer norm scale and shift, the final output projection head, and absolute positional embeddings, and it double-counts in the presence of weight tying. Many modern implementations tie the input embedding $V \times d$ with the output unembedding matrix, so only one $V \times d$ block is counted. The headline 175B figure works out because the omitted and double-counted terms approximately cancel against the neglected layer norm and bias parameters, not because $12d^2 L$ plus two embedding copies is exact.

Intuition

The vast majority of parameters are in the transformer layers, not the embeddings (unless the vocabulary is very large). Within each layer, the FFN contains $2/3$ of the parameters ($8d^2$ vs. $4d^2$ for attention). This is why the FFN layers are where most of the model's knowledge capacity resides.

Why It Matters

Parameter counting is essential for: (1) estimating compute costs for training and inference, (2) understanding scaling laws, (3) comparing architectures, and (4) estimating memory requirements. A model with $N$ parameters in fp16 requires $2N$ bytes of memory just for weights, plus additional memory for activations and optimizer states. Techniques like speculative decoding and quantization reduce these costs at serving time.

A Complete Transformer Block

Putting it all together, one transformer block computes (using pre-norm):

$h = x + \text{MHA}(\text{LayerNorm}(x))$ $\text{output} = h + \text{FFN}(\text{LayerNorm}(h))$

The full model stacks $L$ such blocks, preceded by token embedding + positional encoding and followed by a final layer norm and linear output projection to vocabulary logits.

Common Confusions

Watch Out

Attention is not a learned weight matrix

The attention weights $\alpha_{ij}$ are computed dynamically from the input. They change for every input sequence. The learned parameters are $W_Q, W_K, W_V, W_O$, which determine how attention is computed. This input-dependence is what gives transformers their flexibility compared to fixed-weight architectures.

Watch Out

Multi-head attention does not multiply the cost by h

Each head operates on $d/h$ dimensions, so the total computation across all heads is the same as a single head with full $d$ dimensions. Multi-head attention is a reorganization of computation, not a multiplication.

Watch Out

FlashAttention reduces memory, not FLOPs

FlashAttention computes the same mathematical operation as standard attention. It reduces memory from $O(n^2)$ to $O(n)$ by computing attention in blocks and never materializing the full $n \times n$ matrix. But the number of floating-point operations is unchanged. True compute savings require architectural changes.

Summary

• Self-attention: $\text{Attention}(Q,K,V) = \text{softmax}(QK^\top / \sqrt{d_k})V$
• Multi-head attention: $h$ parallel heads with $d_k = d/h$, concatenated and projected
• Each transformer block: attention + residual + LayerNorm + FFN + residual + LayerNorm
• Attention cost is $O(n^2 d)$. quadratic in sequence length
• FFN cost is $O(n d^2)$. dominates for short sequences
• Per-layer parameters: $\approx 12d^2$ (attention $4d^2$ + FFN $8d^2$). Modern LLMs replace the two-matrix FFN with a SwiGLU or GeGLU gated FFN that has three projections and $d_{\text{ff}} \approx 8d/3$, which preserves the $8d^2$ budget. Mixture-of-experts variants sparsely activate a subset of FFN parameters.
• RoPE gives relative position encoding via rotation of Q and K
• Pre-norm (LayerNorm before sub-layer) is standard in modern LLMs

Exercises

ExerciseCore

Problem

For a transformer with $d = 512$, $h = 8$ heads, and $d_{\text{ff}} = 2048$, compute the number of parameters in one transformer block (ignoring biases and layer norm parameters).

Hint

Reveal Solution

ExerciseCore

Problem

If the sequence length doubles from $n$ to $2n$, by what factor does the attention computation cost increase? By what factor does the FFN computation cost increase?

Hint

Reveal Solution

ExerciseAdvanced

Problem

Show that without positional encoding, self-attention is permutation-equivariant: if you permute the input tokens by a permutation $\sigma$, the output tokens are permuted by the same $\sigma$.

Hint

Reveal Solution

ExerciseAdvanced

Problem

A transformer model has $L = 32$ layers, $d = 4096$, $h = 32$ heads, $d_{\text{ff}} = 11008$ (as in Llama 2 7B, which uses a SwiGLU-gated FFN with $d_{\text{ff}} \approx 8d/3$, rounded to a multiple of 256 for kernel alignment), and vocabulary $V = 32000$. Estimate the total parameter count and the memory required to store weights in fp16.

Hint

Reveal Solution

Related Comparisons

• Autoregressive Models vs. Diffusion Models
• Autoregressive Models vs. JEPA
• Dense Transformers vs. Mixture-of-Experts
• Transformer vs. Mamba vs. TTT

References

Canonical:

• Vaswani et al., "Attention Is All You Need" (2017). The original transformer paper, Sections 3.1-3.3 and 3.5 for architecture and positional encoding.

Current:

• Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding" (2021). RoPE, Sections 3.1-3.4 for the $\theta_i = 10000^{-2i/d}$ construction.
• Shazeer, "GLU Variants Improve Transformer" (2020). Motivates SwiGLU and GeGLU, and recommends shrinking $d_{\text{ff}}$ by $2/3$ to match the 2017 parameter budget.
• Touvron et al., "Llama 2: Open Foundation and Fine-Tuned Chat Models" (2023). SwiGLU-gated FFN with $d_{\text{ff}} = 11008$ at $d = 4096$, Section 2.
• Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (2022).

Mechanistic interpretability:

• Elhage et al., "A Mathematical Framework for Transformer Circuits" (Anthropic, 2021). Previous-token heads and the residual-stream view.
• Olsson et al., "In-context Learning and Induction Heads" (Anthropic, 2022). Induction heads as a mechanism for in-context learning.
• Wang et al., "Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 Small" (2022). Name mover heads.
• Geva et al., "Transformer Feed-Forward Layers Are Key-Value Memories" (EMNLP 2021). FFN-as-memory hypothesis.
• Hase et al., "Does Localization Inform Editing? Surprising Differences in Causality-Based Localization vs. Knowledge Editing in Language Models" (NeurIPS 2023). Evidence that Causal Tracing localizations do not predict editable layers.

Textbooks:

• Jurafsky & Martin, Speech and Language Processing (3rd ed., draft), Chapters 7-12.
• Goodfellow, Bengio, Courville, Deep Learning (2016), Chapters 10-12.

Next Topics

The natural next steps from transformer architecture:

• Mechanistic interpretability: what do the attention heads and FFN layers actually compute?
• Hallucination theory: why the next-token prediction objective leads to confabulation
• RLHF and alignment: fine-tuning the transformer for human preferences
• Vision transformer lineage: how the transformer was adapted for computer vision (ViT, Swin, DINO, CLIP)