Introduction

Exponential decay lowers the learning rate gradually as training progresses. In TensorFlow, the clean modern way to do this is to create a learning-rate schedule object and pass it directly to the optimizer, rather than manually editing the optimizer state in callbacks or ad hoc code.

The Core API

TensorFlow exposes exponential decay through tf.keras.optimizers.schedules.ExponentialDecay.

1import tensorflow as tf
2
3schedule = tf.keras.optimizers.schedules.ExponentialDecay(
4    initial_learning_rate=1e-3,
5    decay_steps=1000,
6    decay_rate=0.96,
7    staircase=False,
8)

This schedule returns a learning rate that shrinks as the optimizer step count increases.

Attach the Schedule to the Optimizer

The schedule object is meant to be passed directly as the optimizer’s learning_rate.

optimizer = tf.keras.optimizers.Adam(learning_rate=schedule)

Then use that optimizer normally when compiling the model.

1model = tf.keras.Sequential([
2    tf.keras.layers.Dense(32, activation="relu", input_shape=(10,)),
3    tf.keras.layers.Dense(1)
4])
5
6model.compile(optimizer=optimizer, loss="mse")

This is the most idiomatic setup in current Keras-based TensorFlow code.

Choosing decay_steps Correctly

The biggest configuration mistake is misunderstanding decay_steps. It is not the number of epochs. It is the number of optimizer steps.

If you want the learning rate to decay roughly once per epoch, compute steps per epoch and use that as a baseline.

1batch_size = 32
2dataset_size = 3200
3steps_per_epoch = dataset_size // batch_size
4
5schedule = tf.keras.optimizers.schedules.ExponentialDecay(
6    initial_learning_rate=1e-3,
7    decay_steps=steps_per_epoch,
8    decay_rate=0.95,
9    staircase=True,
10)

That makes the schedule easier to reason about in training terms.

staircase=True vs Smooth Decay

With staircase=False, the decay is continuous. With staircase=True, it drops in discrete stages.

• smooth decay changes every step
• staircase decay changes every decay_steps block

Neither is universally better. Staircase decay is easier to explain and inspect. Smooth decay can feel more gradual.

Inspect the Learning Rate Values

It helps to print a few schedule values before training so you understand what the curve is actually doing.

for step in [0, 500, 1000, 2000, 5000]:
    print(step, float(schedule(step)))

If those numbers collapse too quickly, your decay_rate or decay_steps is too aggressive.

Match the Schedule to the Optimizer and Dataset

Exponential decay is not automatically a good idea for every training run. Small datasets, very adaptive optimizers, or already-tiny initial learning rates may not benefit much from it. The schedule is most useful when you have a reason to start faster and then gradually become more conservative as optimization settles.

A Common Training Example

1import numpy as np
2import tensorflow as tf
3
4x = np.random.randn(512, 10).astype("float32")
5y = np.random.randn(512, 1).astype("float32")
6
7schedule = tf.keras.optimizers.schedules.ExponentialDecay(
8    initial_learning_rate=1e-3,
9    decay_steps=100,
10    decay_rate=0.9,
11    staircase=True,
12)
13
14optimizer = tf.keras.optimizers.Adam(learning_rate=schedule)
15
16model = tf.keras.Sequential([
17    tf.keras.layers.Dense(32, activation="relu", input_shape=(10,)),
18    tf.keras.layers.Dense(1),
19])
20
21model.compile(optimizer=optimizer, loss="mse")
22model.fit(x, y, epochs=3, batch_size=32, verbose=0)

This is all you need for a basic exponentially decaying learning rate setup.

Common Pitfalls

• Treating decay_steps as epochs instead of optimizer steps makes the schedule decay at the wrong speed.
• Choosing a decay rate that is too aggressive can make learning stall early.
• Combining multiple learning-rate mechanisms without a plan makes behavior harder to interpret.
• Forgetting to inspect actual schedule values makes bad hyperparameters harder to notice.
• Manually mutating the optimizer learning rate in random places is usually less clean than using a schedule object.

Summary

• Use tf.keras.optimizers.schedules.ExponentialDecay for modern TensorFlow learning-rate decay.
• Pass the schedule directly into the optimizer.
• Choose decay_steps based on optimizer steps, not vague intuition.
• Decide whether you want smooth decay or staircase decay.
• Inspect schedule values early so you know whether the decay curve matches your training plan.