The Story — Three Hikers Going Downhill
The careful walker (SGD) feels the slope and takes one step downhill. Every step depends only on the ground under his feet right now. In a narrow valley he zig-zags from wall to wall.
The skateboarder (Momentum) builds up speed. Past slopes keep pushing him forward. The zig-zags cancel out and he rolls along the valley floor fast.
The smart hiker (Adam) has momentum too, but she also measures how bumpy each direction is. On steep, jumpy sides she takes small steps; on gentle, steady ground she strides out.
An optimizer decides how to turn gradients into weight updates. In this lesson you will write all three yourself, check them against Keras, and race them on different loss surfaces.
The Update Rules
| Optimizer | State it remembers | Update for each weight w with gradient g |
|---|---|---|
| SGD | nothing | w ← w − lr·g |
| Momentum | velocity v | v ← β·v − lr·g ; w ← w + v |
| Adam | m (mean of g), s (mean of g²), step t | m ← β₁m + (1−β₁)g ; s ← β₂s + (1−β₂)g²m̂ = m/(1−β₁ᵗ) ; ŝ = s/(1−β₂ᵗ) ; w ← w − lr·m̂/(√ŝ + ε) |
Side-to-side gradients flip sign each step, so momentum cancels them out. The along-the-valley gradient always points the same way, so momentum adds it up.
Writing the Optimizers
Each optimizer is a tf.Module, so its state (velocity, m, s) is tracked and saved in checkpoints — exactly like in Keras. State is created lazily, one slot per variable.
import tensorflow as tf
class SGD(tf.Module):
def __init__(self, lr=0.01):
super().__init__()
self.lr = lr
def apply(self, grads, variables):
for g, v in zip(grads, variables):
v.assign_sub(self.lr * g)
class Momentum(tf.Module):
def __init__(self, lr=0.01, beta=0.9):
super().__init__()
self.lr, self.beta = lr, beta
self.velocity = [] # one slot per variable
def apply(self, grads, variables):
if not self.velocity:
self.velocity = [tf.Variable(tf.zeros_like(v)) for v in variables]
for g, v, vel in zip(grads, variables, self.velocity):
vel.assign(self.beta * vel - self.lr * g)
v.assign_add(vel)
class Adam(tf.Module):
def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-7):
super().__init__()
self.lr, self.b1, self.b2, self.eps = lr, beta1, beta2, eps
self.t = tf.Variable(0.0)
self.m, self.s = [], []
def apply(self, grads, variables):
if not self.m:
self.m = [tf.Variable(tf.zeros_like(v)) for v in variables]
self.s = [tf.Variable(tf.zeros_like(v)) for v in variables]
self.t.assign_add(1.0)
for g, v, m, s in zip(grads, variables, self.m, self.s):
m.assign(self.b1 * m + (1 - self.b1) * g)
s.assign(self.b2 * s + (1 - self.b2) * tf.square(g))
m_hat = m / (1 - self.b1 ** self.t)
s_hat = s / (1 - self.b2 ** self.t)
v.assign_sub(self.lr * m_hat / (tf.sqrt(s_hat) + self.eps))
print("optimizers ready")Check Them Against Keras
If our maths is right, our optimizers must give the same weights as the Keras ones after the same steps.
def loss_fn(w): # a narrow valley: steep in w[1], gentle in w[0]
return w[0] ** 2 + 20.0 * w[1] ** 2
def run(opt, keras_style=False, steps=50):
w = tf.Variable([2.0, 1.0])
for _ in range(steps):
with tf.GradientTape() as tape:
loss = loss_fn(w)
g = tape.gradient(loss, w)
if keras_style:
opt.apply_gradients([(g, w)])
else:
opt.apply([g], [w])
return w.numpy(), float(loss_fn(w))
pairs = [
("SGD ", SGD(0.02), tf.keras.optimizers.SGD(0.02)),
("Momentum", Momentum(0.02, 0.7), tf.keras.optimizers.SGD(0.02, momentum=0.7)),
("Adam ", Adam(0.1), tf.keras.optimizers.Adam(0.1)),
]
for name, ours, keras in pairs:
w_ours, l_ours = run(ours)
w_keras, _ = run(keras, keras_style=True)
print(f"{name} loss after 50 steps: {l_ours:.2e} max diff vs Keras: {abs(w_ours - w_keras).max():.1e}")The differences are tiny float rounding. You have just rebuilt three real optimizers. Now look at the losses: on this narrow valley, plain SGD is by far the slowest. Its step size is limited by the steep direction, so it crawls along the gentle one. Momentum (β = 0.7) races ahead. (With β = 0.9 it overshoots here — try it in the race below.)
Try It — Optimizer Race
Pick a loss surface. Click the map to choose where all three start. Set the learning rates and press Race. Dark = low loss. The white ring is the minimum. Try a large SGD learning rate on the narrow valley and watch it zig-zag or blow up.
| optimizer | step | position | loss | status |
|---|
Choosing an Optimizer
| Optimizer | Typical lr | Strengths | Weaknesses |
|---|---|---|---|
| SGD | 0.01 – 0.1 | Simple; no extra memory | Slow in narrow valleys; lr hard to pick |
| Momentum | 0.01 – 0.1, β = 0.9 | Much faster; smooths noisy gradients; often generalises well | Can overshoot; one extra slot per weight |
| Adam | 0.001 (default) | Works well with little tuning; per-weight step sizes | Two extra slots per weight (3× memory for weights) |
Start with Adam, lr = 0.001. If training is unstable, lower the learning rate. For big vision models trained for a long time, SGD with momentum is still a strong choice. Whatever you pick, the learning rate matters more than the optimizer.
Golden Rules
assign calls.