Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule
Abstract
Most modern optimizers form their momentum as an exponential moving average (EMA) of past gradients, forgetting every direction at one fixed rate. However, the inputs a deep network sees during training can be highly anisotropic, with a few directions queried frequently while most are seen rarely. Preconditioning methods address this anisotropy by wrapping extra processing around this buffer and leave the momentum update itself unchanged. We propose Activation-Keyed Momentum (AK-Momentum), which builds direction-awareness into the momentum update rule. The gradient of a linear layer splits into an input activation that acts as a key and an output-side error that acts as a value. Keying on that activation, AK-Momentum updates the momentum buffer by the canonical delta rule, so each direction is forgotten at a rate set by how often it appears. We prove that it is a valid momentum, that it applies the input-side curvature correction without matrix inversion, and that it clears stale directions faster than EMA under both a fixed and a drifting optimum. It is a drop-in replacement for the momentum buffer of any optimizer, its coefficient transfers across widths under μP, and its extra compute stays between 22.2% and 25.0% of a gated-MLP block's linear cost with no persistent memory. In FineWeb-Edu pretraining, AdamW with AK-Momentum (AK-AdamW) reaches AdamW's validation loss in up to 46.64 ± 4.52% fewer steps at 67M and 22.14 ± 0.81% at 370M over three seeds, and the gain persists at 1B on a Chinchilla-optimal budget. A Muon baseline tuned under the same protocol sits above AK-AdamW at both language-model scales, and the gain holds for SGD, ResNet-18, and ViT-Tiny on CIFAR-10. Training-time diagnostics confirm the predicted mechanism, better gradient tracking and healthier input directions. Code is available at https://github.com/euijinh/AK-Momentum.
Related
- On moments of the Erdős-Hooley Delta-function
- Delta functions, a local coarea formula and marginal distributions
- Delta Redistribution: a variation of State Redistribution for steady-state embedded boundary computations
- Asymptotic delta-like behavior of the weighted taper in Thomson's multitaper F-test
- Gradient Growth for Vlasov-Poisson in background charge distributed by Desingularized Dirac Delta
BibTeX
@misc{hong2026activationkeyed,
title = {{Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule}},
author = {Euijin Hong and Guannan Qu},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2608.19491},
eprint = {2608.19491},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2608.19491},
}
RIS
TY - UNPB
TI - Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule
AU - Euijin Hong
AU - Guannan Qu
PY - 2026
DA - 2026/10/29
PB - arXiv
DO - 10.48550/arxiv.2608.19491
UR - https://arxiv.org/abs/2608.19491
AB - Most modern optimizers form their momentum as an exponential moving average (EMA) of past gradients, forgetting every direction at one fixed rate. However, the inputs a deep network sees during training can be highly anisotropic, with a few directions queried frequently while most are seen rarely. Preconditioning methods address this anisotropy by wrapping extra processing around this buffer and leave the momentum update itself unchanged. We propose Activation-Keyed Momentum (AK-Momentum), which builds direction-awareness into the momentum update rule. The gradient of a linear layer splits into an input activation that acts as a key and an output-side error that acts as a value. Keying on that activation, AK-Momentum updates the momentum buffer by the canonical delta rule, so each direction is forgotten at a rate set by how often it appears. We prove that it is a valid momentum, that it applies the input-side curvature correction without matrix inversion, and that it clears stale directions faster than EMA under both a fixed and a drifting optimum. It is a drop-in replacement for the momentum buffer of any optimizer, its coefficient transfers across widths under μP, and its extra compute stays between 22.2% and 25.0% of a gated-MLP block's linear cost with no persistent memory. In FineWeb-Edu pretraining, AdamW with AK-Momentum (AK-AdamW) reaches AdamW's validation loss in up to 46.64 ± 4.52% fewer steps at 67M and 22.14 ± 0.81% at 370M over three seeds, and the gain persists at 1B on a Chinchilla-optimal budget. A Muon baseline tuned under the same protocol sits above AK-AdamW at both language-model scales, and the gain holds for SGD, ResNet-18, and ViT-Tiny on CIFAR-10. Training-time diagnostics confirm the predicted mechanism, better gradient tracking and healthier input directions. Code is available at https://github.com/euijinh/AK-Momentum.
ER -