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

MCQ·Easy·4.0 · 0·~1 min·Asked atLepton Ai Sierra Together Ai·Relevant atMicrosoft

Attempt it

**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.

TL;DR

Along the key axis (the last dimension). Each query gets a probability distribution over keys; each row sums to 1.

Memory aid

Easy to grasp

Imagine planning a meal where each chef has to decide how much of each ingredient to use. The chef does not normalize 'how much each ingredient gets used across all chefs'. The chef normalizes 'how much of my own meal each ingredient takes up', so for each chef's own dish, the ingredient percentages add up to 100%. Each query token in attention is like one chef: it decides for itself how attention mass is split across all available keys, and those splits add up to 100% for that query.

Key concepts

Softmax Attention score matrix Row-stochastic matrix Attention weight

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Softmax is applied along the key axis (the last dimension of the score matrix), turning each query's row into a probability distribution over keys. The choice is not a convention but a structural consequence of how attention combines value vectors. Flip the axis and the entire attention block breaks: the downstream V matmul becomes malformed, masking loses its renormalization property, and FlashAttention's tiling stops working.

This deep dive derives the axis choice from the convex-combination output, walks why each alternative option fails, and shows how the key-axis convention threads through masking, FlashAttention, cross-attention, and every modern attention modification.

Mental model: each query produces its own distribution over keys, and the distribution sums to 1. Softmax along the key axis is the only choice consistent with the downstream weighted-sum.

Why the key axis is forced by the V matmul

The score matrix per head has shape (T_query, T_key). The downstream computation is:

output = attention_weights @ V

where V has shape (T_key, d_v) and output is (T_query, d_v).

The convex-combination requirement

For each query position i, the output is:

o_i = \sum_{j=1}^{T_k} w_{ij} v_j

This is a convex combination of value rows IF AND ONLY IF the weights satisfy:

w_{ij} >= 0 for every key j.
sum_j w_{ij} = 1.

The sum is over the KEY axis. The constraint is per-query (per row of the weight matrix).

Why this forces the axis choice

Softmax must produce row-stochastic weights, meaning normalization happens along the axis being summed in the V matmul. That axis is the key axis. Any other normalization choice fails the convex-combination property:

Column-stochastic (softmax along query axis): rows do not sum to 1, the per-query weighted sum is not a convex combination.
Batch-stochastic: batch elements get coupled, defeating the independent-sample design.
Head-stochastic: heads get coupled, defeating the multi-head design.

The PyTorch idiom

Production code uses F.softmax(scores, dim=-1). The -1 is the convention for the last dimension, which is the key axis after any batch and head leading dimensions.

Key insight: the axis choice is not a convention. It is forced by the requirement that attention output be a convex combination of value vectors per query.

Why the other axes fail

How masking, FlashAttention, and cross-attention align with the key axis

Common implementation mistakes

Situations where this technique stops working.

2–4 min · Everything important, quickly.

Real products, models, and research that use this idea.

PyTorch's torch.nn.functional.scaled_dot_product_attention computes softmax(scores, dim=-1) internally before multiplying by V.
All modern LLMs (GPT-5.5, Llama 4 Maverick, Claude Opus 4.7, Gemini 3.1 Pro, DeepSeek V4) apply softmax along the key axis in every attention layer.

What an interviewer would ask next. Try answering before peeking at the approach.

QWhat would change if softmax were applied along the query axis instead of the key axis?

Each column would sum to 1, giving each key a distribution over queries. The downstream computation softmax(s, dim=0) @ V would still have shape (T_query, d_v), but the per-query interpretation would be malformed: there is no longer a convex combination of value rows for each query. The output would be a sum of V rows weighted by column-stochastic weights, which has no clean probabilistic meaning. The architecture would have to be redesigned around a different inner product structure.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Applying softmax along the query axis. That would give each key a distribution over queries, not each query a distribution over keys, and the downstream V matmul would be malformed.

60 second bullets to scan on the way to the call.

The shape of the score matrix per head: (T_query, T_key)
Why softmax is applied along the key axis specifically

Primary sources. Browse if you want the original framing.