so I am writing my own neural network library for a class project and I got everything working for a simple 2-class test using the distance (L2) cost function. I wanted to get a similar result using softmax and crossentropy instead.
I did the derivation of the formulas as in this article, but I tried computing the derivative of the cost function w.r.t. its inputs and the Jacobian of the softmax outputs w.r.t. the inputs separately.
This is what I did:
So, let $Y$ be the matrix with rows equal to the expected class probabilities, $\hat{Y}$ the outputs of the softmax layer, $X$ the inputs of the softmax, all of these matrices of size $n\times c$ (samples times classes). The cost function $J(Y,\hat{Y}) = -sum(Y * \log(\hat{Y}))$ where * is element by element multiplication.
Then, if $E$ is the cost function, I first compute
$$\frac{\partial E}{\partial \hat{Y}_j^{(i)}} = -\frac{Y_j^{(i)}}{\hat{Y}_j^{(i)}},\:\: i=1,\cdots,n : j=1,\cdots,c$$
This is a $n\times c$ matrix. Then, being $S$ the softMax function I compute the Jacobian of the softmax at each row of $X$, so a total of $n$ matrices of $c\times c$ dimensions.
$$JS^{(i)} = \left(\frac{\partial \hat{Y}_j^{(i)}}{\partial X_k^{(i)}} \right)_{j,k=1,\cdots,c},\:\: i=1,\cdots,n$$
Finally to compute the partial derivative of the error with respect to the inputs of the softmax I multiply the rows of $\frac{\partial E}{\partial \hat{Y}}$ with each jacobian, to get a new $n\times c$ matrix.
This way I have everything needed to compute:
$$\frac{\partial E}{\partial X_j^{(i)}}= \sum_{k=1}^c \frac{\partial E}{\partial \hat{Y}_k^{(i)}}\frac{\partial \hat{Y}_k^{(i)}}{\partial X_j^{(i)}} ,\:\: i=1,\cdots,n : j=1,\cdots,c$$
Is this correct? Where is the problem if not so? Might the numerical errors make everything fail?
Thanks in advance.