Repository: keras-team/keras

Issue:
Gradient accumulation done wrong
In my opinion gradient accumulation should imitate large batch size (if model changes it's weight only through optimizer - not like in BatchNorm).
So here are 3 tests all with batch_size=12+gradient_accumulation_steps=1 vs batch_size=4+gradient_accumulation_steps=3 https://colab.research.google.com/drive/11MdIRQVQoSobYuev-1bCe7EuoEpU0ghU?usp=sharing 

SGD without weight_decay - works as expected.
SGD with weight_decay - fails, but ok if weight_decay divided by gradient_accumulation_steps
Adam - always fails.

In my opinion issue with SGD occuires because weight decay is applied in every step, but should only after real variable update. This is a performance issue too.

Adam issue as i think has one more reason: this optimizers uses self.iterations which is different in every accumulation step (and shoul use self.iterations // gradient_accumulation_steps).


Discussion:

Patch:
diff --git a/keras/src/backend/jax/optimizer.py b/keras/src/backend/jax/optimizer.py
index fb76e5b389ab..5e0e129688fa 100644
--- a/keras/src/backend/jax/optimizer.py
+++ b/keras/src/backend/jax/optimizer.py
@@ -16,7 +16,7 @@ class JaxOptimizer(base_optimizer.BaseOptimizer):
     def _backend_apply_gradients(self, grads, trainable_variables):
         if self.gradient_accumulation_steps:
             is_update_step = (
-                self.iterations + 1
+                self._iterations + 1
             ) % self.gradient_accumulation_steps == 0
             steps = self.gradient_accumulation_steps
 
@@ -47,6 +47,10 @@ def _backend_apply_gradients(self, grads, trainable_variables):
                 lambda: list(grads),
             )
 
+            # Apply clipping and weight decay.
+            grads = self._clip_gradients(grads)
+            self._apply_weight_decay(trainable_variables)
+
             self._backend_update_step(
                 grads, trainable_variables, self.learning_rate
             )
@@ -71,6 +75,10 @@ def _backend_apply_gradients(self, grads, trainable_variables):
                 g_acc.assign(n_g_acc)
 
         else:
+            # Apply clipping and weight decay.
+            grads = self._clip_gradients(grads)
+            self._apply_weight_decay(trainable_variables)
+
             self._backend_update_step(
                 grads, trainable_variables, self.learning_rate
             )
@@ -101,4 +109,4 @@ def _backend_apply_gradients(self, grads, trainable_variables):
                         + var.value * should_not_overwrite_model_vars_int
                     )
 
-        self.iterations.assign_add(1)
+        self._iterations.assign_add(1)
diff --git a/keras/src/optimizers/base_optimizer.py b/keras/src/optimizers/base_optimizer.py
index 6cb63ae97097..043dff8053fe 100644
--- a/keras/src/optimizers/base_optimizer.py
+++ b/keras/src/optimizers/base_optimizer.py
@@ -110,7 +110,7 @@ def __init__(
                 aggregation="only_first_replica",
             )
         self._track_variable(iterations)
