Implementing Stochastic Gradient Descent with both Weight Decay and Momentum

Question

So I'm trying to implement a neural network using only numpy module in Python. The problem I'm facing is related to the correct implementation of the regularization through weight decay, and also the momentum in stochastic gradient descent.

When I only perform sgd without any weight decay and momentum my neural net performs well on the application scenario I'm facing. If I try to add either of these two terms, my model performs really badly. So I suppose that I'm not understanding them well.

Since my model works fine with sgd, let us assume that my model do it's job in forward and back propagation. The problematic aspect arise in the implementation of the parameters update for a layer of my net.

The method for the parameters update is the following:

def update_params(self, n_samples):
        weight_decay_term = (self.learning_rate * self.weight_decay / n_samples) * self.W
        weight_velocity = self.momentum * self.weight_velocity + self.dW
        self.W = self.W - weight_decay_term - self.learning_rate * weight_velocity
        self.weight_velocity = np.copy(weight_velocity)
        bias_velocity = self.momentum * self.bias_velocity + self.db
        self.b = self.b - self.learning_rate * bias_velocity
        self.bias_velocity = np.copy(bias_velocity)

where 'n_samples' is the size of the mini-batch.

For the weight decay formula I followed the guideline described in: https://github.com/limberc/deeplearning.ai/blob/master/COURSE%202%20Improving%20Deep%20Neural%20Networks%20Hyperparameter%20tuning%2C%20Regularization%20and%20Optimization/week%2001/week%2001%20Setting%20up%20your%20Machine%20Learning%20Application.pptx

For the momentum formula I followed the guideline described in: https://cs230.stanford.edu/files/C2M2.pdf

I also tried the formula described in: Neural Networks: weight change momentum and weight decay without any success.

None of these solutions worked, meaning that setting for example:

self.learning_rate = 0.01
self.momentum = 0.9
self.weight_decay = 0.1

my model performs really badly. I suppose it is related to my understanding of the implementation details of weight decay and momentum, but I really can't wrap my head around this problem. If I set the momentum and weight_decay parameters to zero I get acceptable results with my model, which are similar to those that I obtain with a pytorch model with the same architecture.

I tried to instantiate a pytorch multy layer perceptron with the same architecture that I tried with my model, and used as optimizer:

torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.1)

and the torch net performs greatly on my application scenario.

So, can you spot the problem on my implementation of the update parameters method? Thank you for reading

Answer 1

In short, the n_samples in OP's weight_decay_term equation should be the size of the whole training set instead of mini batch. It was the same error I made.

Below are the explanation of my mistakes and corrections:

I had the same problem, restudy the equations, and found that in my coding, I mixed up the n, size of the whole training set, with the m, size of the mini batch.

Based on my understanding, for the weight decay term calculation, n, size of the whole training set should be used.

Whereas, m, size of the mini batch is used in the approximation of gradient of cost function w.r.t weights or biases using mini batch, which are $\frac{1}{m}\sum_x\frac{\partial C_x}{\partial w}$ and $\frac{1}{m}\sum_x\frac{\partial C_x}{\partial b}$, and I assumed that you have covered in your self.dw and self.db calculations as it's not shown in your codes.

I am following Michael Nielsen's convention so the terms might be different from yours but I think they are the same concept. With some algebra, my equation derivations are as follows:

For weights: $$ v\rightarrow v'=\mu v-\eta\nabla C \\\therefore v\rightarrow v'=\mu v-\frac{\eta}{m}\sum_x\frac{\partial C_x}{\partial w}-\frac{\eta\lambda}{n}w \\w\rightarrow w'=w+v' $$

For biases: $$ v\rightarrow v'=\mu v-\eta\nabla C \\\therefore v\rightarrow v'=\mu v-\frac{\eta}{m}\sum_x\frac{\partial C_x}{\partial b} \\b\rightarrow b'=b+v' $$

where,

$\nabla C$ = gradient of cost function w.r.t weights or biases,

$\mu$ = momentum coefficient, $\eta$ = learning rate, $\lambda$ = regularization parameter in weight decay term,

$m$ = size of mini batch, $n$ = size of the whole training set

After the correction, with $\mu=0.5$, $\eta=0.1$, $\lambda=5$, $m=10$, the model managed to get around 97% of final evaluation accuracy with 100 epochs in the MNIST problem. $n$ will be size of the MNIST training set, which is $50000$.

Please let me know if I made mistakes in the explanation above, coz I am new to neural networks.