TensorFlow Learning Rate Schedules
The learning rate controls how large each weight update is during training. A fixed learning rate is a compromise — high enough to learn quickly, low enough not to overshoot. Learning rate schedules change the rate dynamically during training: high at the start for fast learning, gradually decreasing for fine-grained convergence. The right schedule often provides a 2–5% accuracy improvement over a fixed learning rate with no architectural changes.
Why a Fixed Rate Is Suboptimal
Fixed learning rate journey to the minimum: High LR (0.01): Loss: ╲╱╲╱╲╱╲ (bounces around, never settles at minimum) Low LR (0.00001): Loss: ──────────────────────────────╲ (reaches minimum but takes forever) Schedule (start high, end low): Loss: ╲───────────────────────────╲─ (fast at start, precise at end)
Cosine Decay
Cosine decay starts the learning rate at a maximum value and smoothly decreases it following a cosine curve down to a minimum value. The smooth, gradual decay helps the model converge into a sharper minimum than a step-wise decay.
import tensorflow as tf
cosine_decay = tf.keras.optimizers.schedules.CosineDecay(
initial_learning_rate=0.001, # Start value
decay_steps=10000, # Total steps to decay over
alpha=1e-6 # Minimum learning rate at end
)
optimizer = tf.keras.optimizers.Adam(learning_rate=cosine_decay)
model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
Cosine Decay curve (decay_steps=100):
LR
0.001 │╲
│ ╲
│ ╲
│ ──╲
│ ──────╲
│ ─────────╲____
0.0 └──────────────────────────────── Steps
0 25 50 75 100
Cosine Decay With Warm Restarts
Cosine decay with warm restarts periodically resets the learning rate back to its maximum value before decaying again. Each restart lets the optimizer escape local minima it may have settled into, often finding better solutions than a single long decay.
cosine_restarts = tf.keras.optimizers.schedules.CosineDecayRestarts(
initial_learning_rate=0.001,
first_decay_steps=1000, # Steps in first cycle
t_mul=2.0, # Each cycle is 2× longer than the previous
m_mul=0.9, # Each restart starts at 90% of previous peak LR
alpha=1e-6
)
LR
0.001 │╲ ╲ ╲
│ ╲ ╲ ╲
│ ──╲ ╱──╲ ╱───╲
│ ──╱ ──╱ ──
0.0 └──────────────────────────── Steps
Cycle 1 Cycle 2 Cycle 3 (longer each time)
ExponentialDecay
ExponentialDecay multiplies the learning rate by a fixed decay factor at regular intervals. This is simpler to understand but produces a less smooth decay curve than cosine.
exp_decay = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate=0.01,
decay_steps=1000, # Apply decay every 1000 steps
decay_rate=0.96, # Multiply LR by 0.96 at each decay step
staircase=True # True = step decay; False = continuous decay
)
# After 0 steps: LR = 0.01
# After 1000 steps: LR = 0.01 × 0.96 = 0.0096
# After 2000 steps: LR = 0.01 × 0.96² = 0.0092
# After 10000 steps:LR = 0.01 × 0.96^10 = 0.0066
Linear Warmup + Cosine Decay
Modern training recipes for large models start with a warmup phase where the learning rate increases linearly from near-zero to its target value. This prevents the large initial weight updates that destabilize training when weights are still random. After warmup, cosine decay brings it back down.
import tensorflow as tf
import numpy as np
class WarmupCosineDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
def __init__(self, peak_lr, warmup_steps, total_steps, min_lr=1e-6):
self.peak_lr = peak_lr
self.warmup_steps = warmup_steps
self.total_steps = total_steps
self.min_lr = min_lr
def __call__(self, step):
step = tf.cast(step, tf.float32)
warmup_lr = self.peak_lr * (step / self.warmup_steps)
progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
cosine_lr = self.min_lr + 0.5 * (self.peak_lr - self.min_lr) * \
(1 + tf.cos(np.pi * progress))
return tf.where(step < self.warmup_steps, warmup_lr, cosine_lr)
# 500 warmup steps, 5000 total steps
schedule = WarmupCosineDecay(peak_lr=0.001, warmup_steps=500, total_steps=5000)
optimizer = tf.keras.optimizers.Adam(learning_rate=schedule)
Warmup + Cosine Decay:
LR
0.001 │ ╱╲
│ ╱ ╲
│ ╱ ╲
│ ╱ ──╲
│ ╱ ───────╲___
0.0 └──────────────────────────── Steps
0 500 (warmup) 5000
│
Peak LR
PolynomialDecay
poly_decay = tf.keras.optimizers.schedules.PolynomialDecay(
initial_learning_rate=0.01,
decay_steps=5000,
end_learning_rate=0.0001,
power=2.0, # Quadratic decay (power=1.0 is linear)
cycle=False
)
Step Decay (Manual with LambdaCallback)
# Halve the learning rate every 10 epochs
def step_decay(epoch):
initial_lr = 0.01
drop = 0.5
epochs_drop = 10
lr = initial_lr * (drop ** (epoch // epochs_drop))
return float(lr)
lr_scheduler = tf.keras.callbacks.LearningRateScheduler(step_decay, verbose=1)
model.fit(x_train, y_train, epochs=50, callbacks=[lr_scheduler])
Monitoring the Learning Rate During Training
# Log the learning rate as a metric so it appears in TensorBoard
class LRLogger(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
lr = float(self.model.optimizer.learning_rate)
logs['lr'] = lr
print(f" LR = {lr:.6f}")
model.fit(x_train, y_train,
epochs=50,
callbacks=[LRLogger(),
tf.keras.callbacks.TensorBoard('./logs')])
Schedule Comparison
Schedule Best For ────────────────────────────────────────────────────────────────── Fixed LR Quick experiments, prototypes ExponentialDecay Simple projects; easy to tune CosineDecay Most classification tasks; general purpose CosineDecayRestarts Escaping local minima; longer training runs Warmup + Cosine Large models, Transformers, BERT fine-tuning PolynomialDecay Research-standard decay for comparison papers ReduceLROnPlateau When you cannot predict optimal decay schedule ──────────────────────────────────────────────────────────────────
Practical Starting Recipe
For most projects, start with this recipe: 1. Use Adam optimizer 2. Initial LR = 0.001 3. Add ReduceLROnPlateau(factor=0.5, patience=5) 4. If training for many epochs, switch to CosineDecay For fine-tuning pre-trained models: 1. Use Adam or SGD 2. Initial LR = 1e-5 (very low to protect pre-trained weights) 3. Optionally add 100-step warmup 4. CosineDecay over fine-tuning steps
Learning rate schedules are one of the highest-impact hyperparameters you can tune. The difference between a fixed learning rate and a well-designed schedule often equals the difference between a mediocre model and a state-of-the-art one, without any changes to architecture or data. The next topic covers regularization — techniques that prevent models from memorizing training data so they perform well on unseen examples.