04 · C · Mix

Causal attention

The only place tokens mix. Scale, mask, softmax, GQA, and the softmax Jacobian.

Queries, keys, values

From normalized input $U \in R^{T \times d}$ , three linear maps produce Q, K, V. They are then split into $H$ heads of width $d_{h} = d / H$ . If RoPE is used, Q and K are rotated now.

Q = U W_{Q}, K = U W_{K}, V = U W_{V}

(4.1)

$W_{Q}, W_{K}, W_{V}$: projections in R^{d × d} (or d × H d_h)
$Q_{h}, K_{h}, V_{h}$: head h, each T × d_h

Scores, scale, causal mask, softmax

S_{h} = \frac{Q _{h} K _{h}^{⊤}}{d _{h}} \in R^{T \times T}

(4.2)

Why it mattersA typical coordinate of q and k is O(1), so a typical dot product is O(√dₕ). The √dₕ scale (Vaswani et al., 2017) puts scores back at O(1). Without it, softmax saturates as dₕ grows and the Jacobian vanishes.

Future tokens must be invisible. The causal mask writes $- \infty$ above the diagonal so those entries become zero after softmax.

M_{t, s} = {0 - \infty s \leq t s > t

(4.3)

A_{h} = softmax (S_{h} + M), (A_{h})_{t, s} = \frac{exp (( S _{h} + M ) _{t, s} )}{\sum _{s^{'}} exp (( S _{h} + M ) _{t, s^{'}} )}

(4.4)

$(A_{h})_{t, :}$: a probability distribution over past-and-present positions

Why it mattersAttention is the only place tokens mix. Everything else is applied positionwise. The causal mask is what makes the model a next-token predictor rather than a bidirectional encoder.

O_{h} = A_{h} V_{h}, O = [O_{1} ∣ \dots ∣ O_{H}] W_{O}

(4.5)

$W_{O}$: output projection, d × d
$O$: Attn(U), written back into the residual stream

Softmax Jacobian

For one row $a = softmax (z) \in R^{T}$ :

J_{softmax} = \frac{\partial a}{\partial z} = diag (a) - a a^{⊤} \in R^{T \times T}

(4.6)

\frac{\partial a _{i}}{\partial z _{j}} = a_{i} (δ_{ij} - a_{j})

(4.7)

Why it mattersThis Jacobian is symmetric, has null space along 1 (softmax is shift-invariant), and is positive semidefinite on the simplex tangent space. You will meet it again in the p − y collapse.

Grouped-query and multi-query attention

Multi-head attention uses $H_{q} = H_{k v}$ query and key/value heads. Multi-query (Shazeer, 2019) uses one KV head. Grouped-query (Ainslie et al., 2023) sits in between: Llama-3-class models share each KV head across a group of query heads. The score math is the same; the cache is smaller.

g = H_{q} / H_{k v}, K_{⌊ h / g ⌋} is reused for query head h

(4.8)

$H_{q}$: number of query heads
$H_{k v}$: number of key/value heads, 1 ≤ H_kv ≤ H_q
$g$: group size; g = 1 is MHA, H_kv = 1 is MQA

Why it mattersThe KV cache at inference is O(T · n · H_kv · d_h), not O(T · n · H_q · d_h). That memory, not the matmul, is what bounds generation length.