-        self.iterations = iterations
+        self._iterations = iterations
 
         # Create learning rate (schedule or variable)
         if isinstance(
@@ -139,6 +139,15 @@ def __init__(
             self._track_variable(learning_rate)
             self._learning_rate = learning_rate
 
+    @property
+    def iterations(self):
+        if self.gradient_accumulation_steps:
+            return ops.floor_divide(
+                self._iterations, self.gradient_accumulation_steps
+            )
+
+        return self._iterations
+
     def _track_variable(self, variable):
         self._tracker.add_to_store("variables", variable)
 
@@ -281,7 +290,7 @@ def apply_gradients(self, grads_and_vars):
         grads, trainable_variables = zip(*grads_and_vars)
         self.apply(grads, trainable_variables)
         # Return iterations for compat with tf.keras.
-        return self.iterations
+        return self._iterations
 
     def apply(self, grads, trainable_variables=None):
         """Update traininable variables according to provided gradient values.
@@ -343,10 +352,6 @@ def apply(self, grads, trainable_variables=None):
             if scale is not None:
                 grads = [g if g is None else g / scale for g in grads]
 
-            # Apply clipping and weight decay.
-            grads = self._clip_gradients(grads)
-            self._apply_weight_decay(trainable_variables)
-
             # Apply gradient updates.
             self._backend_apply_gradients(grads, trainable_variables)
             # Apply variable constraints after applying gradients.
@@ -372,7 +377,7 @@ def _backend_apply_gradients(self, grads, trainable_variables):
         """
         if self.gradient_accumulation_steps:
             is_update_step = (
-                self.iterations + 1
+                self._iterations + 1
             ) % self.gradient_accumulation_steps == 0
             # `trainable_variables` might have been filtered in previous
             # processing steps, so we need to ensure the correct mapping between
@@ -388,6 +393,11 @@ def _update_step_fn(grads, trainable_variables):
                 grads = [
                     (g + acc_g) / steps for g, acc_g in zip(grads, acc_grads)
                 ]
+
+                # Apply clipping and weight decay.
+                grads = self._clip_gradients(grads)
+                self._apply_weight_decay(trainable_variables)
+
                 self._backend_update_step(
                     grads, trainable_variables, self.learning_rate
                 )
@@ -401,6 +411,10 @@ def _update_step_fn(grads, trainable_variables):
                 ),
             )
         else:
+            # Apply clipping and weight decay.
+            grads = self._clip_gradients(grads)
+            self._apply_weight_decay(trainable_variables)
+
             # Run udpate step.
             self._backend_update_step(
                 grads, trainable_variables, self.learning_rate
@@ -424,7 +438,7 @@ def _update_step_fn(grads, trainable_variables):
                     lambda: None,
                 )
         # Update iteration counter.
-        self.iterations.assign_add(1)
+        self._iterations.assign_add(1)
 
     def _backend_update_step(self, grads, trainable_variables, learning_rate):
         """Collective update_step that can be overridden by the backend.
@@ -586,7 +600,7 @@ def _get_current_learning_rate(self):
         if isinstance(
             self._learning_rate, learning_rate_schedule.LearningRateSchedule
         ):
-            return self._learning_rate(self.iterations)
+            return self._learning_rate(self._iterations)
         elif callable(self._learning_rate):
             return self._learning_rate()
         return self._learning_rate
@@ -619,7 +633,7 @@ def _overwrite_variables_directly_with_gradients(self, grads, vars):
                 if self.gradient_accumulation_steps:
                     # Utilize a stateless manner for JAX compatibility
                     steps = self.gradient_accumulation_steps
-                    is_update_step = (self.iterations + 1) % steps == 0
+                    is_update_step = (self._iterations + 1) % steps == 0
                     acc_g = self._accumulated_gradients[
                         self._get_variable_index(v)
                     ]
@@ -964,7 +978,10 @@ def _untrack_variable(self, variable):
             value of the gradients since the last update. This is known as
             "gradient accumulation". This can be useful
             when your batch size is very small, in order to reduce gradient
-            noise at each update step.
+            noise at each update step. EMA frequency will look at "accumulated"
+            iterations value (optimizer steps // gradient_accumulation_steps).
+            Learning rate schedules will look at "real" iterations value
+            (optimizer steps).
 """
 
 


Test patch:
diff --git a/keras/src/backend/tensorflow/optimizer_distribute_test.py b/keras/src/backend/tensorflow/optimizer_distribute_test.py
index d93d3f0e1538..2d55c21796c5 100644
--- a/keras/src/backend/tensorflow/optimizer_distribute_test.py
+++ b/keras/src/backend/tensorflow/optimizer_distribute_test.py
@@ -194,16 +194,19 @@ def test_gradient_accumulation(self):
             self.assertAllClose(
                 optimizer._accumulated_gradients[0], [[1.0, 1.0], [2.0, 2.0]]
             )
