I've seen plenty of examples of people doing Sigmoid + MSE backpropagation implementations, yet I do not seem to understand how to implement backpropagation as stated in the title in the case of multi-class classifications.
What confuses me mainly are the matrix-vector shapes and their multiplications, and their implementations in code.
Alright. Consider an ordinary neural network, so, in the last layer, we have, $z^{[L]} = W^{[L]} a^{[L-1]} + b^{[L]}$, where $a^{[L]} = \sigma(z^{[L]})$, where $\sigma$ is the softmax activation: $$ \sigma(\mathbf z)_{i} = \frac{e^{z_i}}{\sum_k e^{z_k}} $$
I think, one of the most effective ways of not to get confused about all these matrices with different dimensions, is to simply get rid of the matrices and do all calculations componentwise. So, calculating the $ij$ component of the jacobian matrix, we have: $$ J_{ij} = \frac{\partial\sigma_i}{\partial z_j} = \frac{e^{z_i}}{\sum_k e^{z_k}}\delta_{ij} - \frac{e^{z_i}e^{z_j}}{(\sum_k e^{z_k})^2} $$
Consider a cost $C$, and we wish to calculate the derivative with respect to the weights of the last layer, that is, we are running the first step of backpropagation. Applying chain rule to the $ij$-component of the weight:
$$ \frac{\partial C}{\partial W^{[L]}_{ij}} = \sum_k\frac{\partial C}{\partial a_{k}}\frac{\partial a^{[L]}_{k}}{\partial W^{[L]}_{ij}} $$
And, with one more chain rule: $$ \frac{\partial a^{[L]}_{k}}{\partial W^{[L]}_{ij}} = \sum_s\frac{\partial a^{[L]}_{k}}{\partial z^{[L]}_s}\frac{\partial z^{[L]}_s}{\partial W^{[L]}_{ij}} $$
Now we can identify exactly where the jacobian matrix is: $$ J_{ij} = \frac{\partial\sigma_i}{\partial z_j} = \frac{\partial a^{[L]}_{i}}{\partial z^{[L]}_j} $$
Thus, the back-propagation step is just: $$ \frac{\partial C}{\partial W^{[L]}_{ij}} = \sum_k\sum_s\frac{\partial C}{\partial a_{k}}\frac{\partial a^{[L]}_{k}}{\partial z^{[L]}_s}\frac{\partial z^{[L]}_s}{\partial W^{[L]}_{ij}} = \sum_k\sum_s\frac{\partial C}{\partial a_{k}}J_{ks}\frac{\partial z^{[L]}_s}{\partial W^{[L]}_{ij}} $$
Now, in component language, it is much simpler to identify where are all the matrix multiplications: The transpose of the gradient with respect to the cost function, multiplied by the jacobian matrix of the softmax activation (or whatever activation), and finally, the last derivative, which will evaluate to something depending on the activation of the previous layer.
=].
