Softmax Jacobian and Cross-Entropy Gradient

1. The Jacobian of the Softmax Function

Let's begin by recalling the softmax function, which maps a vector of logits $\mathbf{z}\in {\mathbb{R}}^K$ to a vector of probabilities $\mathbf{s}\in {\mathbb{R}}^K$:

Since softmax is a vector-valued function of a vector input, its derivative is a Jacobian matrix, $J_S$, where each entry $J_{ij}$ is the partial derivative $\frac{\partial s_i}{\partial z_j}$.

To compute these entries, we must consider two distinct cases:

1. Diagonal elements ($i=j$): The derivative of a softmax output with respect to its corresponding logit.
2. Off-diagonal elements ($i\ne j$): The derivative of a softmax output with respect to a different logit.

The following article provides an exceptionally clear, step-by-step derivation for both cases using the quotient rule.

Let's summarize the results from the reading. The partial derivatives that form the entries of the softmax Jacobian are:

Notice how the derivative of a softmax output always depends on the outputs themselves, not the inputs. This self-referential nature is common in functions involving exponentials.

• When $i=j$, the derivative $\frac{\partial s_i}{\partial z_i}$ is positive. Increasing a logit $z_i$ will always increase its corresponding probability $s_i$.
• When $i\ne j$, the derivative $\frac{\partial s_i}{\partial z_j}$ is negative. Increasing a logit $z_j$ will increase $s_j$, and to keep the total probability sum at 1, all other probabilities $s_i$ must decrease.

2. A Compact Matrix Representation

While the piecewise formula is correct, expressing the Jacobian in a single matrix equation can provide deeper insight. Given your background, you'll likely appreciate this more compact and algebraic view. We can unify the two cases using the Kronecker delta, ${\delta }_{ij}$ (which is 1 if $i=j$ and 0 otherwise):

You can verify that this correctly collapses to $s_i(1-s_i)$ when $i=j$ and to $-s_i{s}_j$ when $i\ne j$.

This component-wise definition leads to a very elegant matrix representation. The entire Jacobian matrix $J_S$ can be written as:

Here:

• $\mathbf{s}$ is the column vector of softmax outputs.
• $\text{diag}(\mathbf{s})$ is a diagonal matrix with the elements of $\mathbf{s}$ on its main diagonal. This matrix captures all the $s_i(1-s_i)$ terms, but only on the diagonal, since $\text{diag}(\mathbf{s}{)}_{ij}=s_i{\delta }_{ij}$. To see this, expand $s_i({\delta }_{ij}-s_j)=s_i{\delta }_{ij}-s_i{s}_j$. The first term $s_i{\delta }_{ij}$ corresponds to $\text{diag}(\mathbf{s})$.
• $\mathbf{s}{\mathbf{s}}^T$ is the outer product of the softmax output vector with itself. This produces a $K\times K$ matrix where the entry at $(i,j)$ is $s_i{s}_j$. This captures the second term, $-s_i{s}_j$.

This form, $\text{diag}(\mathbf{s})-\mathbf{s}{\mathbf{s}}^T$, is not only concise but also reveals that the softmax Jacobian is a symmetric matrix, since both $\text{diag}(\mathbf{s})$ and $\mathbf{s}{\mathbf{s}}^T$ are symmetric. The following resource provides an excellent derivation that arrives at this matrix form directly.

3. Re-deriving the Combined Gradient

Now we have all the machinery needed to formally prove the simplification we used in the last lesson. Our goal is to compute ${\nabla }_{\mathbf{z}}L$ for the cross-entropy loss $L=-{\sum }_k{y}_k\log (s_k)$.

By the chain rule (using numerator layout, where gradients are column vectors):

Since $J_S$ is symmetric, $J_S^T=J_S$, so we have:

Let's compute the two parts and then their product.

1. Gradient of Loss w.r.t. Softmax Outputs (${\nabla }_{\mathbf{s}}L$)
The partial derivative of the loss with respect to a single softmax output $s_i$ is:

The gradient vector ${\nabla }_{\mathbf{s}}L$ is therefore the column vector with entries $-y_i/s_i$.

2. The Matrix-Vector Product
Now we multiply the Jacobian by this gradient vector:

Let's evaluate each term:

• First term: $\text{diag}(\mathbf{s})({\nabla }_{\mathbf{s}}L)$. This is a diagonal matrix multiplying a vector. The $i$-th element of the resulting vector is the product of the $i$-th diagonal element of the matrix ($s_i$) and the $i$-th element of the vector ($-y_i/s_i$).

  $$[\text{diag}(\mathbf{s})({\nabla }_{\mathbf{s}}L){]}_i=s_i\cdot \left(-\frac{y_i}{s_i}\right)=-y_i$$

  So, the first term simplifies to the vector $-\mathbf{y}$.

• Second term: $-(\mathbf{s}{\mathbf{s}}^T)({\nabla }_{\mathbf{s}}L)$. We can re-group this using associativity of matrix multiplication: $\mathbf{s}({\mathbf{s}}^T({\nabla }_{\mathbf{s}}L))$. The term in the parentheses is a scalar dot product:

  $${\mathbf{s}}^T({\nabla }_{\mathbf{s}}L)=\sum _{k=1}^K{s}_k\left(-\frac{y_k}{s_k}\right)=-\sum _{k=1}^K{y}_k$$

  Since $\mathbf{y}$ is a one-hot vector representing the true class, the sum of its elements $\sum y_k$ is exactly 1. So the dot product is -1.
  The second term thus becomes $-\mathbf{s}(-1)=\mathbf{s}$.

Combining the two terms, we get our final result:

Or using $\hat{\mathbf{y}}$ for the predicted probabilities $\mathbf{s}$, we have ${\nabla }_{\mathbf{z}}L=\hat{\mathbf{y}}-\mathbf{y}$. This completes the formal derivation. The intricate dependencies within the softmax Jacobian perfectly cancel with the logarithmic derivatives from the cross-entropy loss, leaving a simple, intuitive error signal.

For a final reinforcement of these steps, the video below walks through the same derivation by combining the derivative terms element by element.

Conclusion

In this lesson, we moved from using a result to deriving it from scratch. We have now rigorously established the mathematical underpinnings of the gradient for the most common classification output layer.

Key Takeaways:

• The Jacobian of the softmax function has a specific structure: positive on the diagonal ($s_i(1-s_i)$) and negative off the diagonal ($-s_i{s}_j$).
• This structure can be expressed compactly in matrix form as $J_S=\text{diag}(\mathbf{s})-\mathbf{s}{\mathbf{s}}^T$, revealing its symmetry.
• The application of the chain rule to the composition of softmax and cross-entropy loss results in a cascade of cancellations, yielding the simple and computationally efficient gradient $\mathbf{s}-\mathbf{y}$. This is a beautiful example of how mathematical components are designed to fit together elegantly in deep learning.

In our upcoming lessons, we will continue to build our matrix calculus toolkit. Our next topic will be deriving the derivatives for several critical matrix operations, such as those involving determinants and quadratic forms. These identities are not just theoretical curiosities; they are the essential tools needed to derive the objective functions for advanced generative models like Variational Autoencoders and Normalizing Flows, bringing us another step closer to implementing those papers.