Softmax Gradient (Vector-Jacobian Product)

Difficulty: Hard · Concept: Activation Functions — Auxiliary Functions

Sigmoid Derivative was elementwise: each output only depends on its own input, so $\sigma'(z_i)$ is a single number per element. Softmax is fundamentally different — every output $s_i = \text{softmax}(z)_i$ depends on all of $z$ (through the shared normalizing sum), so its derivative isn't a list of numbers, it's a full $n \times n$ Jacobian matrix: $\frac{\partial s_i}{\partial z_j} = s_i(\delta_{ij} - s_j)$, where $\delta_{ij}$ is $1$ if $i=j$ else $0$.

Explicitly building that $n \times n$ matrix during backprop would be wasteful — real deep learning frameworks never materialize it. Instead they compute the vector-Jacobian product (VJP) directly: given the upstream gradient $\frac{\partial L}{\partial s}$ (how loss changes with each softmax output), get $\frac{\partial L}{\partial z}$ (how loss changes with each softmax input) in one $O(n)$ pass, using an identity that falls out of the Jacobian formula above:

$\frac{\partial L}{\partial z_j} = s_j\left(\frac{\partial L}{\partial s_j} - \sum_i \frac{\partial L}{\partial s_i} s_i\right)$

Your task: implement softmax_gradient(s, grad_output), where s is a softmax output (already computed, sums to 1) and grad_output is $\partial L/\partial s$. Return $\partial L/\partial z$, the same length as s.

Next: Sample Mean and Variance From Scratch