Transformer architecture

The idea

A transformer is a stack of identical blocks, each doing two things to a sequence of token vectors: mix across positions (multi-head attention), then process each position (a feed-forward network applied identically at every position). Everything else — residuals, layer norms, positional encodings — exists to make that stack trainable and order-aware. One pre-LN block:

$x \leftarrow x + \mathrm{MHA}(\mathrm{LN}(x)), \qquad x \leftarrow x + \mathrm{FFN}(\mathrm{LN}(x))$

Multi-head attention. Rather than one attention pattern over full-width vectors, split into $h$ heads of dimension $d_k = d_{\text{model}}/h$ :

\begin{aligned} \mathrm{head}_i &= \mathrm{Attention}(XW_i^Q,\, XW_i^K,\, XW_i^V), \\ \mathrm{MHA}(X) &= \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_h)\, W^O. \end{aligned}

Same total projection dimensions as one big head, but $h$ separately normalized attention patterns per token. Some trained heads exhibit interpretable positional or syntactic behavior, but specialization is not guaranteed and many heads can be redundant or pruned.

Position-wise FFN. $\mathrm{FFN}(x) = W_2\, \mathrm{GELU}(W_1 x + b_1) + b_2$ with hidden width $4d_{\text{model}}$ . It is where most parameters live, and it acts per token — attention is the only place information crosses positions.

Positional encodings. Attention is permutation-equivariant, so order must be injected. Sinusoidal encodings add a multi-wavelength sin/cos fingerprint of each position to its embedding:

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

Learned absolute embeddings (BERT) and rotary embeddings (most modern LLMs) are the common alternatives.

Why not RNNs? Three answers an interviewer wants, in order: (1) parallelism — RNN training is sequential in $t$ because $h_t$ needs $h_{t-1}$ ; a transformer trains on all positions at once as dense matrix algebra; (2) path length — relating tokens $n$ apart takes $n$ recurrent hops (gradients decay across each) versus exactly one attention hop; (3) scaling — transformers kept converting more data and compute into better loss. The price: $O(n^2)$ attention versus $O(n)$ recurrence, and $O(n)$ inference memory (KV cache) versus an RNN's $O(1)$ state.

Worked example

Parameter budget of one encoder block at $d_{\text{model}} = 512$ , $h = 8$ , FFN width 2048 (the original base config):

Component	Count	Params
$W^Q, W^K, W^V$	$3 \times 512 \times 512$	$786{,}432$
$W^O$	$512 \times 512$	$262{,}144$
FFN in	$512 \times 2048$	$1{,}048{,}576$
FFN out	$2048 \times 512$	$1{,}048{,}576$
2 LayerNorms	$2 \times 2 \times 512$	$2{,}048$
Total		$\approx 3.15$ M

Read off the structure: two thirds of every block is FFN, attention projections are the other third, norms are noise. A 6-block stack is $\approx 18.9$ M parameters before embeddings; a 32k-token vocabulary at $d = 512$ adds another $\approx 16.4$ M — for small models the embedding table rivals the whole stack, which is why input/output embeddings are often tied.

Sanity-check the multi-head shapes on a 10-token sentence: input $10 \times 512$ → per head, $Q, K, V$ are $10 \times 64$ → scores $10 \times 10$ per head (8 of them) → weighted values $10 \times 64$ → concat $10 \times 512$ → $W^O$ → $10 \times 512$ . The sequence length only ever appears in the $10 \times 10$ score matrices — that is the $O(n^2)$ core in the flesh.

Cheat sheet

Block (pre-LN): $x + \mathrm{MHA}(\mathrm{LN}(x))$ , then $x + \mathrm{FFN}(\mathrm{LN}(x))$ ; FFN $= W_2\,\mathrm{GELU}(W_1 x)$ , width $4d$ .

MHA: $h$ heads of $d_k = d/h$ ; params $4d^2$ (Q, K, V, O); compute ≈ one full-width head; heads = independent attention patterns.

Positional encodings: sinusoidal (fixed, extrapolates in principle, relative offsets are linear maps); learned absolute (BERT); rotary/RoPE (relative, inside the QK product — modern default).

	RNN / LSTM	Transformer
Training over length $n$	sequential, $O(n)$ steps	parallel, $O(1)$ depth
Max dependency path	$O(n)$	$O(1)$
Compute per layer	$O(n d^2)$	$O(n^2 d + n d^2)$
Inference state	$O(1)$	KV cache $O(n)$
Long-range gradients	decay through gates	direct

Encoder-only blocks usually use bidirectional self-attention; decoder-only blocks use causal self-attention. Classical encoder-decoder transformers also add decoder cross-attention to encoder outputs, so the architectures differ by more than one mask.
Pre-LN: stable deep training, mild warmup. Post-LN (original): needs warmup, unstable when deep.
Parameter intuition: FFN ≈ 2/3 of block params; embeddings dominate small models.
BERT-base for scale-anchoring: 12 layers, $d = 768$ , 12 heads, ~110M params.
Complexity gotcha: doubling context length quadruples attention compute but leaves parameter count unchanged.
Attention is the only cross-token operation; FFN, LN, residuals are all per-token.