Search papers, labs, and topics across Lattice.
The paper investigates the effectiveness of randomly masking parameter updates in adaptive optimizers for large language model (LLM) training, finding that it surprisingly outperforms state-of-the-art optimizers. They attribute this to a curvature-dependent geometric regularization effect induced by the masking. Based on this, they propose Momentum-aligned gradient masking (Magma), a novel masking technique that modulates updates using momentum-gradient alignment, achieving significant perplexity reductions compared to Adam and Muon in LLM pre-training.
Randomly masking parameter updates in RMSProp delivers state-of-the-art LLM training performance, revealing a surprisingly effective form of geometric regularization.
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.
Google Research