Weird behaviour of softmax derivative?
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.