# Adaptive Memory Momentum via a Model-Based Framework for Deep Learning Optimization

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 83 references that are in the map.
- First posted on arXiv 6 Oct 2025.

## Builds on (most cited first, 30 of 48)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Cyclical Learning Rates for Training Neural Networks](https://mlsanity.com/paper/1506.01186/) (2015)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Stochastic model-based minimization of weakly convex functions](https://mlsanity.com/paper/1803.06523/) (2018)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [An Improved Analysis of Stochastic Gradient Descent with Momentum](https://mlsanity.com/paper/2007.07989/) (2020)
- [Learning-Rate-Free Learning by D-Adaptation](https://mlsanity.com/paper/2301.07733/) (2023)
- [DoG is SGD's Best Friend: A Parameter-Free Dynamic Step Size Schedule](https://mlsanity.com/paper/2302.12022/) (2023)
- [Adaptive Gradient Descent without Descent](https://mlsanity.com/paper/1910.09529/) (2019)
- [MARS: Unleashing the Power of Variance Reduction for Training Large Models](https://mlsanity.com/paper/2411.10438/) (2024)
- [Stochastic (Approximate) Proximal Point Methods: Convergence, Optimality, and Adaptivity](https://mlsanity.com/paper/1810.05633/) (2018)
- [A Unified Analysis of Stochastic Momentum Methods for Deep Learning](https://mlsanity.com/paper/1808.10396/) (2018)
- [The AdEMAMix Optimizer: Better, Faster, Older](https://mlsanity.com/paper/2409.03137/) (2024)

## Built on by (most cited first, 1 of 1)

- [Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets](https://mlsanity.com/paper/2603.16731/) (2026)

## Abstract (arXiv)

The vast majority of modern deep learning models are trained with momentum-based first-order optimizers. The momentum term governs the optimizer's memory by determining how much each past gradient contributes to the current convergence direction. Fundamental momentum methods, such as Nesterov Accelerated Gradient and the Heavy Ball method, as well as more recent optimizers such as AdamW and Lion, all rely on the momentum coefficient that is customarily set to β = 0.9 and kept constant during model training, a strategy widely used by practitioners, yet suboptimal. In this paper, we introduce an textitadaptive memory mechanism that replaces constant momentum with a dynamic momentum coefficient that is adjusted online during optimization. We derive our method by approximating the objective function using two planes: one derived from the gradient at the current iterate and the other obtained from the accumulated memory of the past gradients. To the best of our knowledge, such a proximal framework was never used for momentum-based optimization. Our proposed approach is novel, extremely simple to use, and does not rely on extra assumptions or hyperparameter tuning. We implement adaptive memory variants of both SGD and AdamW across a wide range of learning tasks, from simple convex problems to large-scale deep learning scenarios, demonstrating that our approach can outperform standard SGD and Adam with hand-tuned momentum coefficients. Finally, our work opens doors for new ways of inducing adaptivity in optimization.