-            self.assertAllClose(optimizer.iterations, 1)
+            self.assertAllClose(optimizer._iterations, 1)
+            self.assertAllClose(optimizer.iterations, 0)
             self.strategy.run(lambda: optimizer.apply_gradients([(grads, v)]))
             self.assertAllClose(v, [[1.0, 2.0], [3.0, 4.0]])
             self.assertAllClose(
                 optimizer._accumulated_gradients[0], [[2.0, 2.0], [4.0, 4.0]]
             )
-            self.assertAllClose(optimizer.iterations, 2)
+            self.assertAllClose(optimizer._iterations, 2)
+            self.assertAllClose(optimizer.iterations, 0)
             self.strategy.run(lambda: optimizer.apply_gradients([(grads, v)]))
             self.assertAllClose(v, [[-1.0, 0.0], [-1.0, 0.0]])
             self.assertAllClose(
                 optimizer._accumulated_gradients[0], [[0.0, 0.0], [0.0, 0.0]]
             )
-            self.assertAllClose(optimizer.iterations, 3)
+            self.assertAllClose(optimizer._iterations, 3)
+            self.assertAllClose(optimizer.iterations, 1)
diff --git a/keras/src/optimizers/optimizer_test.py b/keras/src/optimizers/optimizer_test.py
index 5706633a62ad..97dbd1cf2290 100644
--- a/keras/src/optimizers/optimizer_test.py
+++ b/keras/src/optimizers/optimizer_test.py
@@ -195,30 +195,42 @@ def test_gradient_accumulation(self):
             learning_rate=1.0, gradient_accumulation_steps=3
         )
         self.assertEqual(optimizer.gradient_accumulation_steps, 3)
+
+        # Iteration 1
         optimizer.apply_gradients([(grads, v)])
         self.assertAllClose(v, [[1.0, 2.0], [3.0, 4.0]])
         self.assertAllClose(
             optimizer._accumulated_gradients[0], [[1.0, 1.0], [1.0, 1.0]]
         )
-        self.assertAllClose(optimizer.iterations, 1)
+        self.assertAllClose(optimizer._iterations, 1)
+        self.assertAllClose(optimizer.iterations, 0)
+
+        # Iteration 2
         optimizer.apply_gradients([(grads, v)])
         self.assertAllClose(v, [[1.0, 2.0], [3.0, 4.0]])
         self.assertAllClose(
             optimizer._accumulated_gradients[0], [[2.0, 2.0], [2.0, 2.0]]
         )
-        self.assertAllClose(optimizer.iterations, 2)
+        self.assertAllClose(optimizer._iterations, 2)
+        self.assertAllClose(optimizer.iterations, 0)
+
+        # Iteration 3
         optimizer.apply_gradients([(grads, v)])
         self.assertAllClose(v, [[0.0, 1.0], [2.0, 3.0]])
         self.assertAllClose(
             optimizer._accumulated_gradients[0], [[0.0, 0.0], [0.0, 0.0]]
         )
-        self.assertAllClose(optimizer.iterations, 3)
+        self.assertAllClose(optimizer._iterations, 3)
+        self.assertAllClose(optimizer.iterations, 1)
+
+        # Iteration 4
         optimizer.apply_gradients([(grads, v)])
         self.assertAllClose(v, [[0.0, 1.0], [2.0, 3.0]])
         self.assertAllClose(
             optimizer._accumulated_gradients[0], [[1.0, 1.0], [1.0, 1.0]]
         )
-        self.assertAllClose(optimizer.iterations, 4)
+        self.assertAllClose(optimizer._iterations, 4)
+        self.assertAllClose(optimizer.iterations, 1)
 
     @pytest.mark.skipif(backend.backend() != "tensorflow", reason="Requires TF")
     def test_tf_checkpointing(self):
