Adam Optimizer - PyTorch Basics

Adam Optimizer

What you'll build: adam_step(...), one update of the Adam optimizer, returning the new parameters and moment estimates.

Background

Adam (Kingma & Ba, 2015) combines momentum (a running mean of the gradients, m) with RMSProp (a running mean of the squared gradients, v). Dividing by sqrt(v) gives every parameter its own step size, so parameters with small, consistent gradients still move. It is the default optimizer for transformers, usually in its AdamW variant.

Algorithm

At step t (starting at 1), element-wise:

1. First moment: m = beta1 * m + (1 - beta1) * g
2. Second moment: v = beta2 * v + (1 - beta2) * g²
3. Bias correction: m_hat = m / (1 - beta1^t), v_hat = v / (1 - beta2^t)
4. Update: params = params - lr * m_hat / (sqrt(v_hat) + eps)

Why bias correction? m and v start at 0, so early on they are biased towards 0. At t = 1, m = 0.1 g but m_hat = g, and v_hat = g², so the first step is lr * g / |g|: every parameter moves by about lr, whatever the gradient's scale.

Default Hyperparameters

Parameter	Default	Purpose
lr	0.001	step size
beta1	0.9	decay of the gradient average
beta2	0.999	decay of the squared-gradient average
eps	1e-8	avoids division by zero

Inputs

• params, grads, m, v: float arrays of the same shape.
• t: the current step, an integer starting at 1.

Expected Return Format

Return a dictionary with these keys, each the same shape as params:

• 'params': updated parameters
• 'm': updated first moment (before bias correction)
• 'v': updated second moment (before bias correction)

Constraints

• Don't modify the inputs in place; return new arrays.
• PyTorch track: use tensor ops; don't use torch.optim.Adam (the hidden tests compare your trajectory against it).

Follow-up

• AdamW applies weight decay as params -= lr * wd * params, outside the adaptive step. Why is that better than adding wd * params to the gradient?
• Adam keeps two extra tensors per parameter. How much memory is that for a 7B-parameter model in fp32?

Input: adam_step(np.array([1.0, 2.0]), np.array([0.5, -0.3]), np.zeros(2), np.zeros(2), t=1)
Output: {'params': array([0.999, 2.001]), 'm': array([ 0.05, -0.03]), 'v': array([0.00025, 0.00009])}

import numpy as np

def adam_step(params, grads, m, v, t, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
    """
    Perform one step of the Adam optimizer.

    Args:
        params: Current parameters (np.ndarray)
        grads: Current gradients (np.ndarray, same shape as params)
        m: First moment estimate (np.ndarray, same shape as params)
        v: Second moment estimate (np.ndarray, same shape as params)
        t: Current timestep (int, starting from 1)
        lr: Learning rate
        beta1: First moment decay rate
        beta2: Second moment decay rate
        eps: Small constant for numerical stability

    Returns:
        dict with 'params' (updated), 'm' and 'v' (updated moments,
        before bias correction), all the shape of params
    """
    # Your code here
    pass