Support `exclude_from_weight_decay` in AdamW

[jarednielsen]

Describe the feature and the current behavior/state.

The LAMB optimizer has this option, but AdamW does not. This is necessary to train transformer models with Adam.

Relevant information

• Are you willing to contribute it (yes/no): no
• Are you willing to maintain it going forward? (yes/no): no
• Is there a relevant academic paper? (if so, where):
• Is there already an implementation in another framework? (if so, where): Yes, TF 1.
• Was it part of tf.contrib? (if so, where): no

Which API type would this fall under (layer, metric, optimizer, etc.)

Optimizer

Who will benefit with this feature?

User training NLP models with LayerNorm.

[MHStadler]

I think this should be quite easy to implement. follwing the example in LAMB

However, looking at the current DecoupledWeightDecayExtension I'm confused by the additional optional parameter in optimize, that seems to serve exactly this purpose:

def minimize(self, loss, var_list, grad_loss=None, name=None, decay_var_list=None)

Looking at the base minimize, I'm struggling to see how decay_var_list could ever be anything but its default None:

def minimize(self, loss, var_list, grad_loss=None, name=None)

Unless I'm missing something here, this method argument can be removed. Adding a new class member variable instead and it should work out of the box

[dathudeptrai]

@bhack will support ?

[bhack]

/cc @PhilJd

[bhack]

We had this in https://github.com/tensorflow/models/blob/master/official/nlp/optimization.py#L112

[PhilJd]

@MHStadler is right, decay_var_list serves exactly this purpose, the only difference is that you'd have to specify which variables should be decayed instead of which should be excluded.
The base minimize calls the apply_* functions which resolve to the ones the decoupled_weight_decay extension extends. The apply functions access the decay variables stored as a member variable in minimize.

[bhack]

@seanpmorgan Can we close this?