tf.keras.optimizers.experimental.AdamW only support constant weight_decay

[x10000year]

Click to expand!

Issue Type

Feature Request

Source

source

Tensorflow Version

2.8

Custom Code

No

OS Platform and Distribution

No response

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

tf.keras.optimizers.experimental.AdamW only supports constant weight decay. But usually we want the weight_decay value to decay with learning rate schedule.

Standalone code to reproduce the issue

The legacy tfa.optimizers.AdamW supports callable weight_decay, which is much better.

Relevant log output

No response

[gadagashwini]

@x10000year, Thanks for proposing new feature request in TF.
Could you elaborate the feature with code snippet for better understanding. Thanks!

[x10000year]

model = MyModel()
schedule = tf.keras.optimizers.schedules.PolynomialDecay(
    1.0,
    decay_steps=10000)
learning_rate_fn = lambda: learning_rate * schedule(model.iterations.read_value())
weight_decay_fn = lambda: weight_decay * schedule(model.iterations.read_value())
optimizer = tfa.optimizers.AdamW(weight_decay=weight_decay_fn,
                                 learning_rate=learning_rate_fn)

As above code shows, tfa.optimizers.AdamW allows us to specify the schedule of weight decay, which should be in proportional with learning rate schedule. However, I cannot do the same thing with tf.keras.optimizers.experimental.AdamW.

@gadagashwini

[sachinprasadhs]

Thanks for opening this issue. Development of keras moved to separate repository https://github.com/keras-team/keras/issues

Please post this issue on keras-team/keras repo.
To know more see;
https://discuss.tensorflow.org/t/keras-project-moved-to-new-repository-in-https-github-com-keras-team-keras/1999
Thank you!

[google-ml-butler]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler]

Closing as stale. Please reopen if you'd like to work on this further.