Transformers

Scaled dot-product attention#

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

Where:

• $Q \in \mathbb{R}^{n \times d_k}$ — queries, one row per position.
• $K \in \mathbb{R}^{m \times d_k}$ — keys.
• $V \in \mathbb{R}^{m \times d_v}$ — values.
• $n$ — number of query positions; $m$ — number of key/value positions. In self-attention $n = m$.
• $d_k$ — key/query dimension; $d_v$ — value dimension.
• The softmax is applied row-wise, so each row of the $n \times m$ attention matrix sums to $1$.

Why $\sqrt{d_k}$: if the components of $q$ and $k$ are independent with zero mean and unit variance, then $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ has mean $0$ and variance $d_k$. Dividing by $\sqrt{d_k}$ restores unit variance, keeping softmax away from its saturated regime where gradients vanish. This argument is given in footnote 4 of Vaswani et al. (2017).

Multi-head attention#

$$ \text{MHA}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\, W^{O}, \qquad \text{head}_i = \text{Attention}!\left(X W_i^{Q},\, X W_i^{K},\, X W_i^{V}\right) $$

Where $h$ is the number of heads, $W_i^{Q}, W_i^{K} \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^{V} \in \mathbb{R}^{d_{\text{model}} \times d_v}$, and $W^{O} \in \mathbb{R}^{h d_v \times d_{\text{model}}}$. Standard practice sets $d_k = d_v = d_{\text{model}} / h$, so multi-head attention costs the same as single-head attention at full width.

Complexity#

The decisive column is the last one. Attention has constant sequential depth, so an entire sequence is processed in parallel during training. That, more than any accuracy gain, is why transformers displaced RNNs — they use modern hardware fully.

Per transformer layer with $d = d_{\text{model}}$ and feedforward width $4d$:

• Attention projections: $4 d^2$ parameters.
• Feedforward: $8 d^2$ parameters.
• Total: $\approx 12 d^2$ per layer, so the MLP holds two-thirds.

Positional information#

Self-attention is permutation-equivariant. Position must be supplied explicitly.

Sinusoidal (Vaswani et al., 2017):

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

Where $p$ is the position index and $i$ the dimension index. Fixed, not learned; relative offsets are expressible as linear transforms of the encoding.

RoPE (Su et al., 2021, RoFormer) rotates query and key vectors by an angle proportional to position, making the dot product depend only on relative distance $p - q$. It is the standard in Llama, Qwen, Mistral and most current open models, largely because it extrapolates beyond the trained context length more gracefully.

ALiBi (Press et al., 2022) adds a linear distance penalty directly to attention scores, with no positional embedding at all.

Efficiency work that matters in practice#

• FlashAttention (Dao et al., 2022; Dao, 2023 for v2) computes exact attention with $O(n)$ memory by tiling the computation in SRAM and never writing the $n \times n$ matrix to HBM. It is exact, not an approximation — a pure systems optimisation, and it is why long contexts became affordable.
• Multi-query attention (Shazeer, 2019) shares one key/value head across all query heads, shrinking the KV cache by a factor of $h$. Grouped-query attention (Ainslie et al., 2023) interpolates between MQA and MHA and is used in Llama 2 70B onward.
• KV caching is essential at inference. Without it, generating token $t$ recomputes attention over all $t-1$ previous tokens, making generation $O(n^2)$ per token. Cache size is $2 \cdot L \cdot n \cdot h_{kv} \cdot d_{\text{head}} \cdot \text{bytes}$, and it, not the weights, dominates memory during long-context serving.
• Sparse and linear attention — Longformer, BigBird, Performer, Linformer — reduce the asymptotic cost but have largely lost to FlashAttention plus dense attention in practice, since exact attention on modern hardware is fast enough.

Architecture families#

Decoder-only won for general-purpose generation, largely because next-token prediction on raw text scales without labelled data.

Key papers#

• Vaswani et al. (2017), Attention is all you need — the architecture.
• Bahdanau et al. (2015), Neural machine translation by jointly learning to align and translate — attention itself, predating transformers by two years and rarely credited.
• Devlin et al. (2019), BERT — bidirectional pretraining.
• Kaplan et al. (2020) and Hoffmann et al. (2022), Training compute-optimal large language models (Chinchilla) — the scaling laws. Chinchilla corrected Kaplan and showed most large models of that era were badly undertrained for their size.
• Xiong et al. (2020), On layer normalization in the transformer architecture — the pre-norm versus post-norm analysis that changed standard practice.

Honest open problems#

Quadratic scaling is a real architectural limit, softened but not removed by FlashAttention. Interpretability remains weak: attention weights are not reliable explanations of model behaviour (Jain and Wallace, 2019, Attention is not explanation), a point frequently ignored in applied write-ups. State-space models such as Mamba (Gu and Dao, 2023) offer linear scaling and are competitive at moderate scale, though they have not displaced transformers at the frontier.

What to learn next#

• Attention — the mechanism on its own, in more depth.
• What is a large language model? — what gets built from stacked transformer blocks.
• PyTorch basics — the tools for building one yourself.