Along which axis of the QK^T score matrix is softmax applied inside attention?

**The correct answer is B: along the key axis (the last dimension), so each query gets a probability distribution over keys.** The score matrix `s = QK^T / sqrt(d_k)` has shape `(T_query, T_key)` per head. Softmax is applied row-wise, where each row corresponds to one query token and the entries within that row correspond to all the keys that query can attend to. After softmax, each row is a probability distribution: every entry lies in `[0, 1]`, and the entries sum to exactly 1. For each query token, the model has decided how to distribute its attention mass across all the keys. In PyTorch the call is `softmax(scores, dim=-1)` because the key axis is the last dimension of the score tensor. **Why the other options are wrong.** - **Option A (query axis):** Would give each key a probability distribution over queries, which is backwards. The downstream computation `attention_weights @ V` requires each query to have a weighting over keys (and therefore over rows of V), not each key to have a weighting over queries. - **Option C (batch axis):** Softmax over batch is nonsensical because batch elements are independent samples; you do not want one batch element's scores to depend on another's. - **Option D (head axis):** Heads are independent attention computations; softmax over heads would couple them and destroy the multi-head design. **The downstream consequence.** The attention output is `o = softmax(s, dim=-1) @ V`, shape `(T_query, d_v)`. For each query, the output is a convex combination of value vectors with the row's attention weights as mixing coefficients. This only works because each row sums to 1.

**Along the key axis (the last dimension of the score matrix).** Each query gets its own probability distribution over keys, and each row of the post-softmax weight matrix sums to exactly 1. The score matrix per head has shape `(T_query, T_key)`. The softmax operator is applied row-wise: $$w_{ij} = \frac{e^{s_{ij}}}{\sum_k e^{s_{ik}}}$$ where the sum in the denominator runs over the key axis. In PyTorch this is `softmax(scores, dim=-1)`. **Why row-wise specifically.** The downstream computation is: `output = attention_weights @ V` where `V` has shape `(T_key, d_v)` and the output is `(T_query, d_v)`. For each query, the output is a convex combination of value rows. That convex combination requires the weights to sum to 1 over the key axis, which is exactly what row-wise softmax produces. **Why not the other axes.** - **Query axis (option A):** would give each key a distribution over queries, which is the wrong tensor shape for the downstream matmul. The output would be `T_key`-dimensional in a meaningless way. - **Batch axis (option C):** batch elements are independent samples in a minibatch; softmax across batch would couple them, which is never what you want. - **Head axis (option D):** heads are designed to be independent retrieval channels; softmax across heads would destroy that independence. **The systems perspective.** Score-level operations (masking via -inf, temperature scaling, ALiBi bias) all operate BEFORE softmax and affect the same key axis. Causal masking sets `s_{ij} = -inf` for `j > i` along the key axis, so each row's mask is applied along the key dimension and softmax automatically renormalizes the unmasked entries. All of this fits the 'each query gets a distribution over keys' picture; flipping the axis would break the whole pipeline.