Introduction

Gradient accumulation lets you simulate a larger effective batch size by summing gradients over several smaller batches before applying an optimizer step. In TF.Keras, the cleanest way to keep model.fit while adding this behavior is to subclass keras.Model and override train_step.

Why Gradient Accumulation Helps

Suppose your GPU can only hold a batch size of 8, but you want the optimization behavior of batch size 32. You can process four micro-batches of 8, accumulate their gradients, and then update the weights once.

That gives an effective batch size of:

micro_batch_size * accumulation_steps

This is useful when:

• memory limits prevent a larger real batch
• training is noisy with very small batches
• you want to preserve the fit API, callbacks, and metrics

Custom train_step Approach

The main idea is:

1. compute gradients for each micro-batch
2. add them into accumulation buffers
3. apply the average or summed gradients every N steps
4. reset the buffers

A compact implementation looks like this:

1import tensorflow as tf
2from tensorflow import keras
3
4
5class AccumModel(keras.Model):
6    def __init__(self, accumulation_steps=4, **kwargs):
7        super().__init__(**kwargs)
8        self.accumulation_steps = accumulation_steps
9        self.accum_step_counter = tf.Variable(0, trainable=False, dtype=tf.int64)
10        self.gradient_accumulators = []
11
12    def compile(self, optimizer, loss, metrics=None, **kwargs):
13        super().compile(optimizer=optimizer, loss=loss, metrics=metrics, **kwargs)
14        self.gradient_accumulators = [
15            tf.Variable(tf.zeros_like(v), trainable=False)
16            for v in self.trainable_variables
17        ]
18
19    def train_step(self, data):
20        x, y = data
21
22        with tf.GradientTape() as tape:
23            y_pred = self(x, training=True)
24            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
25
26        gradients = tape.gradient(loss, self.trainable_variables)
27
28        for acc, grad in zip(self.gradient_accumulators, gradients):
29            acc.assign_add(grad)
30
31        self.accum_step_counter.assign_add(1)
32
33        if tf.equal(self.accum_step_counter % self.accumulation_steps, 0):
34            scaled = [g / tf.cast(self.accumulation_steps, g.dtype) for g in self.gradient_accumulators]
35            self.optimizer.apply_gradients(zip(scaled, self.trainable_variables))
36            for acc in self.gradient_accumulators:
37                acc.assign(tf.zeros_like(acc))
38
39        self.compiled_metrics.update_state(y, y_pred)
40        return {m.name: m.result() for m in self.metrics} | {"loss": loss}
41
42
43inputs = keras.Input(shape=(10,))
44x = keras.layers.Dense(32, activation="relu")(inputs)
45outputs = keras.layers.Dense(1)(x)
46model = AccumModel(inputs=inputs, outputs=outputs, accumulation_steps=4)
47model.compile(optimizer="adam", loss="mse", metrics=["mae"])

This keeps the familiar fit workflow.

Using It With fit

1import numpy as np
2
3x = np.random.randn(128, 10).astype("float32")
4y = np.random.randn(128, 1).astype("float32")
5
6model.fit(x, y, batch_size=8, epochs=3)

With batch_size=8 and accumulation_steps=4, the optimizer updates as if the effective batch were 32.

Important Detail: Final Partial Accumulation

If the epoch ends before the accumulation counter reaches the exact step boundary, you may have leftover gradients that were never applied. In production code, handle that carefully by flushing remaining accumulated gradients at epoch end or designing dataset size and step count so the remainder is acceptable.

Loss Scaling and Metrics

A common choice is to divide accumulated gradients by accumulation_steps before applying them. That keeps the update magnitude comparable to a real larger batch.

Metrics usually update every micro-batch, which is fine for most training dashboards. The main thing to be consistent about is the optimizer step schedule, not the UI frequency.

Distributed Training Considerations

Gradient accumulation is different from distributed data parallelism. If you later combine it with strategies such as tf.distribute.MirroredStrategy, think carefully about where gradients are already being reduced and how many samples each update truly represents.

Common Pitfalls

A common mistake is applying gradients every batch and still calling it gradient accumulation. If the optimizer steps every micro-batch, you are not accumulating.

Another mistake is forgetting to divide by accumulation_steps, which changes the effective learning rate.

Developers also often overlook the leftover-gradient case at the end of an epoch, which can silently skip updates.

Summary

• Gradient accumulation simulates a larger batch using smaller micro-batches.
• In TF.Keras, overriding train_step is the cleanest way to keep model.fit.
• Accumulate gradients for N steps, then apply them once.
• Scale accumulated gradients if you want behavior closer to a true larger batch.
• Handle leftover gradients at epoch boundaries deliberately.