Guides / Matrix calculus

The softmax Jacobian: diag(s) − ssᵀ, a backward pass without the matrix, and s − y for cross-entropy

Softmax turns a vector of scores z into probabilities s_i = e^z_i/Σ_k e^z_k. Every s_i shares the same denominator, so every output depends on every input, and the Jacobian is a full matrix, diag(s) − ss^T. A backward pass never needs to build it: its product with a vector takes one dot product. Put cross-entropy on top and the whole thing collapses to s − y. The mistakes come from keeping only the diagonal, from stopping the chain rule at s, and from computing the exponentials without subtracting the largest score.

Why it works

Write Σ = Σ_k e^z_k, so s_i = e^z_i/Σ. Raising z_j raises Σ by e^z_j, which lowers every s_i in proportion to s_is_j; for i = j the numerator also rises, adding s_i. That is s_i(δ_ij − s_j). The probabilities always add to 1, so a change in z can only move probability between classes, and every column of J sums to 0.

J is symmetric, so J^Tg = Jg = s ⊙ g − s(s^Tg): one dot product, then an elementwise step, instead of an n × n matrix. For cross-entropy, L = −log s_c = −z_c + log Σ for the target class c, and its derivative in z_j is s_j − δ_jc. Adding the same constant to every score multiplies the numerator and Σ by the same factor, which is why the shift by the maximum is free.

Example 1: softmax of three scores

Find softmax(z) for z = (2, 1, 0).

1. e², e¹, e⁰The largest score is 2; subtracting it gives (0, −1, −2) and the same answer. These are small enough either way.
2. ≈ 7.389, 2.718, 1The three exponentials.
3. Σ ≈ 11.107Their sum, the shared denominator.
4. s ≈ (0.6652, 0.2447, 0.0900)Each exponential over the sum. They add to 1. A gap of 1 in the scores is a factor of e ≈ 2.718 in the probabilities.

Example 2: the Jacobian, entry by entry

Find ∂s/∂z at the scores of Example 1.

1. s₁(1 − s₁) ≈ 0.2227The diagonal entries, s_i(1 − s_i) …
2. 0.1848 and 0.0819… for i = 2 and 3.
3. −s₁s₂ ≈ −0.1628The off-diagonal entries, −s_is_j …
4. −0.0599 and −0.0220… for (1, 3) and (2, 3). J is symmetric, so these fill both sides.
5. J ≈ .2227−.1628−.0599−.1628.1848−.0220−.0599−.0220.0819Each row and column sums to 0. A finite-difference Jacobian of softmax at z agrees to every digit shown.

Example 3: the cross-entropy gradient and one step

With the scores of Example 1, the target is class 2, so y = (0, 1, 0). Find the loss, ∇_zL, and the loss after one gradient step of size 1.

1. L = −log s₂Only the target's term survives.
2. ≈ −log 0.2447 ≈ 1.4076From Example 1.
3. ∇_zL = s − ySoftmax and cross-entropy together.
4. ≈ (0.6652, −0.7553, 0.0900)The target's entry is negative; the others are their probabilities. They sum to 0.
5. z − ∇_zLOne step downhill.
6. ≈ (1.3348, 1.7553, −0.0900)The target's score went up, and the others went down.
7. s₂ ≈ 0.5511, L ≈ 0.5959Softmax of the new scores. The step moved the target from the second most likely class to the most likely.

Example 4: a backward pass without the matrix

Inside a network, softmax receives the gradient g = (1, 0, −1) from the layer above. With s from Example 1, find J^Tg without forming J.

1. s^Tg ≈ 0.6652 − 0.0900One dot product.
2. ≈ 0.5752The s-weighted average of g.
3. g − 0.5752 · 1Subtract it from every entry.
4. ≈ (0.4248, −0.5752, −1.5752)g, centered at its s-weighted average.
5. s ⊙ (g − (s^Tg)1)Multiply entry by entry by s.
6. ≈ (0.2826, −0.1408, −0.1418)The same as J times g with the matrix of Example 2, in O(n) work instead of O(n²). The entries sum to 0, as every output of J must.

Common mistakes, and the exact line they happen on

1. Keeping only the diagonal

The backward pass of Example 4.

1. ∂s_i/∂z_i = s_i(1 − s_i)Right so far: that is the diagonal, and it looks just like the sigmoid's derivative.
2. s ⊙ (1 − s) ⊙ gThis is the mistake. The sigmoid acts elementwise, so its Jacobian is diagonal. Softmax doesn't: every s_i depends on every z_j through Σ. This gives (0.2227, 0, −0.0819), which doesn't sum to 0. The right answer is (0.2826, −0.1408, −0.1418).

2. Stopping the chain rule at s

The cross-entropy gradient of Example 3.

1. ∇_sL = −y/sRight so far: the derivative of −log s₂ with respect to s.
2. ∇_zL ≈ (0, −4.086, 0)This is the mistake. That is still the gradient with respect to s. One more step through softmax, J^T(−y/s), gives s − y ≈ (0.6652, −0.7553, 0.0900). The wrong version never pushes the other scores down, and its size, 1/s₂, blows up as the prediction gets worse.

3. Exponentiating large scores directly

Scores of z = (1000, 999, 998), which differ by the same gaps as Example 1.

1. s_i = e^z_i/ΣRight so far as math.
2. e¹⁰⁰⁰ = ∞This is the mistake. In floating point, e^z overflows once z passes about 709 (64-bit) or 88 (32-bit), and ∞/∞ is NaN. Subtract the largest score first: (0, −1, −2) gives the answer of Example 1, (0.6652, 0.2447, 0.0900), because softmax ignores a shift of every score.

Practice

1. Find softmax(0, 0, 0, 0).

   Show answer

   (1/4, 1/4, 1/4, 1/4).

2. Find softmax(ln 3, 0).

   Show answer

   (3/4, 1/4): the exponentials are 3 and 1.

3. For s = (0.5, 0.3, 0.2), find ∂s₁/∂z₂ and ∂s₁/∂z₁.

   Show answer

   −s₁s₂ = −0.15 and s₁(1 − s₁) = 0.25.

4. A classifier outputs s = (0.7, 0.2, 0.1) and the target is class 3. Find ∇_zL for cross-entropy.

   Show answer

   s − y = (0.7, 0.2, −0.9).

5. A batch of 4 examples uses the mean cross-entropy. One row has s = (0.6, 0.4) with target class 1. What is that row of ∇_ZL?

   Show answer

   (s − y)/4 = (−0.1, 0.1): the mean divides every row by the batch size.