PyTorch Adam: optimizer steps, state, and weight decay

PyTorch Adam updates parameters using running estimates of gradients and squared gradients. To use it in a training loop, create the optimizer once, clear gradients at the intended boundary, compute the loss, call backward, and then call step. The optimizer’s history affects what happens on the next update.

This guide follows five small examples through their actual outputs. We ran 41 checks with Python 3.12.14 and PyTorch 2.15.0+cpu, comparing updates against independent arithmetic. The fixtures use float64 CPU tensors and disable foreach and fused execution. They demonstrate optimizer behavior, not GPU speed or training quality.

Start with one Adam update

Our parameter has two values. Squaring and summing them produces a loss whose gradient is easy to calculate: twice each value. We use a learning rate of 0.1 to make the movement visible, with the default betas and epsilon. This rate is a teaching choice, not a recommendation for your model.

import torch

p = torch.nn.Parameter(torch.tensor([1.0, -2.0], dtype=torch.float64))
opt = torch.optim.Adam([p], lr=0.1, foreach=False, fused=False)
opt.zero_grad(set_to_none=True)
loss = p.square().sum()
loss.backward()
gradient = p.grad.clone()
opt.step()
state = opt.state[p]
print("gradient:", gradient.tolist())
print("parameter:", [round(x, 6) for x in p.detach().tolist()])
print("first moment:", [round(x, 6) for x in state["exp_avg"].tolist()])
print("second moment:", [round(x, 6) for x in state["exp_avg_sq"].tolist()])
print("step:", int(state["step"].item()))

gradient: [2.0, -4.0]
parameter: [0.9, -1.9]
first moment: [0.2, -0.4]
second moment: [0.004, 0.016]
step: 1

The gradient is [2, -4]. The stored first moment is [0.2, -0.4], and the stored second moment is [0.004, 0.016]. At this first step, dividing those estimates by 0.1 and 0.001 respectively corrects their initial bias. The corrected values are the gradient and its square.

For the first component, the update is 1 - 0.1 * 2 / (2 + 1e-8). For the second, it is -2 + 0.1 * 4 / (4 + 1e-8). The printed values are rounded to six decimal places; the tests compare the full values with a tolerance.

The two components move by almost the same absolute amount here even though their gradients differ in magnitude. That observation describes this first update from zero moment estimates. It does not mean every later Adam update has size equal to the learning rate.

Know which settings the example uses

The PyTorch 2.15 Adam reference lists defaults of lr=0.001, betas=(0.9, 0.999), eps=1e-8, and weight_decay=0. The beta values control the running estimates; epsilon is added in the update denominator. This guide leaves AMSGrad and maximization disabled.

In a model, pass the parameters you intend to train, typically through model.parameters(). Keep the optimizer outside the batch loop so its state survives between steps. If you change the learning rate during training, our PyTorch learning-rate scheduler guide covers that separate operation.

Clear gradients deliberately

Calling step does not clear the gradient. The next example performs two backward calls on newly computed losses without changing the parameter between them.

import torch

p = torch.nn.Parameter(torch.tensor(2.0, dtype=torch.float64))
opt = torch.optim.Adam([p], lr=0.1, foreach=False, fused=False)
p.square().backward()
first = p.grad.item()
p.square().backward()
accumulated = p.grad.item()
opt.zero_grad(set_to_none=True)
cleared = p.grad is None
p.square().backward()
fresh = p.grad.item()
opt.step()
after_step = p.grad.item()
print("first, accumulated:", first, accumulated)
print("cleared to None:", cleared)
print("fresh, after step:", fresh, after_step)

first, accumulated: 4.0 8.0
cleared to None: True
fresh, after step: 4.0 4.0

The first backward call gives 4. The second adds another 4, producing 8. After clearing to None, another backward call gives 4 again. That value remains available after the optimizer step.

If each batch should produce its own update, clear the previous gradients before accumulating the next batch’s contribution. If several batches should contribute to one update, choose that boundary deliberately. Our gradient-accumulation guide explains the larger training-loop pattern. An accidental leftover gradient is not evidence that accumulation is configured correctly.

Preserve Adam’s history across steps

To isolate the optimizer from a model, the remaining examples assign fixed gradients directly. Each assignment replaces the previous gradient tensor. This is a controlled fixture, not a replacement for backward in ordinary training.

import torch

def run(reset_each_step):
    p = torch.nn.Parameter(torch.tensor(1.0, dtype=torch.float64))
    opt = torch.optim.Adam([p], lr=0.1, foreach=False, fused=False)
    trajectory = []
    for g in [1.0, -1.0, 0.5]:
        if reset_each_step:
            opt = torch.optim.Adam([p], lr=0.1, foreach=False, fused=False)
        p.grad = torch.tensor(g, dtype=torch.float64)
        opt.step()
        trajectory.append(p.item())
    return trajectory, opt.state[p]

