This thing called Weight Decay. Learn how to use weight decay to train... | by Dipam Vasani | Towards Data Science

In my previous article, I mentioned that data augmentation helps deep learning models generalize well. That was on the data side of things. What about the model side of things? What can we do while training our models, that will help them generalize even better.
We do weight decay.
We start by looking at the image above. We see that we have a bunch of data points and that we cannot fit them well with a straight line. Hence, we use a 2nd degree polynomial to do so. We also notice that if increase the degree of the polynomial beyond a certain point, then our model becomes too complex and starts to overfit.
This means that in order to prevent overfitting, we shouldn’t allow our models to get too complex. Unfortunately, this has led to a misconception in deep learning that we shouldn’t use a lot of parameters (in order to keep our models from getting overly complex).
First of all, real world data is not going to be as simple as the one shown above. Real world data is complex and in order to solve complex problems, we need complex solutions.
Having fewer parameters is only one way of preventing our model from getting overly complex. But it is actually a very limiting strategy. More parameters mean more interactions between various parts of our neural network. And more interactions mean more non-linearities. These non-linearities help us solve complex problems.
However, we don’t want these interactions to get out of hand. Hence, what if we penalize complexity. We will still use a lot of parameters, but we will prevent our model from getting too complex. This is how the idea of weight decay came up.
We’ve seen weight decay in my article on collaborative filtering. In fact, every learner in the fastai library has a parameter called weight decay.
One way to penalize complexity, would be to add all our parameters (weights) to our loss function. Well, that won’t quite work because some parameters are positive and some are negative. So what if we add the squares of all the parameters to our loss function. We can do that, however it might result in our loss getting so huge that the best model would be to set all the parameters to 0.
To prevent that from happening, we multiply the sum of squares with another smaller number. This number is called weight decay or wd.
Our loss function now looks as follows:
Loss = MSE(y_hat, y) + wd * sum(w^2)
When we update weights using gradient descent we do the following:
w(t) = w(t-1) - lr * dLoss / dw
Now since our loss function has 2 terms in it, the derivative of the 2nd term w.r.t w would be:
d(wd * w^2) / dw = 2 * wd * w (similar to d(x^2)/dx = 2x)
That is from now on, we would not only subtract the learning rate * gradient from the weights but also 2 * wd * w . We are subtracting a constant times the weight from the original weight. This is why it is called weight decay.
Generally a wd = 0.1 works pretty well. However, the folks at fastai have been a little conservative in this respect. Hence the default value of weight decay in fastai is actually 0.01 .
The reason to choose this value is because if you have too much weight decay, then no matter how much you train, the model never quite fits well enough whereas if you have too little weight decay, you can still train well, you just have to stop a little bit early.
I’ve demonstrated this concept in this jupyter notebook.
It is a multi-class (and not a multi-label) classification problem where we try to predict the class of plant seedlings.
I’ve used 3 values for weight decay, the default 0.01 , the best value of 0.1 and a large value of 10 . In the first case our model takes more epochs to fit. In the second case it works best and in the final case it never quite fits well even after 10 epochs. (see difference b/w training and validation loss.)
That will be it for this article. You can learn about other regularization techniques here.
If you want to learn more about deep learning check out my series of articles on the same: