Introduction

Gradient accumulation lets you simulate a larger effective batch size by summing gradients across several smaller mini-batches before applying an optimizer step. In TF.Keras, the clean way to do this while still using model.fit() is not to rewrite the whole training loop, but to override train_step in a custom Model subclass.

Why Override train_step Instead of Replacing fit

model.fit() already handles callbacks, metrics, validation, progress bars, distribution strategy integration, and data iteration. If you replace it entirely, you throw away a lot of useful framework behavior.

Overriding train_step gives you control over one batch update while keeping the rest of the Keras training stack. That is usually the right level for gradient accumulation.

A Working Pattern

The model below accumulates gradients for accum_steps mini-batches and applies them only when the accumulation counter reaches that value. For clarity, this version uses eager execution during training.

1import tensorflow as tf
2
3
4class AccumModel(tf.keras.Model):
5    def __init__(self, accum_steps, *args, **kwargs):
6        super().__init__(*args, **kwargs)
7        self.accum_steps = tf.constant(accum_steps, dtype=tf.int64)
8        self.accum_step_counter = tf.Variable(0, dtype=tf.int64, trainable=False)
9        self.gradient_accumulators = []
10
11    def compile(self, optimizer, loss, metrics=None, **kwargs):
12        super().compile(optimizer=optimizer, loss=loss, metrics=metrics, **kwargs)
13        self.gradient_accumulators = [
14            tf.Variable(tf.zeros_like(var), trainable=False)
15            for var in self.trainable_variables
16        ]
17
18    def train_step(self, data):
19        x, y = data
20
21        with tf.GradientTape() as tape:
22            y_pred = self(x, training=True)
23            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
24            loss = loss / tf.cast(self.accum_steps, loss.dtype)
25
26        gradients = tape.gradient(loss, self.trainable_variables)
27
28        for acc, grad in zip(self.gradient_accumulators, gradients):
29            if grad is not None:
30                acc.assign_add(grad)
31
32        self.accum_step_counter.assign_add(1)
33
34        if tf.equal(self.accum_step_counter % self.accum_steps, 0):
35            self.optimizer.apply_gradients(zip(self.gradient_accumulators, 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}
41
42
43inputs = tf.keras.Input(shape=(4,))
44x = tf.keras.layers.Dense(16, activation="relu")(inputs)
45outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
46
47model = AccumModel(accum_steps=4, inputs=inputs, outputs=outputs)
48model.compile(
49    optimizer=tf.keras.optimizers.Adam(),
50    loss=tf.keras.losses.BinaryCrossentropy(),
51    metrics=[tf.keras.metrics.BinaryAccuracy()],
52    run_eagerly=True,
53)
54
55x_train = tf.random.normal((64, 4))
56y_train = tf.cast(tf.reduce_sum(x_train, axis=1) > 0, tf.float32)
57
58model.fit(x_train, y_train, batch_size=4, epochs=2)

The important detail is dividing the loss by the accumulation step count before computing gradients. Without that scaling, each optimizer update would be too large.

Effective Batch Size

If your physical mini-batch size is 4 and accum_steps is 4, the effective batch size is 16. That is the main point of the technique: you can mimic larger-batch training without fitting all samples in memory at once.

This is especially useful for large models, long sequences, or high-resolution inputs where the real bottleneck is memory rather than compute.

Metrics and Update Timing

Metrics usually update every mini-batch, not only when gradients are applied. That is fine, but you should remember that the displayed metric count and the optimizer step count are no longer the same thing.

This matters when scheduling learning rates or logging per-step information. If your schedule assumes one optimizer step per batch, you need to adapt it to the accumulation interval.

About Current TensorFlow Optimizers

In current TensorFlow and Keras releases, some optimizers expose built-in gradient accumulation options. When that is available and matches your needs, it is simpler than custom code. The custom train_step approach is still valuable when you need full control, want optimizer-independent behavior, or need to combine accumulation with custom loss logic.

Common Pitfalls

The biggest pitfall is forgetting to divide the loss by the accumulation step count before computing gradients. That changes optimization dynamics and usually destabilizes training.

Another mistake is failing to reset the accumulator tensors after applying gradients. If the old values stay around, the next update includes stale gradients.

Developers also overlook the final partial accumulation at the end of an epoch. If the batch count is not divisible by accum_steps, you may want logic to flush the remaining gradients instead of dropping them.

Finally, be careful with schedules, logging, and callbacks that assume one optimizer update per batch. Gradient accumulation changes that assumption.

Summary

• In TF.Keras, the clean way to add gradient accumulation while keeping model.fit() is to override train_step.
• Accumulate gradients across several mini-batches and apply them only at the chosen interval.
• Scale the loss before gradient calculation so the effective update matches the intended batch size.
• Reset accumulator tensors after each optimizer step.
• Watch out for learning-rate schedules and end-of-epoch remainder batches, because accumulation changes the meaning of a training step.