persistent, state = run(False)
reset, reset_state = run(True)
print("persistent:", [round(x, 6) for x in persistent])
print("reset:", [round(x, 6) for x in reset])
print("final step counts:", int(state["step"].item()), int(reset_state["step"].item()))

persistent: [0.9, 0.905263, 0.887791]
reset: [0.9, 1.0, 0.9]
final step counts: 3 1

Both runs receive gradients 1, -1, and 0.5 in that order. They begin with the same parameter and make the same first update. At the second step, the persistent optimizer reaches about 0.905263, while the recreated optimizer returns to about 1.0.

The recreated optimizer treats every supplied gradient as its first observation. Its final step count is one; the persistent optimizer’s count is three. Our separate recurrence checks reproduce both paths, including the persistent first and second moments.

Saving model weights alone does not preserve this optimizer history. For a training restart, see our PyTorch checkpoint guide, which covers saving and loading optimizer state alongside model state. Recreating Adam with the same learning rate does not recover its earlier moment estimates.

Distinguish a zero gradient from no gradient

After a first update with gradient 1, we perform a second step without another backward call. One branch clears the existing gradient to None; the other turns it into a tensor containing zero.

import torch

def second_step(use_none):
    p = torch.nn.Parameter(torch.tensor(1.0, dtype=torch.float64))
    opt = torch.optim.Adam([p], lr=0.1, foreach=False, fused=False)
    p.grad = torch.tensor(1.0, dtype=torch.float64)
    opt.step()
    before = p.item()
    opt.zero_grad(set_to_none=use_none)
    opt.step()
    return before, p.item(), int(opt.state[p]["step"].item())

skipped = second_step(True)
zero = second_step(False)
print("None: before, after, step:", tuple(round(x, 6) for x in skipped))
print("zero: before, after, step:", tuple(round(x, 6) for x in zero))

None: before, after, step: (0.9, 0.9, 1)
zero: before, after, step: (0.9, 0.832994, 2)

With None, the parameter stays at approximately 0.9 and its step count stays at one. With a zero gradient, its existing moment history produces another update, reaching about 0.832994, and the step count advances to two. Weight decay is zero in both branches, so decay does not explain the movement.

The zero_grad reference documents this skip-versus-update distinction. Here, set_to_none=False zeros an already existing gradient. Do not interpret it as a promise to allocate gradients for parameters that have never received one.

Compare Adam and AdamW with matching decay settings

Use the same learning rate and weight-decay coefficient when checking the difference. In our runtime, Adam defaults to zero weight decay, while AdamW defaults to 0.01. Comparing untouched defaults would change two things at once.

import torch

def decay_run(kind):
    p = torch.nn.Parameter(torch.tensor(2.0, dtype=torch.float64))
    options = dict(lr=0.1, weight_decay=0.1, foreach=False, fused=False)
    if kind == "AdamW":
        opt = torch.optim.AdamW([p], **options)
    else:
        opt = torch.optim.Adam(
            [p], decoupled_weight_decay=(kind == "decoupled Adam"), **options
        )
    values = []
    for g in [0.25, -0.5, 0.0]:
        p.grad = torch.tensor(g, dtype=torch.float64)
        opt.step()
        values.append(p.item())
    return values, opt.state[p]

results = {kind: decay_run(kind) for kind in ["Adam", "AdamW", "decoupled Adam"]}
for kind, (values, _) in results.items():
    print(kind + ":", [round(x, 6) for x in values])

Adam: [1.9, 1.887059, 1.856736]
AdamW: [1.88, 1.89781, 1.907132]
decoupled Adam: [1.88, 1.89781, 1.907132]

We explicitly set weight decay to 0.1 and supply the same three gradients to all branches. In the first coupled Adam update, the gradient used for the moments is 0.25 + 0.1 * 2 = 0.45. The resulting parameter is approximately 1.9.

AdamW first applies a separate shrinkage of 2 * (1 - 0.1 * 0.1) = 1.98. Its moment calculation uses 0.25 without adding that decay term, and the adaptive update then brings the parameter to approximately 1.88. The difference persists across the next two steps.

With decoupled_weight_decay=True, Adam matched AdamW across all three tested steps, including both moment tensors. That option is present in the tested PyTorch 2.15 API; check your installed version before using the same call. We also tested a fresh AdamW parameter with no gradient: it stayed unchanged and acquired no optimizer state. Supplying a zero gradient instead applied the isolated decay and changed 2 to 1.98.

Inspect the update before tuning a full run

When an Adam training run behaves unexpectedly, inspect a parameter’s gradient before step, its value afterward, and the optimizer’s step count and moment estimates. Confirm that you are clearing at the intended boundary and keeping the same optimizer instance. Record the decay mode explicitly when comparing runs.

The examples establish exact behavior for small, fixed CPU inputs. They do not identify the best optimizer or learning rate for a dataset, measure convergence, or compare CUDA execution paths. Use a task-specific validation metric for those decisions, after checking that the update sequence is the one you intended.