AdamPM: Memory-Efficient LLM Optimization via Moment Pooling

Xiaowei Ji; Lishan Xu; Jianing An; Jie Luo; Wenjun Wu; Lei Huang

cs.LGOct 17, 2026

AdamPM: Memory-Efficient LLM Optimization via Moment Pooling

Authors: Xiaowei Ji, Lishan Xu, Jianing An, Jie Luo, Wenjun Wu, Lei Huang

Organizations: SKLCCSE, School of Artificial Intelligence, Beihang University, Beijing, China · Hangzhou International Innovation Institute, Beihang University, Hangzhou, China

Abstract

Adam is widely used for training large language models (LLMs) owing to its strong optimization performance and training stability, but its element-wise estimation for the first and second moments incur substantial memory overhead. Although numerous memory-efficient optimizers have been proposed, many incur additional computational overhead or require additional tuning of hyperparameters. Consequently, Adam remains a dominant and reliable choice for LLM pre-training. We propose AdamPM (Adam with Pooled Moments), a simple and memory-efficient variant of Adam that groups adjacent moment entries along the row or column dimension and represents each group by a single shared mean. AdamPM can well inherit Adam's hyperparameter settings and is readily integrated into existing Adam optimizer implementations, with less memory overhead. We systematically investigate pooling of both first and second moment estimates and identify second moment pooling as the effective design. With a group size of $s$ , AdamPM reduces second moment memory to $1/s$ of the original requirement. Extensive experiments on language model pretraining, fine-tuning and multimodal learning demonstrate that AdamPM maintains model performance and training speed while substantially reducing optimizer-state memory. Our code is available at https://github.com/jixiaowei95/AdamPM.

Figures & tables

Figure 1: Evolution of effective ranks and coefficients of variation during LLaMA 130M pre-training. (a) Effective ranks of the gradients, first and second moments for embed and key . (b) Coefficients of variation of the first and second moments for value and gate . (c) Coefficients of variation of the preconditioner for gate .

Model Size	60M	130M	350M	1.3B
Training Tokens	1.1B	2.2B	6.4B	13.1B
Adam	30.30 (0.325G)	23.23 (0.749G)	17.79 (2.05G)	14.83 (7.48G)
AdamPM-2	30.31 (0.270G)	23.47 (0.624G)	17.69 (1.71G)	14.70 (6.24G)
AdamPM-4	30.25 (0.243G)	22.99 (0.562G)	17.64 (1.54G)	14.66 (5.61G)
AdamPM-8	30.61 (0.230G)	22.87 (0.531G)	17.75 (1.46G)	14.73 (5.30G)
AdamPM-16	30.57 (0.223G)	23.17 (0.515G)	17.86 (1.41G)	14.93 (5.14G)

Table 1: Comparison of AdamPM, Adam, and representative memory-efficient optimizers for pre-training LLaMA models of various sizes on the C4 dataset. Validation perplexity is reported, with the estimated BF16 memory cost of model parameters and optimizer states. * denotes that the results are taken directly from the original papers, as they use the same pre-training setup, model scales, and token budgets.

Figure 2: Learning rate sweeps for LLaMA 60M and 130M pre-training on C4

Model Size	60M	130M	350M
Adam	30.30	23.23	17.79
AdamPM-3	30.58	23.36	17.63
AdamPM-5	30.54	23.09	17.59
AdamPM-col-2	30.45	23.40	17.56
AdamPM-col-3	30.36	23.34	17.53
AdamPM-M-2	59.99	206.35	49.26

Table 2: Ablation study of pooling strategy.

Figure 3: Perplexity comparison between AdamPM and vanilla Adam during LLaMA 1.3B pre-training.

Figure 4: Comparison of the preconditioners

\frac{1}{\sqrt{\widehat{V}_{t}}+\epsilon}

(denoted by pre ) and preconditioned updates

\frac{\widehat{M}_{t}}{\sqrt{\widehat{V}_{t}}+\epsilon}

(denoted by upd ) before and after applying second moment pooling during LLaMA 130M pre-training. (a) Relative errors of the preconditioners and preconditioned updates for query and the MLP down projection matrix down , together with the averages across all weight matrices in the network. (b) Corresponding cosine similarities. (c) Average cosine similarity of preconditioned updates across all layers under different group sizes

s

. (d) Average cosine similarity across all layers for different model sizes.

	Memory	CoLA	STS-B	MRPC	RTE	SST2	MNLI	QNLI	QQP	Avg
Adam	748M	64.09	91.18	93.21	79.78	94.49	87.30	92.87	92.06	86.87
Adam-mini	499M	54.68	90.46	86.32	52.70	94.49	86.01	92.34	91.52	81.06
AdamPM-2	623M	64.29	91.16	92.28	80.14	94.83	87.28	93.02	91.97	86.87
AdamPM-4	561M	63.26	91.16	93.52	79.06	94.83	87.32	93.15	91.99	86.79
AdamPM-8	529M	63.87	91.20	93.28	82.31	95.06	87.53	92.87	91.86	87.25
AdamPM-16	514M	62.92	91.20	93.05	79.78	95.29	87.36	92.75	91.87	86.78

