Using relu makes train accuracy to drop in the middle of training and MSE to fluctuate

Question

I am building a neural network using numpy, in python. it has one hidden layer with 16 neurons. I train it on MNIST. the learning rate is 0.1, the activation function for the hidden layer re RELU and for the output layer is softmax. the cost is MSE. the values are normelized between 0 to 1. weights and biases initializtions:

bound = (6 ** 0.5) / ((input_size + self.hidden_size)**0.5)
self.w1_2 = np.random.uniform(-bound,bound,(input_size, self.hidden_size))
bound = (6 ** 0.5) / ((self.hidden_size + output_size)**0.5)
self.w2_3 = np.random.uniform(-bound,bound,(self.hidden_size, output_size))
self.b2 = np.zeros(self.hidden_size)
self.b3 = np.zeros(output_size)

the mse and train accuracy change with the training as follows:

you can see that after 400 epochs things start to mess up. when I use sigmoid instead of relu in the hidden layer, things works out fine and there is no problem.

can anyone think of possible reasons for this problem?

Answer 1

I guess the learning rate is high since the accuracy increases slowly (and loss decreases slowly) in the initial epochs and the loss drastically fluctuates (with dip in accuracy) long into training. Could you try reducing the learning rate by a factor of 3 or 10 to 0.03, 0.01 or lower and train the model again?

Note: For finding a good learning rate hyperparameter, you could only check the model's performance for initial few epochs (say, only first 50 epochs and need not train for 400 epochs to compare learning rates) and it would be intuitive if the direction of change in learning rate is helping to train the model.

PS: ReLU almost never harms the performance of such a network, so I don't think the problem is with ReLU. Correct me if I'm wrong.