Define a single attention head, what does one head actually own and produce?

A single **attention head** is the smallest independent attention computation inside a multi-head attention layer. It owns three things and produces one thing. What it owns: its own slice of the query, key, and value projections. For a model with `d_model = 4096` and 32 heads, each head gets a head dimension of `d_head = 128`. Inside one head you have three small projection matrices, W_Q, W_K, and W_V, each of shape `(d_model, d_head) = (4096, 128)`. What it computes: using those projections it produces its own Q, K, V matrices of shape `(seq_len, 128)` per token, then runs scaled dot-product attention: $$\text{Attn}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$ What it produces: a `d_head = 128` dimensional context vector per token. All heads run fully in parallel and never look at each other's outputs inside the attention layer. Their per-token outputs get concatenated along the feature axis to width `d_model`, then a shared output projection W_O of shape `(d_model, d_model)` mixes them back into the residual stream.

A single **attention head** is one independent (Q, K, V) computation inside multi-head attention. It owns its three projection matrices, runs its own softmax, and writes out a vector at its head width. For `d_model = 4096` and `n_heads = 32`, the standard recipe gives `d_head = d_model / n_heads = 128`. Each head owns: - `W_Q^h` of shape `(d_model, d_head)` = (4096, 128) - `W_K^h` of shape `(d_model, d_head)` = (4096, 128) - `W_V^h` of shape `(d_model, d_head)` = (4096, 128) Given input X of shape (seq_len, d_model), the head computes Q, K, V of shape (seq_len, 128) and then `softmax(QK^T / sqrt(d_head)) V`. The output is a (seq_len, 128) context tensor. Heads run in parallel. They never communicate inside the attention layer, this is a load-bearing property because it is what makes multi-head attention embarrassingly parallel across GPU lanes. The integration happens after concatenation: - Concatenate all head outputs along the feature axis: shape becomes `(seq_len, d_model)`. - Apply the shared output projection `W_O` of shape `(d_model, d_model)`. - Add to the residual stream. That output projection is the only place where information from different heads mixes. In implementations the per-head W_Q matrices are usually packed into a single `(d_model, d_model)` W_Q tensor, then split, but the per-head semantics are identical.