🎓 A1Learner
🔍

PyTorch What Adam Remembers

Adam keeps two running averages for every parameter it updates. The averages live in the optimizer state under the names exp_avg and exp_avg_sq.

The state after one step

One step with a fixed gradient fills both averages. The state keys name them directly.

import torch

p = torch.tensor([1.0], requires_grad=True)
optimizer = torch.optim.Adam([p], lr=0.1)
p.grad = torch.tensor([0.5])
optimizer.step()
state = optimizer.state[p]
print(sorted(state.keys()))
print(f"exp_avg: {state['exp_avg'].item():.4f}")
print(f"exp_avg_sq: {state['exp_avg_sq'].item():.6f}")

One hand update of m

The first update of m blends the gradient 0.5 with the starting value 0. Beta1 keeps nine tenths of the old value.

m = 0.9 * 0.0 + 0.1 * 0.5
print(m)

Bias correction undoes the slow start

Early averages lean toward zero because they start at zero. Adam divides by one minus the beta raised to the step count.

m = 0.05
corrected = m / (1 - 0.9 ** 1)
print(f"{corrected:.4f}")

Example

import torch

torch.manual_seed(0)
p = torch.tensor([1.0], requires_grad=True)
optimizer = torch.optim.Adam([p], lr=0.1)

grads = [0.5, -0.25, 0.75]
beta1, beta2 = 0.9, 0.999
m, v = 0.0, 0.0
print("Step | Hand m    | State m   | Hand v    | State v")
for t, g in enumerate(grads, start=1):
    p.grad = torch.tensor([g])
    optimizer.step()
    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * g * g
    state = optimizer.state[p]
    sm = state["exp_avg"].item()
    sv = state["exp_avg_sq"].item()
    print(f"{t}    | {m:.6f}  | {sm:.6f}  | {v:.6f}  | {sv:.6f}")

state = optimizer.state[p]
print(f"m match to 6 decimals: {round(m, 6) == round(state['exp_avg'].item(), 6)}")
print(f"v match to 6 decimals: {round(v, 6) == round(state['exp_avg_sq'].item(), 6)}")
print(f"Parameter after 3 steps: {p.item():.4f}")

Key points

  • m: the running average of the gradients, stored as exp_avg
  • v: the running average of the squared gradients, stored as exp_avg_sq
  • Beta1: the blending rate for m, with default value 0.9
  • Beta2: the blending rate for v, with default value 0.999
  • Use optimizer.state[p] to read the stored averages for one parameter
  • Each step blends the newest gradient into both averages with those rates

In this example, one parameter starts at 1.0 and Adam runs at rate 0.1. Three gradients are fed in by hand, one per step. The hand values of m and v use plain Python floats and the two blending formulas. The state values come from optimizer.state after each real step. Both columns agree to six printed decimals on all three steps. The small float32 and float64 gap hides below that precision. The parameter itself ends at 0.8098 after the three Adam updates.

Note

Adam stores its averages in float32, and plain Python floats use float64

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

After one step with gradient 0.5 and beta1 of 0.9, what is m?