@@ -281,7 +293,8 @@ def test_overwrite_with_gradient_with_gradient_accumulation(self):
 
         # Iteration 1
         optimizer.apply_gradients([(grad_ones, v), (grad_ones, v2)])
-        self.assertAllClose(optimizer.iterations, 1)
+        self.assertAllClose(optimizer._iterations, 1)
+        self.assertAllClose(optimizer.iterations, 0)
         self.assertAllClose(v, [[1.0, 2.0], [3.0, 4.0]])
         self.assertAllClose(v2, [[1.0, 2.0], [3.0, 4.0]])
         self.assertAllClose(
@@ -290,9 +303,11 @@ def test_overwrite_with_gradient_with_gradient_accumulation(self):
         self.assertAllClose(
             optimizer._accumulated_gradients[1], [[1.0, 1.0], [1.0, 1.0]]
         )
+
         # Iteration 2
         optimizer.apply_gradients([(grad_twos, v), (grad_twos, v2)])
-        self.assertAllClose(optimizer.iterations, 2)
+        self.assertAllClose(optimizer._iterations, 2)
+        self.assertAllClose(optimizer.iterations, 1)
         self.assertAllClose(v, [[2.0, 2.0], [2.0, 2.0]])
         self.assertAllClose(v2, [[-0.5, 0.5], [1.5, 2.5]])
         self.assertAllClose(
@@ -301,9 +316,11 @@ def test_overwrite_with_gradient_with_gradient_accumulation(self):
         self.assertAllClose(
             optimizer._accumulated_gradients[1], [[0.0, 0.0], [0.0, 0.0]]
         )
+
         # Iteration 3
         optimizer.apply_gradients([(grad_ones, v), (grad_ones, v2)])
-        self.assertAllClose(optimizer.iterations, 3)
+        self.assertAllClose(optimizer._iterations, 3)
+        self.assertAllClose(optimizer.iterations, 1)
         self.assertAllClose(v, [[2.0, 2.0], [2.0, 2.0]])
         self.assertAllClose(v2, [[-0.5, 0.5], [1.5, 2.5]])
         self.assertAllClose(
@@ -313,6 +330,49 @@ def test_overwrite_with_gradient_with_gradient_accumulation(self):
             optimizer._accumulated_gradients[1], [[1.0, 1.0], [1.0, 1.0]]
         )
 
+    @parameterized.parameters(
+        [
+            ("adam",),
+            ("sgd",),
+            ("adamw",),
+            ("adagrad",),
+            ("rmsprop",),
+            ("adadelta",),
+            ("adamax",),
+            ("lion",),
+            ("nadam",),
+            ("ftrl",),
+            ("adafactor",),
+        ]
+    )
+    def test_gradient_accumulation_with_weigth_decay(self, optimizer):
+        optimizer1 = optimizers.get(
+            {"class_name": optimizer, "config": {"weight_decay": 0.05}}
+        )
+        optimizer3 = optimizers.get(
+            {
+                "class_name": optimizer,
+                "config": {
+                    "weight_decay": 0.05,
+                    "gradient_accumulation_steps": 3,
+                },
+            }
+        )
+        variable1 = backend.Variable([[0.9], [0.5]])
+        variable3 = backend.Variable([[0.9], [0.5]])
+
+        for epoch in range(8):
+            grads3 = np.random.random([3, 2, 1]).astype("float32")
+
+            grads1 = backend.convert_to_tensor(grads3.mean(axis=0))
+            optimizer1.apply_gradients([(grads1, variable1)])
+
+            for batch in range(3):
+                grads3_ = backend.convert_to_tensor(grads3[batch])
+                optimizer3.apply_gradients([(grads3_, variable3)])
+
+        self.assertAllClose(variable1, variable3)
+
     def test_setting_lr_to_callable_untracks_lr_var(self):
         adam = optimizers.Adam(learning_rate=0.001)
         self.assertLen(adam.variables, 2)