Guides / Matrix calculus

Attention backward: AᵀG for V, a row-wise softmax for the scores, and the 1/√d on Q and K

Scaled dot-product attention is two matrix products with a row-wise softmax between them: scores S = QK^T/√d, weights A = softmax(S) row by row, and output O = AV. The backward pass runs the same pieces in reverse. Each matrix product's gradient is another matrix product with one factor transposed, as in the batched backprop guide. The softmax step is the softmax Jacobian guide's row formula, applied to every row at once. The conventions are those of the matrix calculus conventions guide, with rows as positions and each gradient shaped like its variable, the same rules the one-hidden-layer backprop guide uses for a whole network. The mistakes come from a missing transpose, from a softmax Jacobian that mixes rows, and from dropping the 1/√d on the way back.

Short answer

With G = ∇_OL (n × d_v):

∇_VL = A^TG and G̃ = ∇_AL = GV^T.

∇_SL = A ⊙ (G̃ − r1^T), where r = rowsum(A ⊙ G̃) holds one number per query row.

∇_QL = (∇_SL)K/√d and ∇_KL = (∇_SL)^TQ/√d.

Every row of ∇_SL sums to 0, a masked score gets exactly 0, and the 1/√d appears only in the gradients of Q and K.

Why it works

For a product C = PR with upstream gradient ∇_CL, the right factor gets P^T∇_CL and the left factor gets (∇_CL)R^T. In each case the transposed matrix is the one whose index the sum runs over. O = AV gives the gradients of V and A directly. S = QK^T/√d gives the gradients of Q and K, with the constant 1/√d carried along, because it scales every score.

In between, row i of A depends only on row i of S, so the softmax's backward pass is the vector formula a ⊙ (g̃ − (a^Tg̃)1) for each row, with its own a^Tg̃. Stacking the rows gives one matrix expression, and no m × m Jacobian is ever built. A softmax row sums to 1 whatever its scores, so adding a constant to a row of S changes nothing, and each row of ∇_SL sums to 0. That is a check worth running on every answer.

Example 1: the forward pass

Take two queries and two keys of width d = 4, so √d = 2: query rows q₁ = (1, 1, 1, 1) and q₂ = (1, −1, 0, 0), key rows k₁ = (1, 1, 0, 0) and k₂ = (0, 0, 1, −1), and V = 1002. Find S, A and O.

QK^T = 2000Entry (i, j) is q_i · k_j: only q₁ · k₁ is nonzero.
S = 1000Divide by √d = 2.
row 1 of A ≈ (0.7311, 0.2689)The softmax of (1, 0) is (e, 1)/(e + 1).
row 2 of A = (0.5, 0.5)Equal scores give equal weights.
O = AV ≈ 0.73110.53790.51Row i of O mixes the rows of V with the weights in row i of A: 0.2689 × 2 = 0.5379.

Example 2: the gradients of V and A

The layer above sends back G = 1−101. Find ∇_VL and G̃ = ∇_AL.

∇_VL = A^TGV is the right factor of AV. (m × n)(n × d_v), the shape of V.
row 1: 0.7311(1, −1) + 0.5(0, 1)Key j collects G's rows, weighted by how much each query attended to it: column j of A.
∇_VL ≈ 0.7311−0.23110.26890.2311Row 2 is 0.2689(1, −1) + 0.5(0, 1).
G̃ = GV^TA is the left factor. (n × d_v)(d_v × m), the shape of A.
G̃ = 1−202Entry (i, j) is row i of G dotted with value row j: (1, −1) · (0, 2) = −2.

Example 3: through the softmax, row by row

Find ∇_SL = A ⊙ (G̃ − r1^T).

