AMSGrad

AMSGrad is a small but important modification to Adam, designed to stop Adam’s per-parameter step sizes from becoming unexpectedly large later in training. Think of it as Adam with a “never forget the worst recent uncertainty” rule: once a parameter has shown highly variable gradients, AMSGrad remains cautious about updating it.

Like Adam, AMSGrad keeps two running statistics for each parameter: momentum m, an average gradient direction, and second moment v, an average of squared gradients. Adam divides the momentum by √v, giving parameters with consistently large gradients smaller effective learning rates. AMSGrad adds a third quantity, maximum second moment, which records the largest bias-corrected v seen so far:

• Adam uses the current gradient-scale estimate, v̂_t.
• AMSGrad uses v̄_t = max(v̄_t−1, v̂_t).
• The update divides by √v̄_t + ε, so its adaptive denominator cannot decrease.

In Adam, the squared-gradient estimate can fall after a quiet stretch. That shrinks the denominator and raises the effective learning rate, even if a later gradient is noisy or misleading. On carefully constructed optimization problems, this behavior can prevent convergence. AMSGrad avoids that failure mode by ensuring adaptive learning rates only stay level or become more conservative. The original analysis provides a convergence guarantee under conditions where Adam’s standard update lacks one.

In a framework such as PyTorch, AMSGrad is enabled with torch.optim.Adam(..., amsgrad=True). It costs one extra state tensor per parameter, increasing optimizer memory beyond Adam’s already substantial momentum and variance buffers. It can help when loss curves show repeated late-stage spikes or unstable oscillation. Its caution also has a cost: an early burst of large gradients can permanently make later updates too small, causing slow progress or a plateau. For this reason, standard Adam or AdamW remains more common, while AMSGrad is a principled choice when stability matters more than rapid adaptation.