Weird behaviour of softmax derivative?

Question

I have been implementing some neural networks in MATLAB and recently I noticed a weird thing while implementing softmax derivative:

Setting the derivative to one, rather than using the actual derivative converges faster.

I have to admit that the derivative of softmax in particular confused me quite a bit, since the actual derivative requires the Jacobian as opposed to other activation functions that only depend on the input. That is, for output vector of size nx1, the derivative of the activation function is also nx1 (see ReLU, tanh etc.) but for softmax is nxn.

Anyway, the simple case is to consider the diagonal of the Jacobian which is “out*(1-out)” as derived and proposed in a couple of sources around [1-3]. However setting the derivative to “1” shows faster convergence.

I am not fully sure to say how this phenomenon generalises but I tried with different settings, architectures and optimisers and the issue seems to persist.

I made an example of a 2x5x5x2 feedforward network with tanh in all but last layer, where I’ve used softmax, classifying a typical XOR problem using cross-entropy loss with ADAM optimizer.

Does anyone have any hunch why this would happen?

Refs:

1: https://eli.thegreenplace.net/2016/the-softmax-function-and-its-derivative/
2: https://math.stackexchange.com/questions/945871/derivative-of-softmax-loss-function
3: https://datascience.stackexchange.com/questions/29735/how-to-apply-the-gradient-of-softmax-in-backprop

P.S. I am aware that showing this behaviour in only one problem does not mean that it happens always, but in a couple of experiments I have done, this behaviour seems to be persistent

Answer 1

I'm not sure what you mean by "the simple case is to consider the diagonal of the Jacobian". Are you perhaps doing some unusual variation of gradient descent where you compute only the diagonal of the Jacobian, but not the other entries? If you are, then don't be surprised if it performs poorly. Don't do that. I wouldn't expect that to work well.

Instead, I recommend you follow the standard approach: use backpropagation to compute the gradient and use gradient descent to update the weights. There are many standard tutorials on how to do that. Realistically, we never need to compute the Jacobian matrix for the softmax directly. Instead, we look at the output of the loss function, and we compute the gradient of that with respect to the weights. That doesn't require a Jacobian, because we are computing the gradient of a function that has a single output. Then, we use the gradient to update the weights (that's the gradient descent part).

If you're implementing this yourself from scratch, there are many things that can go wrong. I recommend implementing a gradient check to make sure you are computing the gradient properly. But, more realistically, I recommend using some existing framework for training neural networks (e.g., Keras, Torch, etc.) rather than trying to implement it from scratch. If you implement from scratch, there are many ways you can create some small implementation error that will lead to results that are subpar but hard to debug.