# MoMo: Momentum Models for Adaptive Learning Rates

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

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 54 references that are in the map.
- First posted on arXiv 12 May 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Online Learning: A Modern Introduction Using Convex Optimization](https://mlsanity.com/paper/1912.13213/) (2019)
- [AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://mlsanity.com/paper/2010.07468/) (2020)
- [Adaptive Gradient Methods with Dynamic Bound of Learning Rate](https://mlsanity.com/paper/1902.09843/) (2019)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [Stochastic model-based minimization of weakly convex functions](https://mlsanity.com/paper/1803.06523/) (2018)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [Painless Stochastic Gradient: Interpolation, Line-Search, and Convergence Rates](https://mlsanity.com/paper/1905.09997/) (2019)
- [Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers](https://mlsanity.com/paper/2007.01547/) (2020)
- [Three Mechanisms of Weight Decay Regularization](https://mlsanity.com/paper/1810.12281/) (2018)
- [Learning-Rate-Free Learning by D-Adaptation](https://mlsanity.com/paper/2301.07733/) (2023)
- [The Cost of Training NLP Models: A Concise Overview](https://mlsanity.com/paper/2004.08900/) (2020)
- [SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation](https://mlsanity.com/paper/2006.10311/) (2020)
- [Stochastic (Approximate) Proximal Point Methods: Convergence, Optimality, and Adaptivity](https://mlsanity.com/paper/1810.05633/) (2018)
- [Understanding AdamW through Proximal Methods and Scale-Freeness](https://mlsanity.com/paper/2202.00089/) (2022)
- [Training Neural Networks for and by Interpolation](https://mlsanity.com/paper/1906.05661/) (2019)
- [Dynamics of SGD with Stochastic Polyak Stepsizes: Truly Adaptive Variants and Convergence to Exact Solution](https://mlsanity.com/paper/2205.04583/) (2022)
- [A Stochastic Proximal Polyak Step Size](https://mlsanity.com/paper/2301.04935/) (2023)
- [Accelerated, Optimal, and Parallel: Some Results on Model-Based Stochastic Optimization](https://mlsanity.com/paper/2101.02696/) (2021)
- [The Power of Factorial Powers: New Parameter settings for (Stochastic) Optimization](https://mlsanity.com/paper/2006.01244/) (2020)

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

- [Implicit Bias and Fast Convergence Rates for Self-attention](https://mlsanity.com/paper/2402.05738/) (2024)
- [An Exploration of Non-Euclidean Gradient Descent: Muon and its Many Variants](https://mlsanity.com/paper/2510.09827/) (2025)
- [Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance](https://mlsanity.com/paper/2406.04142/) (2024)
- [Stochastic Gradient Descent with Preconditioned Polyak Step-size](https://mlsanity.com/paper/2310.02093/) (2023)
- [New Perspectives on the Polyak Stepsize: Surrogate Functions and Negative Results](https://mlsanity.com/paper/2505.20219/) (2025)
- [SANIA: Polyak-type Optimization Framework Leads to Scale Invariant Stochastic Algorithms](https://mlsanity.com/paper/2312.17369/) (2023)
- [Batch size invariant Adam](https://mlsanity.com/paper/2402.18824/) (2024)
- [Analysis of an Idealized Stochastic Polyak Method and its Application to Black-Box Model Distillation](https://mlsanity.com/paper/2504.01898/) (2025)
- [Algorithm- and Data-Dependent Generalization Bounds for Diffusion Models](https://mlsanity.com/paper/2506.03849/) (2025)
- [Enhancing Policy Gradient with the Polyak Step-Size Adaption](https://mlsanity.com/paper/2404.07525/) (2024)
- [How far away are truly hyperparameter-free learning algorithms?](https://mlsanity.com/paper/2505.24005/) (2025)
- [Taking the Road Less Scheduled with Adaptive Polyak Steps](https://mlsanity.com/paper/2511.07767/) (2025)
- [Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients](https://mlsanity.com/paper/2512.02342/) (2025)
- [Gradient Clipping Beyond Vector Norms: A Spectral Approach for Matrix-Valued Parameters](https://mlsanity.com/paper/2605.11838/) (2026)
- [Adaptive Memory Momentum via a Model-Based Framework for Deep Learning Optimization](https://mlsanity.com/paper/2510.04988/) (2025)
- [Adaptive Sharpness-Aware Minimization with a Polyak-type Step size: A Theory-Grounded Scheduler](https://mlsanity.com/paper/2606.01827/) (2026)
- [Regularized PolyKervNets: Optimizing Expressiveness and Efficiency for Private Inference in Deep Neural Networks](https://mlsanity.com/paper/2312.15229/) (2023)
- [Learning Developmental Age from 3D Infant Kinetics Using Adaptive Graph Neural Networks](https://mlsanity.com/paper/2402.14400/) (2024)
- [RDM: Recurrent Diffusion Model for Human Motion Generation](https://mlsanity.com/paper/2406.07169/) (2024)
- [Step-Size Stability in Stochastic Optimization: A Theoretical Perspective](https://mlsanity.com/paper/2602.09842/) (2026)
- [Byzantine-Robust Optimization under (L0, L1)-Smoothness](https://mlsanity.com/paper/2603.12512/) (2026)
- [Zero-order Parameter-free Optimization for LMO-based Methods: Novel Approach for Efficient Fine-tuning](https://mlsanity.com/paper/2606.14970/) (2026)
- [Optimizers for Diffusion Models: A Controlled Benchmark](https://mlsanity.com/paper/2609.23055/) (2026)
- [Why Clipping Matters in AdaGrad? Toward a High-Probability Theory under Generalized Smoothness](https://mlsanity.com/paper/2609.30276/) (2026)

## Abstract (arXiv)

Training a modern machine learning architecture on a new task requires extensive learning-rate tuning, which comes at a high computational cost. Here we develop new Polyak-type adaptive learning rates that can be used on top of any momentum method, and require less tuning to perform well. We first develop MoMo, a Momentum Model based adaptive learning rate for SGD-M (stochastic gradient descent with momentum). MoMo uses momentum estimates of the losses and gradients sampled at each iteration to build a model of the loss function. Our model makes use of any known lower bound of the loss function by using truncation, e.g. most losses are lower-bounded by zero. The model is then approximately minimized at each iteration to compute the next step. We show how MoMo can be used in combination with any momentum-based method, and showcase this by developing MoMo-Adam, which is Adam with our new model-based adaptive learning rate. We show that MoMo attains a O(1/√K) convergence rate for convex problems with interpolation, needing knowledge of no problem-specific quantities other than the optimal value. Additionally, for losses with unknown lower bounds, we develop on-the-fly estimates of a lower bound, that are incorporated in our model. We show that MoMo and MoMo-Adam improve over SGD-M and Adam in terms of robustness to hyperparameter tuning for training image classifiers on MNIST, CIFAR, and Imagenet, for recommender systems on Criteo, for a transformer model on the translation task IWSLT14, and for a diffusion model.