Make API for custom optimizer wrappers more consistent

[reedwm]

There are several optimizer classes which wrap an existing optimizer and add functionality, but they have slightly different APIs The Lookahead, MovingAverage, and SWA each take an optimizer in its constructor and adds additional functionality. Similarly, LossScaleOptimizer from Keras, an optimizer which can prevent numeric underflow by scaling loss and gradients, takes an optimizer in its constructor as well. On the other hand, the extend_with_decoupled_weight_decay function instead creates a new dynamically-created class which subclasses from the original optimizer's class but adds functionality. I think we should make the API for all these classes consistent, so that all optimizer wrappers have the same API.

LossScaleOptimizer is being made non-experimental in TF 2.4, so I need to decide on the final API soon. Ideally, it would be consistent with the optimizer APIs in TF-Addons, which is why I am bringing up this issue. There is a mixed precision RFC describing the proposed changes (no need to read it, I will summarize everything important here).

I summarize the two approaches an optimizer extension can take below:

1. Continue wrapping optimizers with another optimizer class:

This approach continues having optimizers like Lookahead and LossScaleOptimizer wrap other optimizers. An OptimizerWrapper helper class can be created, which automatically deals with wrapping logic such as implementing get_config properly so each optimizer wrapper can focus on just the added functionality.

One question is how to whether __getattribute__ and __setattr__ should be delegate hyperparameters (and potentially other attributes) to the inner optimizer, which could be done in OptimizerWrapper. I proposed doing this in the mixed precision RFC, but can change it. Doing this delegation would allow attributes of the inner optimizer to be accessed Lookahead, LossScaleOptimizer, and other wrappers:

opt = tf.keras.optimizers.SGD(0.1, momentum=0.1)
opt = tfa.optimizers.Lookahead(opt)
...  # Do some training
print(opt.momemtum)
opt.momentum = 0.05

Delegating all attributes in __getattribute__ is tempting, but it does have a significant flaw: it can cause the custom logic in optimizer wrappers not to be called:

class MyCustomOptimizer(tf.keras.optimizers.SGD):

  def apply_gradients_zero_min(self, grads_and_vars):
    grads_and_vars = [(tf.nn.relu(g), v) for g, v in grads_and_vars]
    self.apply_gradients(grads_and_vars)

opt = MyCustomOptimizer(1.)
opt = tf.keras.mixed_precision.LossScaleOptimizer(opt)
opt.apply_gradients_zero_min(...)  # Bug! Does not call LossScaleOptimizer version of apply_gradients.

Custom logic in the function LossScaleOptimizer.apply_gradients is not called, as only the apply_gradients of the inner optimizer is called. Similar issues occur with any optimizer wrapper.

2. Dynamically create a subclass

This approach is used by extend_with_decoupled_weight_decay. It consists of dynamically creating a subclass with two superclasses: The original optimizer (e.g. Adam) and a special extension class which adds the extra optimizer functionality. For example, see the implementation of extend_with_decoupled_weight_decay. Other optimizers could use the same approach.

We want to allow users to, e.g., create an Adam optimizer then pass it to a function which returns an optimizer of the new type. To allow this, we could have a function which takes in the old optimizer, creates the new class and returns an optimizer of the new class, e.g.:

opt = tf.keras.optimizers.SGD(0.1, momentum=0.1)
opt = tfa.optimizers.copy_with_decoupled_weight_decay(opt)

In this case, copy_with_decoupled_weight_decay would create the new optimizer class then use opt.get_config to serialize the old optimizer and from_config to deserialize it into the new optimizer. One flaw is this requires the optimzier to be serializable with get_config and from_config Another flaw is that users might accidentally modify the old optimizer, not realizing that this does not modify the new optimizer.

Alternatively, we could have a function monkey-patch opt.__class__ to the new optimizer class, so the optimizer does not need to be serializable with get_config and from_config:

opt = tf.keras.optimizers.SGD(0.1, momentum=0.1)
# Monkey-patches opt.__class__ to an optimizer class which subclasses 
# both DecoupledWeightDecayExtension and SGD
tfa.optimizers.modify_class_to_use_decoupled_weight_decay(opt)

Which to choose?

Does anyone have any opinions on which approach is superior, and whether all the optimizers should be unified to use the same approach? I think I prefer dynamically creating a subclass because it completely emulates the API of the original optimizer's class.

/CC @CyberZHG, @bhack, @Squadrick, @shreyashpatodia, @PhilJd, who have worked on optimizers in TF-Addons
/CC @omalleyt12, @tomerk, who have worked on Keras optimizers

[PhilJd]

