weight decay multiplied by learning rate
Opened this issue · 8 comments
If you look carefully at the formula in the article, the weight decay (w) is not multiplied by the learning rate (alpha), but rather by the schedule coefficient (eta). In the pytorch implementation eta seems to be always assumed to be one. In your code, however, you implicitly multiply w by alpha (step_size). I suggest changing the corresponding line to
p.data.add_(torch.mul(p.data, group['weight_decay']).addcdiv(-step_size, exp_avg, denom) )
Thank you for your advice. I would check it.
Note that if the implementation is changed so that the weight decay is not multiplied by the learning rate you should subtract the weight decay from the weights instead of adding them (otherwise the weights will become larger, not smaller). Thus the update rule could be:
p.data.add_(-torch.mul(p.data, group['weight_decay']).addcdiv(-step_size, exp_avg, denom) )
or equivalently
p.data.add_(-group["weight_decay"], p.data).addcdiv_(-step_size, exp_avg, denom)
Another thing to note is that at least for the problem I tried I had to modify the weight decay significantly from the value used when using Adam optimizer. It is probably a good idea to do a sweep on just the weight decay value when switching from Adam to AdamW.
@ousou Great comment! Sorry guys, I was too busy to dig more here though it will be really interesting. I should make it a rule like 10% time to do things that have no deadlines.
Another thing is, I think the weight decay should not be divided by \sqrt(v_t)?
Another thing is, I think the weight decay should not be divided by \sqrt(v_t)?
@WilliamLwj Yes, that's correct. Neither of the solutions above should do that, since the addcdiv-part is only using denom and exp_avg, not the weight decay. I noticed though that in my previous comment the suggested implementations actually seem to return different values, and I'm not really sure why. I think the second one is correct, but I can't be sure.
I noticed though that in my previous comment the suggested implementations actually seem to return different values, and I'm not really sure why. I think the second one is correct, but I can't be sure.
It seems that the issue in the first version is with the -torch.mul part getting the minus wrong and thus getting the sign wrong in what's added to the weights. The correct version would be
p.data.add_((-torch.mul(p.data, group['weight_decay'])).addcdiv(-step_size, exp_avg, denom))
but it's easier to just use the second version.
@ousou Should the second version of the code also multiply the learning rate, ie.
p.data.add_(-group['weight_decay'] * group['lr'], p.data)
p.data.addcdiv_(-step_size, exp_avg, denom)
@ousou Should the second version of the code also multiply the learning rate, ie.
p.data.add_(-group['weight_decay'] * group['lr'], p.data)
p.data.addcdiv_(-step_size, exp_avg, denom)
@arunava555 No, the weight decay should not be multiplied by the learning rate - the AdamW paper specifies the update rule this way. In the current version of the implementation in this repository the weight decay is implicitly multiplied by the learning rate, and in this issue there is mentioned how to modify the update rule to not do this. In the AdamW paper it is mentioned that the weight decay should not be multiplied by learning date to decouple the two from each other. Decoupling them means that you can find the optimal learning rate and weight decay independently, i.e. changing the learning rate does not change the optimal weight decay (and vice versa).
In the original Adam the weight decay was also implicitly bound to the learning rate, so you had to find a new optimal weight decay when changing the learning rate. AdamW fixes this issue.