On Surprising Effectiveness of Masking Updates in Adaptive Optimizers
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
| 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 | 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 |
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 |
| 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 |
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