On Surprising Effectiveness of Masking Updates in Adaptive Optimizers

Taejong Joo; Wenhan Xia; Cheolmin Kim; Ming Zhang; Eugene Ie

cs.LGFeb 17, 2026

On Surprising Effectiveness of Masking Updates in Adaptive Optimizers

Authors: Taejong Joo, Wenhan Xia, Cheolmin Kim, Ming Zhang, Eugene Ie

Organizations: Northwestern University · Work done while the author was a Student Researcher at Google. · Google

Abstract

Training large language models (LLMs) relies almost exclusively on dense adaptive optimizers with increasingly sophisticated preconditioners. We challenge this by showing that randomly masking parameter updates can be highly effective, with a masked variant of RMSProp consistently outperforming recent state-of-the-art optimizers. Our analysis reveals that the random masking induces a curvature-dependent geometric regularization that smooths the optimization trajectory. Motivated by this finding, we introduce Momentum-aligned gradient masking (Magma), which modulates the masked updates using momentum-gradient alignment. Extensive LLM pre-training experiments show that Magma is a simple drop-in replacement for adaptive optimizers with consistent gains and negligible computational overhead. Notably, for the 1B model size, Magma reduces perplexity by over 19% and 9% compared to Adam and Muon, respectively.

Figures & tables

Figure 1: Pre-training performance on C4 across model scales. Despite discarding half of updates, SkipUpdate yields substantial improvements over state-of-the-art dense optimizers.

Method	60M	130M	350M	1B
Adam	30.79	24.77	18.42	16.35
C-Adam	29.70	23.59	18.58	15.92
Adam+SGG $\dagger$	30.31	22.18	17.28	14.30
Adam+Magma	29.09	22.08	16.41	13.71
LaProp	29.98	23.07	18.56	16.38
LaProp+Magma	29.05	22.16	16.37	13.82

Table 1: Llama 2 pre-training results on the C4 dataset. Validation perplexity is reported for four model scales (60M to 1B).

\dagger

denotes results from Li et al. (2025) ; remaining results are ours. RMSProp diverges for the 1B model within the learning rate search space.

Figure 2: Optimization trajectories of pre-training the Nano MoE model on OpenWebText.

Panel of Figure 3 — Figure 3: Magma on light-tailed and heavy-tailed data distributions. Top: Optimization trajectories for Adam and Magma. Bottom: Robust condition number defined as the ratio between the maximum and median eigenvalues of the loss Hessian.

Figure 4: Magma on homogeneous and heterogeneous quadratics. Top: Optimization trajectories for AdamW and Magma on quadratic objectives with identical eigenspectra but different block structures. Bottom: Average gradient–momentum alignment per block.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

	Baseline	Attn Only	Attn + MLP	All
Eval Perplexity	22.64	21.92	21.65	21.94

Appendix

Table A1: Validation perplexity (

\downarrow

) for different masking components.

Method	Element	Row	Column	Block
Uniform Sampling	21.73	21.76	21.78	21.81
Damping	21.97	21.95	21.91	21.92
Uniform Sampling + Damping	21.58	21.62	21.61	21.65
Momentum-Gradient Alignment-based Sampling	21.77	21.78	21.75	21.78
Momentum-Gradient Alignment-based Sampling + Damping	21.63	21.60	21.61	21.65

Appendix

Table A2: Validation perplexity (

\downarrow

) for different masking granularities using Uniform vs. Momentum-Gradient Alignment-based sampling schemes (with and without damping). The baseline (RMSProp) excluding both sampling and damping achieves the perplexity of 22.64.

Figure A3: Sensitivity analysis of learning rate on evaluation perplexity. Unlike Adam and C-Adam, which exhibit narrow optimal windows, Adam+Magma maintains consistent performance, demonstrating stability across a significantly wider hyperparameter range.

Deep Learning Optimization

Deep learning optimization develops algorithms and training strategies that help neural networks converge, generalize, and use compute efficiently.

See all 199 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
Omni-Masked Gradient Descent: Memory-Efficient Optimization via Mask Traversal with Improved Convergence
Mar 6, 2026Hui Yang, Tao Ren, Jinyang Jiang +2Gradient Descent Memory-Efficient Training
1Guanghua School of Management, Peking University · 2Wangxuan Institute of Computer Technology, Peking University.
Boosting Large Language Models with Mask Fine-Tuning
Mar 27, 2025Mingyuan Zhang, Yue Bai, Huan Wang +4LLM Fine-Tuning
College of Engineering, Northeastern University · Khoury College of Computer Science, Northeastern University

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

Deep Learning Optimization

Explore similar work

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

Omni-Masked Gradient Descent: Memory-Efficient Optimization via Mask Traversal with Improved Convergence

Boosting Large Language Models with Mask Fine-Tuning