Derivation of backpropagation for Softmax

Question

So, after a couple dozen tries I finally implemented a standalone nice and flashy softmax layer for my neural network in numpy. All works well, but I have a question regarding the maths part because there's just one tiny point I can't understand, like at all.

Having any kind of activation function in the output layer, back-propagation looks like:

$$ z = x\cdot w+b \\ a = f(z) \\ E = \frac{1}{N} \sum^{N}_{i=1} (a_{i} - y_{i})^{2} \\ \frac{\partial E}{\partial w} := \frac{\partial E}{\partial a_{i}} \cdot\frac{\partial a_{i}}{\partial z} \cdot \frac{\partial z}{\partial w} \\ \frac{\partial E}{\partial w} := (a_{i} - y_{i}) \cdot f^{'}(z) \cdot x $$

So by differentiating $ a_{l} $ with respect to $ z_{l} $, the result is the derivative of the activation function with $ z_{l} $ itself. Now, with Softmax in the final layer, this does not apply. If I use $ Softmax'(z_{l}) $ I get incorrect results, but I rather need $ Softmax'(a_{l}) $. Looking at a couple online materials like LINK in the Backpropagation phase section in the python code, the author also uses the activation itself in the argument for softmax's derivative and I've been differentiating there and back for the last 2 weeks, ... I don't understand why it works like that.

In short, my question is: With specifically Softmax in the output layer, why is my $ \frac{\partial a_{l}}{\partial z_{l}} $ becomes $ Softmax'(a_{l}) $ instead of $ Softmax'(z_{l}) $ like with pretty much any other normal activation function that's out there?

Answer 1

We let $$a=\operatorname{Softmax}(z)$$

that is $$a_i = \frac{e^{z_i}}{\sum_{j=1}^N e^{z_j}}.$$

$a$ is indeed a function of $z$ and we want to differentiate $a$ with respect to $z$. The interesting thing is we are able to express this final outcome as an expression of $a$ in an elegant fashion.

If you look at the section of "Derivative of Softmax Function" in your link, using quotient rule:

We can see that

if $i=m$, \begin{align}\frac{\partial a_i}{\partial z_m} &=\left(\frac{e^{z_i}}{\sum_{j=1}^N e^{z_j}}\right)\left(\frac{\sum_{j=1}^Ne^{z^j}-e^{z_m}}{\sum_{j=1}^N e^{z_j}} \right)=a_i(1-a_m)\end{align}

if $i\ne m$, \begin{align}\frac{\partial a_i}{\partial z_m} &=-\left(\frac{e^{z_m}}{\sum_{j=1}^N e^{z_j}}\right)\left(\frac{e^{z_i}}{\sum_{j=1}^N e^{z_j}} \right)=-a_m(a_i)\end{align}

If you want to evaluate things in terms of $z$, you can still use the middle term though using the formula

$$\frac{da}{dz}=ae^T\circ (I-ea^T)$$

is more elegant.