# Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule

arXiv: https://arxiv.org/abs/2608.19491 · HTML: https://mlsanity.com/paper/2608.19491/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 30 references that are in the map.
- First posted on arXiv 19 Aug 2026.

## Builds on (most cited first, 21 of 21)

- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings](https://mlsanity.com/paper/1909.00512/) (2019)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [Momentum-Based Variance Reduction in Non-Convex SGD](https://mlsanity.com/paper/1905.10018/) (2019)
- [An Investigation into Neural Net Optimization via Hessian Eigenvalue Density](https://mlsanity.com/paper/1901.10159/) (2019)
- [Empirical Analysis of the Hessian of Over-Parametrized Neural Networks](https://mlsanity.com/paper/1706.04454/) (2017)
- [All-but-the-Top: Simple and Effective Postprocessing for Word Representations](https://mlsanity.com/paper/1702.01417/) (2017)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Nested Learning: The Illusion of Deep Learning Architectures](https://mlsanity.com/paper/2512.24695/) (2025)
- [Measurements of Three-Level Hierarchical Structure in the Outliers in the Spectrum of Deepnet Hessians](https://mlsanity.com/paper/1901.08244/) (2019)
- [It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization](https://mlsanity.com/paper/2504.13173/) (2025)
- [Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales](https://mlsanity.com/paper/2512.05620/) (2025)
- [The Newton-Muon Optimizer](https://mlsanity.com/paper/2604.01472/) (2026)
- [Momentum-based variance-reduced proximal stochastic gradient method for composite nonconvex stochastic optimization](https://mlsanity.com/paper/2006.00425/) (2020)
- [Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss](https://mlsanity.com/paper/2606.06418/) (2026)

## Abstract (arXiv)

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.39 ± 4.32% fewer steps at 67M and 22.12 ± 0.80% 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.