I don't have a strong opinion on this, as you stated, both approaches have their pros and cons. Note though that I haven't worked with Keras lately, so I'm not familiar with the config/scheduler interface and might miss important points.

I'm guilty of introducing the dynamic subclass approach, to keep the API consistent by having an optimizer class that can be instantiated like all standard optimizers. This comes at the cost of somewhat reduced readability due to multiple inheritance. Since then the optimizer logic has become more complex with initialization from config/schedulers etc., making this harder to maintain and read, so I'm very open to an alternative.
However, if going with 1, it's probably safest to allow custom optimizers to overwrite the wrapper classes *apply* functions only to prevent the problem you mentioned and have the optimizer fail loudly otherwise. I don't know if this would be flexible enough to cover all use cases. Unfortunately, the effort to clean this up (tensorflow/community/pull/234) and make the API consistent is on hold.

One thing I'm trying to understand is the need for creating the optimizer from an existing instance. What's the advantage over creating a wrapped optimizer class and directly instantiating this class?
I guess when wrapping an instance, (1) is the more intuitive approach, while (2) works better for the latter.

[bhack]

What about the Multioptimizer for the dynamic approach case? https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/discriminative_layer_training.py

[reedwm]

I think the dynamic subclass approach is pretty clean compared to wrapping an optimizer, as it completely emulates the interface of the original optimizer. It's a bit hard to read, but I think it will be easy to maintain. The DecoupledWeightDecayExtension.from_config method can be modified to call super().from_config, which I think would then make it work even if the base Optimizer changes its from_config implementation.

tensorflow/community/pull/234 will likely be approved soon and the transform_gradients and aggregate_gradients constructor arguments have been implemented (they have been renamed to gradient_transformers and gradient_aggregator respectively). However, the implementation is not flexible enough to support the existing optimizer wrappers in Addons and Keras.

One thing I'm trying to understand is the need for creating the optimizer from an existing instance. What's the advantage over creating a wrapped optimizer class and directly instantiating this class?
I guess when wrapping an instance, (1) is the more intuitive approach, while (2) works better for the latter.

I think it's a better API to convert an existing optimizer, especially when a model supports specify which optimizer to use based on a flag. Compare the two approaches for using decoupled weight decay:

# Approach 1: Use existing API
if FLAGS.optimizer == 'sgd':
  opt_type = tf.keras.optimizers.SGD
  kwargs = {'learning_rate': 0.1}
else:
  opt_type = tf.keras.optimziers.Adam(0.1, beta_1=0.99, beta_2=0.999)
  kwargs = {'learning_rate': 0.1, 'beta_1': 0.99, 'beta_2': 0.999}

if FLAGS.use_dwd:
  opt_type = tfa.optimizers.extend_with_decoupled_weight_decay(opt_type)
  kwargs['weight_decay'] = 0.1

opt = opt_type(**kwargs)

# Approach 2: API which uses existing optimizer instance
if FLAGS.optimizer == 'sgd':
  opt = tf.keras.optimizers.SGD(learning_rate=0.1)
else:
  opt = tf.keras.optimizers.Adam(learning_rate=0.1, beta_1=0.99, beta_2=0.990)

opt = tfa.optimizers.copy_with_decoupled_weight_decay(opt, weight_decay=0.1)

Approach 2 is both shorter and clearer.

What about the Multioptimizer for the dynamic approach case? https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/discriminative_layer_training.py

