AdamW as a torch Optimizer

Write AdamW as a real torch.optim.Optimizer subclass, so it can drive any model the same way the built-in one does.

An optimiser keeps its hyperparameters in self.param_groups, a list of dicts that each hold some "params" and their settings. It keeps per-parameter memory in self.state[p], a dict that starts empty. step() walks every parameter that has a gradient and updates it in place. For AdamW, with t the number of steps this parameter has taken so far, including this one:

p = p * (1 - lr * weight_decay)          # decoupled weight decay
m = beta1 * m + (1 - beta1) * g
v = beta2 * v + (1 - beta2) * g * g
m_hat = m / (1 - beta1 ** t)
v_hat = v / (1 - beta2 ** t)
p = p - lr * m_hat / (sqrt(v_hat) + eps)

The weight decay is the point of AdamW. Adam with L2 adds weight_decay * p to the gradient, so the decay then goes through the adaptive scaling and parameters with large gradients are barely decayed. AdamW shrinks the weights directly and keeps the decay out of m and v.

Your results must match torch.optim.AdamW to float precision. Do not call or subclass the built-in Adam or AdamW.

Function signature

class AdamW(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2):
    def step(self, closure=None):

Example

Input

p = torch.nn.Parameter(torch.tensor([1.0, -2.0]))
opt = AdamW([p], lr=0.1, weight_decay=0.5)
p.grad = torch.tensor([0.5, 0.5])
opt.step()
p

Output

tensor([ 0.8500, -2.0000])

Reasoning

The decay multiplies p by 1 - 0.1 * 0.5 = 0.95, which gives [0.95, -1.9]. On the first step m_hat = g and sqrt(v_hat) = |g|, so the Adam step is lr * sign(g) = 0.1 for both entries. That gives [0.85, -2.0].

Constraints and notes

• Call super().__init__(params, defaults) so that parameter groups work, including groups with their own lr or weight_decay.
• Skip a parameter whose .grad is None. Do not decay it and do not create state for it.
• Keep m, v and the step count in self.state[p], so that state_dict() and load_state_dict() save and restore them.
• If a closure is given, call it with gradients enabled and return its loss.