Table 3: Evaluating AdamPM for memory-efficient fine-tuning on GLUE benchmark using pretrained RoBERTa-base. We report the overall (matched and mismatched) accuracy for MNLI, Matthew’s correlation coefficient for CoLA, Pearson correlation for STS-B, F1-score for MRPC, and accuracy for the remaining tasks.

	Memory	VQAv2	MM-Vet	GQA	SQA	TextVQA	POPE	MME	MMMU
Adam	4.38G	68.95	20.6	51.38	54.02	42.88	85.22	1330.60	29.80
AdamPM-V-2	3.91G	69.27	21.6	51.66	53.88	42.34	85.21	1329.66	30.30
AdamPM-V-4	3.68G	69.55	19.5	52.06	54.14	42.41	85.58	1336.84	31.10
AdamPM-V-8	3.57G	69.70	19.6	52.23	53.95	42.75	85.61	1320.59	31.30
AdamPM-V-16	3.51G	69.84	19.2	52.61	54.07	42.41	85.82	1313.27	31.10

Table 4: Performance comparison of AdamPM with different group sizes on TinyLLaVA-1B.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Model Size	60M	130M	350M	1.3B
Hidden dimension	512	768	1024	2048
Intermediate dimension	1376	2048	2736	5461
Attention Heads	8	12	16	32
Decoder layers	8	12	24	24
Learning Rate	$4\times 10^{-3}$	$2\times 10^{-3}$	$6\times 10^{-4}$	$3\times 10^{-4}$
$\beta_{1}$	$0.9$

Appendix

Table A1: The hyperparameters for LLaMA pre-training.

Optimizer	Memory Cost
Adam	38.11 G
AdamPM-2	36.73 G
AdamPM-4	35.96 G
AdamPM-8	35.43 G
AdamPM-16	35.19 G
AdamPM-inf	35.14 G

Appendix

Table A2: Measured memory usage per GPU during LLaMA-1.3B pre-training.

Hyperparameters	MNLI	SST-2	MRPC	CoLA	QNLI	QQP	RTE	STS-B
Batch size	64	32	16	32	64	64	16	32
Epochs	30	30	30	30	30	30	30	30
$lr$	2e-5	1e-5	2e-5	2e-5	1e-5	2e-5	2e-5	2e-5
$\beta_{1}$	$0.9$
$\beta_{2}$	$0.95$
Max seq. len.	$512$

Appendix

Table A3: Hyperparameters for fine-tuning RoBERTa-base on GLUE.

Efficient Language Model Training

Efficient language model training studies how to reduce the data, compute, memory, time, energy, and cost needed to build capable language models.

See all 157 papers

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers
May 9, 2026Aditya RanganathDeep Learning Optimization LLM Input Optimization
Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA
Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching
Oct 3, 2026Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi +3Nyström Method Second-Order Optimization
Mila & DIRO, Université de Montréal · Critical ML, University of Waterloo · Independent Researcher +2
AdaMeZO: Adam-style Zeroth-Order Optimizer for LLM Fine-tuning Without Maintaining the Moments
May 1, 2026Zhijie Cai, Haolong Chen, Guangxu ZhuMemory-Efficient Fine-Tuning Zeroth-Order Optimization
Shenzhen Research Institute of Big Data · The Chinese University of Hong Kong-Shenzhen · Shenzhen Loop Area Institute

May 9, 2026J/K move · Enter open · S save

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers

Aditya Ranganath

Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA

Deep Learning Optimization LLM Input Optimization Memory-Efficient Optimization Matrix Optimization LLM Training

Open paper AskPDF arXiv

Training large language models requires optimization algorithms that are not only statistically effective, but also computationally and memory efficient at extreme scale. Although Adam remains the dominant optimizer for large-scale language-model pretraining and fine-tuning, recent work has revisited nearly every component of the optimization stack: adaptive moment estimation, decoupled weight decay, memory footprint, curvature approximation, sign-based updates, large-batch stability, low-rank gradient structure, and matrix-wise orthogonalized updates. This survey reviews optimizer design for large language models through a systems-and-optimization lens. We organize the literature into classical first-order optimizers, adaptive optimizers, memory-efficient variants, second-order and curvature-aware methods, sign-based and discovered optimizers, low-rank and projection-based methods, and matrix-based optimizers such as Muon. We also discuss benchmarking methodology, including hyperparameter fairness, scale dependence, wall-clock efficiency, token efficiency, memory overhead, and downstream evaluation. We argue that optimizer research for LLMs is entering a new phase: moving from single-algorithm speedup claims toward rigorous, scale-aware comparisons that jointly evaluate convergence, stability, memory, and implementation complexity.

Abstract

Efficient Language Model Training

Explore similar work

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers

Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching

AdaMeZO: Adam-style Zeroth-Order Optimizer for LLM Fine-tuning Without Maintaining the Moments