Adaptive Memory Momentum via a Model-Based Framework for Deep Learning Optimization
Kristi Topollai, Anna Choromanska
TL;DR
This paper tackles the suboptimality of fixed momentum in first-order optimizers like HB, NAG, and AdamW by introducing Adaptive Memory (AM), a principled, model-based framework that makes the momentum coefficient $\beta_t$ adaptive and online. By approximating the loss with a two-plane proximal model—one plane from the current gradient and another from the accumulated descent direction—the authors derive a closed-form update for $\beta_t$ and show that the standard momentum step becomes $d_{t+1}=\frac{\beta_t+\lambda}{1+\lambda}d_t+\frac{1-\beta_t}{1+\lambda}g_t$, with $x_{t+1}=x_t-\eta d_{t+1}$, while enabling easy integration with preconditioning and decoupled weight decay to produce AM-AdamW. Theoretical guarantees in convex and non-convex settings show the method preserves the usual $O(1/\sqrt{T})$ convergence rate under standard assumptions, despite potential correlations introduced by adaptivity. Empirically, AM variants outperform fixed-momentum baselines across convex problems, image classification, and large-language-model pretraining, with negligible computational overhead and promising warmup-free stability for large models. The work thus provides a practical, tuning-free path to momentum adaptivity that can simplify training pipelines and inspire further adaptive strategies in optimization.
Abstract
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 \textit{adaptive 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.