why not AdamW style weight decay

Hello,

While translating your optimizer to Flax (here), I noticed that you are using a traditional weight decay were you add the weight decay to the gradient (here in your implementation):

grad += weight_decay * parameters

Rather than an AdamW style weight decay (which, I believe, is now the default for most optimizers) were you would subtract the weight decay time the learning rate just before returning the parameters:

updated_parameters -= learning_rate * weight_decay * param

Is there a particular reason for that decision ?

[adefazio]

I will look into adding the adamw style weight decay as an option, thanks for the discussion and results!

[adefazio]

I’m going to add adamw style averaging to the implementation this week as it seems popular based on the comments here.