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