ZEMA: confidence-calibrated signed descent with only one state

Abstract

Recent work characterizes the Adam optimizer as sign descent modulated by gradient uncertainty: the update is the sign of the batch gradient, scaled by how strongly the batch supports that sign. We apply this view directly and introduce ZEMA (Z-score Evidence Momentum Adaptation), which measures that support with a per-coordinate $z$-score computed from the per-example gradients in the current batch, converts it into a sign probability, and smooths the resulting update with a single exponential moving average. Because the probability is formed before smoothing rather than after, ZEMA keeps one state per parameter instead of Adam's two, at a $4.6$ increase in per-step time; the per-example variance it requires can also be estimated from a small fraction of the batch. In fine-tuning and pretraining experiments, ZEMA is competitive with Adam, reduces the variability of the validation loss, and in some settings recovers the sign of the full-dataset gradient more consistently. We close with two directions this opens: extending confidence adaptation beyond the sign geometry, to methods such as Muon, and to differentially private training.