Softmax Jacobian

Differentiate softmax with respect to its input.

Softmax maps a vector to a vector, so its derivative is a full (n, n) matrix rather than a vector. Writing s = softmax(x):

J[i, j] = s_i * (delta_ij - s_j)

which is s_i * (1 - s_i) on the diagonal and -s_i * s_j off it. In one expression:

J = diag(s) - outer(s, s)

Every output depends on every input, because the normalising sum contains them all. That coupling is why the off-diagonal terms are negative: pushing one logit up pushes every other probability down.

Two properties worth checking, and both are tested. J is symmetric, and every row sums to zero, since the outputs always sum to one no matter how you perturb the input.

Return the (n, n) Jacobian for a single input vector of shape (n,).

Function signature

def softmax_jacobian(x):

Example

Input

softmax_jacobian(np.array([0.0, 0.0]))

Output

array([[ 0.25, -0.25],
       [-0.25,  0.25]])

Reasoning

Equal inputs give s = [0.5, 0.5]. The diagonal is 0.5 * (1 - 0.5) = 0.25, and the off-diagonal is -0.5 * 0.5 = -0.25. Each row sums to 0.

Constraints and notes

• The input is one vector of shape (n,).
• The Jacobian is symmetric, and its rows sum to 0.