I didn't know about Multioptimizer until now. Multioptimizer will not work with the dynamic approach case. However, it also doesn't fully work with my proposed optimizer wrapper case, as we cannot delegate __getattribute__ (because there are multiple inner optimizers so we don't know which to delegate to). I think we should treat Multioptimizer as a special case that doesn't have to have the same API as other optimizer wrappers. Most other wrappers extend the functionality of an optimizer in some way, but Multioptimizer instead is an API to use different optimizers for different sets of weights

[bhack]

Cannot we make an hypotesis about extending copy_with_decoupled_weight_decay to handle also the input types that already handle the Multioptimizer?
It is a little bit true that It could be a special case but at the end other then "special inputs" we just overide/wrap apply_gradients. So anyhting so special.

[reedwm]

By "special inputs" do you mean to have copy_with_decoupled_weight_decay also allow a list of (optimizer, layers) pairs to be passed, like as MultiOptimizer.__init__, instead of the single optimizer? The issue with doing that is that copy_with_decoupled_weight_decay would create a dynamic subclass that subclasses from the original optimizer's class. But if there are multiple original optimizers, there is no single class to subclass from.

[bhack]

For "special inputs" I meant that but if it is really a "very special case" cannot we dynaically subclass to Multioptimizer directly? Or we could have other more general use cases for this kind of inputs?

[reedwm]

Multioptimizer cannot be implemented as a dynamic subclass. Since what would its superclasses be? We have multiple original optimizers for Multioptimizer, but we cannot subclass all of them, just one.

I don't think its worth having a general use case for wrapping multiple optimizers, as in Multioptimizer, since it isn't very common. There are many wrapper-like classes that wrap just one optimizer (Lookahead, extend_with_decoupled_weight_decay, LossScaleOptimizer) but just one that wraps multiple optimizers (Multioptimizer). Also, Multioptimizer is unlike the other wrappers in other way as well: it does not actually modify the logic in how gradients are applied to variables, but just directly delegates to one of the inner optimizers for each weight.

[bhack]

For the second part ok It could be just an evaluation about that multioptimizer will be really the only case to handle this kind of inputs or that It Is just different cause It Is only a router on optimizers/layers.

But for the first part probably I don't understand what copy_with_decoupled_weight_decay will return for other optimizers?

I think that the base class for Multioptimizer is optimizer.
But if we don't have other use cases in mind for the layer/optimizer input other then Multioptimizer "wrapper/routing" activity It could be plausibile to maintain it as Is.

[reedwm]

I think its feasible to maintain Multioptimizer as is. If we wish to implement more optimizers that wrap multiple optimizers, we can create a generic MultiOptimizerWrapper that could be used for all of them.

But for the first part probably I don't understand what copy_with_decoupled_weight_decay will return for other optimizers?

copy_with_decoupled_weight_decay would roughly look like this (the example uses DecoupledWeightDecayExtension which is defined here):

def copy_with_decoupled_weight_decay(optimizer, weight_decay):
  class OptimizerWithDecoupledWeightDecay(DecoupledWeightDecayExtension, optimizer.__class__):
    pass

  config = optimizer.get_config()
  config['weight_decay'] = weight_decay
  return OptimizerWithDecoupledWeightDecay.from_config(config)    

copy_with_decoupled_weight_decay copies the optimizer with get_config and from_config, but the new optimizer also subclasses from DecoupledWeightDecayExtension to get the extra functionality.

[WindQAQ]

Hi @reedwm, thanks for bringing this up. I would prefer option 2. BTW, do you have the timeline that the new loss scale optimizer API will land in OSS TensorFlow? We usually pin to 2.4rc few days after it releases, so we want to have some time to prepare in advance. If you don't mind, we can discuss or even reveal the designated API here. Thank you!

Gentle ping to @tensorflow/sig-addons-maintainers because of API change. We should also update the related documentation as well.

Also, I want to notice that because this is a huge API change, and combined with #2122, we can discuss whether we should bump our major version to 1.0.0 to inform users that there are no more backward compatibility.

[reedwm]

The 2.4 branch is being cut in about two weeks. At that point, we cannot make API changes to TensorFlow anymore, but there will still be a few weeks after that before it is released.

If we do end up creating a helper class like OptimizerWrapper, it will likely not make it to TF 2.4 but instead TF 2.5. Which means TF Addons won't be able to use it until 2.5 is released. But the OptimizerWrapper can always be implemented in Addons first. In any case, I don't think we need to rush to make any API changes to Addons, but we do need to hurry in deciding the TensorFlow LossScaleOptimizer API as the TF 2.4 branch is being cut soon.

/CC @omalleyt12 any thoughts between options 1 and 2?

[omalleyt12]

Thanks for opening this discussion!

My preference is for option (1). Conceptually, what these optimizers do is "Delegate most logic to a user-provided optimizer (of any class), while adding additional functionality". Composition seems like the best pattern for this.

IMO, option (2) falls too much into the "magic" realm. A few issues I think we'd run into with approach (2) off the top of my head:

1. What happens to the original optimizer when a user does make_dynamic_subclass(my_optimizer)?

If there is no connection b/t the state of the my_optimizer instance and the new instance, this will be surprising to users. In that case, why not use the class name rather than an instance? But if we use the class name, then it becomes a pain to pass parameters to the dynamic subclass

2. Two optimizers of the same "dynamic class" won't have the same __class__, which will break isinstance checks

3. Extending the dynamic subclasses will be more difficult. Rather than subclassing the optimizer, the user will have to make their own function that creates a dynamic subclass, which seems a lot more difficult

Re option (1), composition of Layers is a fairly established pattern in Keras. For instance, the TimeDistributed and RNN layers use this pattern.

IMO when we have the use case "I want to delegate some functionality of my object to an object of type X, but I don't care what subclass of X is used," this is a classic use case for composition