Adam Optimizer

Adam (Adaptive Moment Estimation) combines Momentum with per-parameter adaptive learning rates:

m_t = \beta_1 m_{t-1} + (1-\beta_1)\nabla L_t \quad \text{(1st moment — momentum)}

v_t = \beta_2 v_{t-1} + (1-\beta_2)(\nabla L_t)^2 \quad \text{(2nd moment — gradient variance)}

Bias correction (critical for early steps):

\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t}

Update:

\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

Default hyperparameters: $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\epsilon = 10^{-8}$

Why Adam beats vanilla SGD in most cases:

Per-parameter learning rates: parameters with large gradients get smaller updates; sparse/small gradient parameters get larger updates
Momentum handles consistent gradient directions
Works well out-of-the-box with less learning rate tuning

Memory cost: Adam stores $\theta$ , $g$ , $m$ , and $v$ per parameter → roughly 4× the parameter size in optimizer state (vs 1× for plain SGD).

AdamW decouples weight decay from the gradient, applying it directly to the parameters:

\theta \leftarrow \theta - \eta\frac{\hat m}{\sqrt{\hat v} + \epsilon} - \eta\lambda\theta

def step(self):
    for group in self.param_groups:
        lr, (beta1, beta2) = group["lr"], group["betas"]
        eps, wd = group["eps"], group["weight_decay"]
        for p in group["params"]:
            if p.grad is None:
                continue
            state = self.state[p]
            t = state.get("t", 0)
            m = state.get("m", torch.zeros_like(p.data))
            v = state.get("v", torch.zeros_like(p.data))
 
            p.data -= lr * wd * p.data              # decoupled weight decay
            g = p.grad.data
            m = beta1 * m + (1 - beta1) * g
            v = beta2 * v + (1 - beta2) * g**2
            m_hat = m / (1 - beta1 ** (t + 1))      # bias correction
            v_hat = v / (1 - beta2 ** (t + 1))
            p.data -= lr * m_hat / (v_hat.sqrt() + eps)
 
            state.update(t=t + 1, m=m, v=v)

Typically weight decay is applied to weights but not to biases or LayerNorm/RMSNorm parameters (use separate parameter groups).

Limitations:

Can generalize worse than well-tuned SGD+momentum on some tasks (especially image classification)
AdamW (decoupled weight decay) fixes the interaction between Adam and L2 Regularization — preferred in transformer training

Adam Optimizer

Linked from