Neural Network from Scratch · lesson 06/7

Adam Optimizer

6 min readModule 05

Plain SGD moves every weight by the same learning rate times its own gradient, so parameters with tiny gradients crawl while large-gradient parameters lurch. Adam keeps two running averages per parameter — the gradient and its square — and uses them to give every parameter a similar step size.

The idea

At step $t$ , with gradient $g_t$ :

m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t, \qquad v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2

$m$ is the first moment — a smoothed gradient, which is momentum. It cancels the noise that flips sign step to step.
$v$ is the second moment — a smoothed squared gradient, which measures how large this parameter's gradients usually are.

Both start at zero, so the early estimates are biased toward zero. Bias correction divides that out:

\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}, \qquad \theta_t = \theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

The ratio $\hat{m}/\sqrt{\hat{v}}$ is close to ±1 whenever the gradient is consistent, so the parameter moves about $\eta$ per step regardless of gradient scale. Defaults: β₁ = 0.9, β₂ = 0.999, ε = 1e-8.

One caveat: mixing L2 regularization into the gradient (adding λθ to g) also gets rescaled by $\sqrt{v}$ . AdamW decouples it by shrinking the weight directly, which is why transformers use it.

Worked example

Feed a constant gradient of 1.0 to a parameter with η = 0.1. Start from m = v = 0.

Step	m	v	m̂	v̂	θ
1	0.100	0.001	1.0	1.0	0.9
2	0.190	0.002	1.0	1.0	0.8
3	0.271	0.003	1.0	1.0	0.7

Both estimates correct to 1.0, so every step is exactly η = 0.1. Without bias correction the first step would use 0.1/√0.001 ≈ 3.162, a step of 0.3162 — three times too large.

In code

import torch

p = torch.nn.Parameter(torch.tensor([1.0]))
opt = torch.optim.Adam([p], lr=0.1)
for _ in range(3):
    opt.zero_grad()
    (1.0 * p).sum().backward()      # constant gradient of 1.0
    opt.step()
    print(round(p.item(), 6))       # 0.9, then 0.8, then 0.7

m = v = 0.0
w, b1, b2, lr, eps = 1.0, 0.9, 0.999, 0.1, 1e-8
for t in range(1, 4):
    g = 1.0
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g ** 2
    m_hat, v_hat = m / (1 - b1 ** t), v / (1 - b2 ** t)
    w -= lr * m_hat / (v_hat ** 0.5 + eps)
    print(t, round(w, 6), round(m_hat, 6), round(v_hat, 6))   # t 0.9 1.0 1.0, then 0.8, then 0.7

Check yourself

What quantity does v estimate, and why does dividing by its square root make the step size independent of gradient scale?
A parameter has had one gradient so far, equal to 0.1. What are m̂ and v̂, and what step does Adam take?
Why does AdamW exist as a separate optimizer from Adam with weight_decay set?

Key takeaways

Adam stores a smoothed gradient and a smoothed squared gradient for every parameter.
Bias correction is what makes the first few steps the right size instead of 0.1× or 30× too big.
The normalizer gives each parameter a step near η, which is why Adam tolerates less tuning than SGD.