# AdaPM: a Partial Momentum Algorithm for LLM Training

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 41 references that are in the map.
- First posted on arXiv 10 Oct 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [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)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [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)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [UltraFeedback: Boosting Language Models with Scaled AI Feedback](https://mlsanity.com/paper/2310.01377/) (2023)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [Gradient Descent Happens in a Tiny Subspace](https://mlsanity.com/paper/1812.04754/) (2018)
- [ReLoRA: High-Rank Training Through Low-Rank Updates](https://mlsanity.com/paper/2307.05695/) (2023)
- [ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models](https://mlsanity.com/paper/2310.10505/) (2023)
- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be](https://mlsanity.com/paper/2304.13960/) (2023)
- [A Spectral Condition for Feature Learning](https://mlsanity.com/paper/2310.17813/) (2023)
- [Continual Learning in Low-rank Orthogonal Subspaces](https://mlsanity.com/paper/2010.11635/) (2020)
- [CAME: Confidence-guided Adaptive Memory Efficient Optimization](https://mlsanity.com/paper/2307.02047/) (2023)
- [Subspace Optimization for Large Language Models with Convergence Guarantees](https://mlsanity.com/paper/2410.11289/) (2024)
- [Neural Tangent Kernel: A Survey](https://mlsanity.com/paper/2208.13614/) (2022)
- [Low-rank Gradient Approximation For Memory-Efficient On-device Training of Deep Neural Network](https://mlsanity.com/paper/2001.08885/) (2020)
- [When and Why Momentum Accelerates SGD:An Empirical Study](https://mlsanity.com/paper/2306.09000/) (2023)
- [Risk Bounds of Accelerated SGD for Overparameterized Linear Regression](https://mlsanity.com/paper/2311.14222/) (2023)
- [Optimal Algorithms in Linear Regression under Covariate Shift: On the Importance of Precondition](https://mlsanity.com/paper/2502.09047/) (2025)

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

- [Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation](https://mlsanity.com/paper/2602.24283/) (2026)
- [PRAC: Principal-Random Subspace for LLM Activation Compression and Memory-Efficient Training](https://mlsanity.com/paper/2602.23111/) (2026)
- [Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling](https://mlsanity.com/paper/2606.00888/) (2026)

## Abstract (arXiv)

In the training of large language models, momentum is widely used and often demonstrated to achieve significant acceleration. However, storing momentum typically presents memory challenges. In this paper, we propose AdaPM, an adaptive training strategy that leverages partial momentum to implement a memory-efficient optimizer. To this end, AdaPM utilizes a non-uniform momentum design: for most blocks, full momentum is not necessary to preserve the performance of the optimization. In the momentum design of AdaPM, to mitigate the bias and performance loss caused by partial momentum, we enhance the partial momentum by a bias correction technique. Empirically, we verify that our approach reduces memory by over 90% in momentum while maintaining both efficiency and performance for pretraining various language models ranging from 60M to 1.5B, as well as for supervised fine-tuning and RLHF. AdaPM can further reduce memory by up to 95% in optimizer states by combining the memory-efficient technique on the second-order statistic, saving over 30% GPU hours for pretraining GPT-2 1.5B.