Softmax layer in a neural network

Question

I'm trying to add a softmax layer to a neural network trained with backpropagation, so I'm trying to compute its gradient.

The softmax output is $h_j = \frac{e^{z_j}}{\sum{e^{z_i}}}$ where $j$ is the output neuron number.

If I derive it then I get

$\frac{\partial{h_j}}{\partial{z_j}}=h_j(1-h_j)$

Similar to logistic regression. However this is wrong since my numerical gradient check fails.

What am I doing wrong? I had a thought that I need to compute the cross derivatives as well (i.e. $\frac{\partial{h_j}}{\partial{z_k}}$) but I'm not sure how to do this and keep the dimension of the gradient the same so it will fit for the back propagation process.

Answer 1

I feel a little bit bad about providing my own answer for this because it is pretty well captured by amoeba and juampa, except for maybe the final intuition about how the Jacobian can be reduced back to a vector.

You correctly derived the gradient of the diagonal of the Jacobian matrix, which is to say that

$ {\partial h_i \over \partial z_j}= h_i(1-h_j)\;\;\;\;\;\;: i = j $

and as amoeba stated it, you also have to derive the off diagonal entries of the Jacobian, which yield

$ {\partial h_i \over \partial z_j}= -h_ih_j\;\;\;\;\;\;: i \ne j $

These two concepts definitions can be conveniently combined using a construct called the Kronecker Delta, so the definition of the gradient becomes

$ {\partial h_i \over \partial z_j}= h_i(\delta_{ij}-h_j) $

So the Jacobian is a square matrix $ \left[J \right]_{ij}=h_i(\delta_{ij}-h_j) $

All of the information up to this point is already covered by amoeba and juampa. The problem is of course, that we need to get the input errors from the output errors that are already computed. Since the gradient of the output error $\nabla h_i$ depends on all of the inputs, then the gradient of the input $x_i$ is

$[\nabla x]_k = \sum\limits_{i=1} \nabla h_{i,k} $

Given the Jacobian matrix defined above, this is implemented trivially as the product of the matrix and the output error vector:

$ \vec{\sigma_l} = J\vec{\sigma_{l+1}} $

If the softmax layer is your output layer, then combining it with the cross-entropy cost model simplifies the computation to simply

$ \vec{\sigma_l} = \vec{h}-\vec{t} $

where $\vec{t}$ is the vector of labels, and $\vec{h}$ is the output from the softmax function. Not only is the simplified form convenient, it is also extremely useful from a numerical stability standpoint.

Answer 2

The derivative is wrong. It should be,

$$\frac{\partial h_{j}}{\partial z_{k}} = h_{j}\delta_{kj}-h_{j}h_{k}$$

check your calculations again. Also, the expression given by amoeba for the cross-entropy is not entirely correct. For a set of data samples drawn from $C$ different classes, it reads,

$$-\sum_{n}\sum_{k=1}^{C}t_{k}^{n}\ln y_{k}(\boldsymbol{x}^{n})$$

where the superindex runs over the sample set, $t_{k}^{n}$ is the value of the k-th component of the target for the n-th sample. Here it is assumed that you are using a 1-of-C coding scheme, that is, $t_{k}^{n}$. In such case all t's are zero except for the component representing its corresponding class, which is one.

Note, that the t's are constant. Hence minimizing this functional is equivalent to minimizing,

$$-\sum_{n}\sum_{k=1}^{C}t_{k}^{n}\ln y_{k}(\boldsymbol{x}^{n}) + \sum_{n}\sum_{k=1}^{C}t_{k}^{n}\ln t_{k}^{n} = -\sum_{n}\sum_{k=1}^{C}t_{k}^{n}\ln \frac{y_{k}(\boldsymbol{x}^{n})}{t_{k}^{n}}$$

which has the advantage that the Jacobian takes a very convenient form, namely,

$$\frac{\partial E}{\partial z_{j}} = h_{j}-t_{j}$$

I would recommend you to get a copy of Bishop's Neural Networks for Pattern Recognition. IMHO still the best book on neural networks.

Answer 3

Each output of the softmax depends on all the inputs, so the gradient is indeed a whole Jacobian matrix. You correctly computed $\partial_j h_j = \frac{\partial h_j}{\partial z_j}=h_j(1-h_j)$, but you also need $\partial_k h_j=-h_jh_k$ if $j \neq k$. I guess if you can derive the first expression, you should easily be able to derive the second one as well.

I am not sure what problem you see with back-propagating: in the softmax layer you have $j$ outputs and $j$ inputs, so an error from each output should be propagated to each input, and that is precisely why you need the whole Jacobian. On the other hand, usually you would have a cost function associated with the softmax output, e.g. $$C=-\sum_j t_j \log h_j, $$ where $t_j$ are your desired outputs (when you do classification, then often one of them is equal to 1, and others to 0). Then in fact you are interested in $\frac{\partial C}{\partial z_j}$, which can be computed with a chain rule resulting in a neat expression, and is indeed a vector (not a matrix).