Repository: tensorflow/addons

Issue:
Lookahead might divide weights with mixed_float16
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
- TensorFlow version and how it was installed (source or binary): 2.3.0-rc2 (source)
- TensorFlow-Addons version and how it was installed (source or binary): Just lookahead extracted from the source (master)
- Python version: 3.6.9
- Is GPU used? (yes/no): yes

When tf.keras.mixed_precision.experimental.set_policy('mixed_float16') is set, using LookAhead divides the layers weights after a few gradient updates (for example 7 in my case).

After investigating, I think the reason is that when mixed_float16 is used with automatic loss scaling, the loss scale is initialized very high. Thus the first few gradient updates do nothing (see https://github.com/tensorflow/tensorflow/blob/582c8d236cb079023657287c318ff26adb239002/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L714 for how updates are skipped).

As a result, LookAhead ends up with its 'slow vars' initialized to 0 instead of the original layer weight. Indeed the vars are initialized to the correct value only on the first iteration during the gradient update (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lookahead.py#L106).

Thus when the gradient update finally does something, the first variable update does average them with the slow variables (initialized to zero). Which would explain the weight division I'm seeing.

I don't have a short code to replicate the bug, but I think the above explanation and links should highlight well enough the issue with the current way LookAhead's slow variables are initialized.

Lookahead might divide weights with mixed_float16
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
- TensorFlow version and how it was installed (source or binary): 2.3.0-rc2 (source)
- TensorFlow-Addons version and how it was installed (source or binary): Just lookahead extracted from the source (master)
- Python version: 3.6.9
- Is GPU used? (yes/no): yes

When tf.keras.mixed_precision.experimental.set_policy('mixed_float16') is set, using LookAhead divides the layers weights after a few gradient updates (for example 7 in my case).

After investigating, I think the reason is that when mixed_float16 is used with automatic loss scaling, the loss scale is initialized very high. Thus the first few gradient updates do nothing (see https://github.com/tensorflow/tensorflow/blob/582c8d236cb079023657287c318ff26adb239002/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L714 for how updates are skipped).

As a result, LookAhead ends up with its 'slow vars' initialized to 0 instead of the original layer weight. Indeed the vars are initialized to the correct value only on the first iteration during the gradient update (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lookahead.py#L106).

Thus when the gradient update finally does something, the first variable update does average them with the slow variables (initialized to zero). Which would explain the weight division I'm seeing.

I don't have a short code to replicate the bug, but I think the above explanation and links should highlight well enough the issue with the current way LookAhead's slow variables are initialized.



Discussion:

Resolving patch:
diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py
--- a/tensorflow_addons/optimizers/lookahead.py
+++ b/tensorflow_addons/optimizers/lookahead.py
@@ -87,7 +87,7 @@ def _create_slots(self, var_list):
             var_list=var_list
         )  # pylint: disable=protected-access
         for var in var_list:
-            self.add_slot(var, "slow")
+            self.add_slot(var, "slow", initializer=var)
 
     def _create_hypers(self):
         self._optimizer._create_hypers()  # pylint: disable=protected-access
@@ -103,17 +103,6 @@ def apply_gradients(self, grads_and_vars, name=None, **kwargs):
         )  # pylint: disable=protected-access
         return super().apply_gradients(grads_and_vars, name, **kwargs)
 
-    def _init_op(self, var):
-        slow_var = self.get_slot(var, "slow")
-        return slow_var.assign(
-            tf.where(
-                tf.equal(self.iterations, tf.constant(0, dtype=self.iterations.dtype)),
-                var,
-                slow_var,
-            ),
-            use_locking=self._use_locking,
-        )
-
     def _look_ahead_op(self, var):
         var_dtype = var.dtype.base_dtype
         slow_var = self.get_slot(var, "slow")
@@ -138,24 +127,22 @@ def weights(self):
         return self._weights + self._optimizer.weights
 
     def _resource_apply_dense(self, grad, var):
-        init_op = self._init_op(var)
-        with tf.control_dependencies([init_op]):
-            train_op = self._optimizer._resource_apply_dense(
-                grad, var
-            )  # pylint: disable=protected-access
-            with tf.control_dependencies([train_op]):
-                look_ahead_op = self._look_ahead_op(var)
-        return tf.group(init_op, train_op, look_ahead_op)
+        train_op = self._optimizer._resource_apply_dense(
+            grad, var
+        )  # pylint: disable=protected-access
+        with tf.control_dependencies([train_op]):
+            look_ahead_op = self._look_ahead_op(var)
+        return tf.group(train_op, look_ahead_op)
 
     def _resource_apply_sparse(self, grad, var, indices):
-        init_op = self._init_op(var)
-        with tf.control_dependencies([init_op]):
-            train_op = self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access
+        train_op = (
+            self._optimizer._resource_apply_sparse(  # pylint: disable=protected-access
                 grad, var, indices
             )
-            with tf.control_dependencies([train_op]):
-                look_ahead_op = self._look_ahead_op(var)
-        return tf.group(init_op, train_op, look_ahead_op)
+        )
+        with tf.control_dependencies([train_op]):
+            look_ahead_op = self._look_ahead_op(var)
+        return tf.group(train_op, look_ahead_op)
 
     def get_config(self):
         config = {


Test patch: