Chapter 3: Transformer Architecture

The projection matrices $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$ transform each token into three roles. The attention score $q_i^T k_j = x_i^T W_Q^T W_K x_j$ is a bilinear form that measures compatibility between tokens $i$ and $j$ through the learned metric $W_Q^T W_K$. This is more expressive than simple dot-product similarity because the model learns what aspects of tokens should determine attention.

Self-attention computes $\text{softmax}(QK^T/\sqrt{d_k})V$. The $\sqrt{d_k}$ scaling prevents dot products from growing with dimension, keeping softmax gradients useful. Without scaling, if $q$ and $k$ are random vectors with entries $\sim N(0,1)$, then $q^T k$ has variance $d_k$. For $d_k = 64$: std $\approx 8$, and $\text{softmax}([8, -8, 2, ...])$ is nearly one-hot, producing vanishing gradients for all but the top-scoring key.

The causal mask $M_{ij} = -\infty \cdot \mathbb{1}[j > i]$ ensures $\text{softmax}(\ldots + M)_{ij} = 0$ for $j > i$, because $e^{-\infty} = 0$. This enforces the autoregressive property: $P(x_t \mid x_{<t})$ cannot depend on $x_{>t}$. The mask is the only difference between encoder (bidirectional) and decoder (causal) attention — the same attention mechanism with a different mask produces fundamentally different model capabilities.

Computing the $n \times n$ attention matrix $QK^T$ requires $n^2 d$ multiply-adds. Storing this matrix takes $n^2$ floats. For $n = 128{,}000$ (Llama's context): the matrix has $1.64 \times 10^{10}$ entries, requiring $\sim 64$ GB in FP32. FlashAttention avoids materializing this matrix by computing attention in tiles that fit in SRAM, reducing memory from $O(n^2)$ to $O(n)$ while maintaining exact computation.