Softmax Gradients for ANN

It is common to see Artificial Neural Networks with a sofmax function in the last layer. This activation function outputs a probability distribution among the different classes that the networks is trained to predict on. Softmax, for the case of a single class looks like this: $$ \pi_{\theta}(a | s) = \frac{e^{z_{a}}}{\sum_{k}e^{z_{k}}} $$ In order to train a NN with softmax as its last activation function, the cost function known as the cross-entropy error function has to be minimized.

Cross-entropy derivatives

Let $ \mathbf{y} $ denote the output of a neural network with a softmax function as the activation of its last layer and $ \mathbf{t} $ a 'one-hot encoded' vector. Then, the error function is: $$ E = -\sum_{i} t_{i} \log{ y_{i} } $$ By estimating the gradient of the error function with respect to the parameters of a neural network, and updating these parameters with these gradients, the network can learn.

In order to estimate the gradient of the error function with respect to the parameters of the last layer of a network, $ L $ will denote the number of layers, $ \mathbf{h}^{(L)} $ the output of the previous layer and $ \mathbf{z}^{(L)} $ the result of the multiplication of $ \mathbf{h}^{(L)} $ by the weight matrix $ \mathbf{w}^{(L)} $. $$ y_{i} = \frac{e^{z_{i}^{(L)}}}{\sum_{k} e^{z_{k}^{(L)}}} \\ \\ \\ \mathbf{z}^{(L)} = \mathbf{w}^{(L)} \mathbf{h}^{(L)} \quad \small\text{where} \quad \mathbf{w}^{(L)} \in \mathbb{R}^{i\times j} $$
Therefore, for any given $ w_{ij}^{(L)} \in \mathbf{ w }^{(L)} $,
$$ \begin{align*} \frac{\partial{E}}{\partial{w_{ij}^{(L)}}} &= \sum_{i} \frac{\partial{E}}{\partial{y_{i}}} \Bigg( \sum_{k} \frac{\partial{y_{i}}}{\partial{z_{k}^{(L)}}} \frac{\partial{z_{k}^{(L)}}}{\partial{w_{ij}^{(L)}}} \Bigg) \\ &= \sum_{i} \frac{ \partial{E} }{ \partial{y_{i}} } \frac{ \partial{y_{i}} }{ \partial{z_{j}^{(L)}} } \frac{ \partial{z_{j}^{(L)}} }{ \partial{w_{ij}^{(L)}} } \quad \tiny\texttt{True since $ \frac{\partial{z_{k}^{(L)}}}{\partial{w_{ij}^{(L)}}}=0 \ \ \forall \ \ k \neq j$.} \\ &= h_{i}^{(L)} (y_{j}\sum_{i}t_{i} - t_{j}) \\ &= h_{i}^{(L)}(y_{j} - t_{j}) \quad \tiny\texttt{True since $\sum_{i}t_{i} = 1$, $\mathbf{t}$ is a 'one-hot' encoded vector.} \end{align*} $$
In vectorized notation, $$ \begin{align*} \nabla_{\mathbf{w}^{(L)}}E &= \mathbf{h}^{\intercal(L)}(\mathbf{y}-\mathbf{t}) \\ \\ &= \begin{bmatrix} h_{1}^{(L)}(y_{1} - t_{1}) & h_{1}^{(L)}(y_{2} - t_{2}) & \dots \\ \vdots & \ddots & \\ h_{i}^{(L)}(y_{1} - t_{1}) & \dots & h_{i}^{(L)}(y_{j} - t_{j}) \end{bmatrix} \end{align*} $$
And the derivative with respect to a previous weight matrix can be obtained as follows. $$ \begin{align*} \frac{\partial{E}}{\partial{w_{pq}^{(L-1)}}} &= \sum_{i} \frac{\partial{E}}{\partial{y_{i}}} \Bigg( \sum_{k} \frac{\partial{y_{i}}}{\partial{z_{k}^{(L)}}} \frac{\partial{z_{k}^{(L)}}}{\partial{h_{q}^{(L)}}} \frac{\partial{h_{q}^{(L)}}}{\partial{z_{q}^{(L-1)}}} \frac{\partial{z_{q}^{(L-1)}}}{\partial{w_{pq}^{(L-1)}}}\Bigg) \\ \\&= x_{p}h_{q}^{(L)}(1-h_{q}^{(L)})\sum_{q}w_{pq}^{(L)}(y_{q}-t_{q}) \end{align*} $$
Expressed as a a matrix, where $ h_{q}^{\prime(L)} = h_{q}^{(L)}(1 - h_{q}^{(L)}) $ and $ \delta_{q} = y_{q} - t_{q} $,
$$ \begin{align*} \nabla_{\mathbf{w}^{(L-1)}}E &= \mathbf{x}^{\intercal} \Big( \mathbf{h}^{(L)} \circ (1 - \mathbf{h}^{(L)}) \circ \big[ (\mathbf{y} - \mathbf{t})\mathbf{w}^{\intercal(L)} \big] \Big) \\ \\ &= \begin{bmatrix} x_{1}h_{1}^{\prime(L)} \sum_{q}w_{1q}^{(L)}\delta_{q} & x_{1}h_{2}^{\prime(L)}\sum_{q}w_{1q}^{(L)}\delta_{q} & \dots \\ \vdots & \ddots & \\ x_{p}h_{1}^{\prime(L)}\sum_{q}w_{pq}^{(L)}\delta_{q} & \dots & x_{p}h_{q}^{\prime(L)}\sum_{q}w_{pq}^{(L)}\delta_{q} \end{bmatrix} \end{align*} $$