r₁ = 0.7311(1) + 0.2689(−2)Row 1 of A ⊙ G̃, summed: the weighted average of row 1 of G̃ …
r₁ ≈ 0.1932… which every entry of that row is measured against.
0.7311 × (1 − 0.1932)Entry (1, 1): the weight times its entry of G̃ minus r₁ …
0.2689 × (−2 − 0.1932)… and entry (1, 2) the same way.
row 1 ≈ (0.5898, −0.5898)It sums to 0. For two keys it is always ±a₁a₂(g̃₁ − g̃₂).
r₂ = 0.5(0) + 0.5(2) = 1Row 2 uses its own weights and its own row of G̃.
∇_SL ≈ 0.5898−0.5898−0.50.5Row 2 is 0.5 × (0 − 1) and 0.5 × (2 − 1). Central finite differences of L in S agree.

Example 4: the gradients of Q and K

Finish with ∇_QL and ∇_KL, both 2 × 4.

∇_QL = (∇_SL)K/2Q is the left factor of QK^T, and the 1/√d comes along.
row 1 = (0.5898k₁ − 0.5898k₂)/2Query 1 is pushed toward key 1, the key whose value it should use more.
≈ 0.2949 (1, 1, −1, 1)k₁ − k₂ = (1, 1, −1, 1).
row 2 = 0.25(−k₁ + k₂)(−0.5k₁ + 0.5k₂)/2 …
= −0.25 (1, 1, −1, 1)… so query 2 is pushed the other way.
∇_KL = (∇_SL)^TQ/2K is transposed in the forward pass, so the sum runs over the query index.
row 1 = (0.5898q₁ − 0.5q₂)/2Column 1 of ∇_SL weights the query rows …
≈ 0.2949q₁ − 0.25q₂… which is (0.0449, 0.5449, 0.2949, 0.2949). Row 2 is its negative, since each row of ∇_SL sums to 0. Finite differences of L in Q and K agree with every entry.

Common mistakes, and the exact line they happen on

1. The gradient to V without the transpose

The first step of Example 2.

O = AV, G givenRight so far.
∇_VL = AGThis is the mistake. "V is multiplied by A, so its gradient is A times the gradient." In self-attention n = m, so the product runs. Here it gives rows (0.7311, −0.4621) and (0.5, 0). The sum runs over A's row index, the queries, so it is A^TG.

2. One softmax over the whole matrix

The softmax step of Example 3.

a ⊙ (g̃ − (a^Tg̃)1) for one rowRight so far.
flatten A and G̃, use one a^Tg̃This is the mistake. That subtracts the whole matrix's total, r₁ + r₂ ≈ 1.1932, from every row. Row 1 becomes (−0.1412, −0.8588), which sums to −1, not 0. Rows don't interact in a row-wise softmax, so each row uses its own r_i.

3. Dropping the 1/√d

The first step of Example 4.

∇_SL from Example 3Right so far.
∇_QL = (∇_SL)KThis is the mistake. The scale is a constant, so it needs no gradient of its own, but it still multiplies every score, and so every gradient through S. Here row 1 comes out as 0.5898 (1, 1, −1, 1), twice the answer; at d = 64 it is 8 times. The direction is right, so a check that compares only directions passes.

Practice

Self-attention runs on n = m = 1024 positions with d = 64. How many entries does S have, and how many multiply-adds does it take to form?
Show answer
1024² = 1,048,576 entries, and 1024² × 64 = 67,108,864 multiply-adds.
Find ∇_VL for A = 0.50.50.250.75 and G = 2004.
Show answer
A^TG has rows (1, 1) and (1, 3).
One row has weights a = (0.5, 0.25, 0.25) and g̃ = (2, 0, 4). Find that row of ∇_SL.
Show answer
a^Tg̃ = 2, so a ⊙ (g̃ − 2) = (0, −0.5, 0.5).
With a causal mask, the first query can see only the first key. What is its row of A, and its row of ∇_SL?
Show answer
A's row is (1, 0, …, 0), and the gradient row is all zeros. The masked entries have weight 0, and the one unmasked weight is 1 whatever its score, so g̃₁ − r₁ = 0.
If the entries of a query and a key are independent with mean 0 and variance 1, and d = 64, what is the standard deviation of q^Tk, and of the scaled score?
Show answer
The variance is d = 64, so the standard deviation is 8. Divided by √d = 8, the score has standard deviation 1.