An Update Noise Scale for Stateful Optimizers

Gradient noise scales (GNS) aim to measure the ratio between noise and signal in stochastic gradients. Prior work has focused on modelling this quantity for stochastic gradient descent (SGD), fixed preconditioners, and non-Euclidean geometry. However, this does not account for the fact that Adam's update depends on previous optimiser history. We introduce an analogous update noise scale, which is conditioned on previous optimiser history and uses the mean and variance of Adam's next counterfactual update. We show the scale is approximately invariant to the probe size used to estimate it for both Adam and AdamW and admits the fixed preconditioned and standard GNS as special cases. We also show that it changes substantially when we change the momentum parameter, as opposed to other scales that do not. Using the update noise scale, we define a stability metric and study the trade-off between momentum and staleness. Our findings show that our quantity is better calibrated to optimisers with internal memory than existing GNS.