Loss randomly changing, incorrect output (even for low loss) when trying to overfit on a single set of input and output

Question

I am trying to make a neural network framework from scratch in C++ just for fun, and to test my backpropagation, I thought it would be an easy way to test the functionality if I give it one input - a randomized size 10 vector, and one output: a size 5 vector containing all 1s, and train it a bunch of times to see if the loss will decrease. Essentially trying to make it overfit

The problem is that for each run that I do, the loss either shoots up and goes to nan, or reduces a lot, going to 0.000452084 or other similar small values. However even in the low end of things, my output (which should be close to all 1s, as the "ground truth") is something like:

0.000263654
1e-07
8.55893e-05
1e-07
0.999651

The only close value close to 1 being the last element.

My network consists of the input layer 10 neurons, one 10 neuron dense layer with RELU activation, and another 5 neuron dense layer for output, with SoftMax activation. I am using categorical cross entropy as my loss function, and I am normalizing my gradient by dividing it by the norm of my gradient if it is over 1.0. I initialize my weights to be random values between -0.1 and 0.1

To calculate the gradient of the loss function, I use -groundTruth/predictedOutput. To calculate the other derivatives, I dot the derivative of that layer with the gradient of the previous layer with respects to its activation function.

Before this problem I was having exploding gradients, which the gradient scaling fixed, however it was very weird that that would even happen on a very small network like this, which could be related to the problem I am currently having. Is the implementation not correct or am I missing something very obvious?

Any ideas about this weird behavior, and where I should look first? I am not sure how to show a minimal reproduceable example as that would require me to paste the whole codebase, but I am happy to show pieces of code with explanation. Any advice welcomed!!

Answer 1

Softmax activation always adds up to 1, because it's designed to deal with probabilities (in problems of classification, those probabilities represent how likely the network thinks an object belongs to a specific class). You can verify that by summing up the numbers of your output layer. So currently your network is trying to do the impossible, to produce output that sums up to 5, instead of 1. Therefore, loss will never become stable. Since you want your output layer to produce all ones, you need to use some other activation, for example, linear. Linear activation does not have the same constraint that Softmax does.