On Surprising Effectiveness of Masking Updates in Adaptive Optimizers
Abstract
Training large language models (LLMs) relies almost exclusively on dense adaptive optimizers with increasingly sophisticated preconditioners.
We overturn this prevailing paradigm by showing that randomly masking parameter updates can be highly effective for LLM pre-training, with a masked variant of RMSProp consistently outperforming recent state-of-the-art optimizers.
We analyze this behavior and show that the randomly update 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 masking probabilities using momentum-gradient alignment.
Experiments on LLM pre-training demonstrate that Magma is a simple drop-in replacement for adaptive optimizers providing consistent gains with negligible computational overhead.
We overturn this prevailing paradigm by showing that randomly masking parameter updates can be highly effective for LLM pre-training, with a masked variant of RMSProp consistently outperforming recent state-of-the-art optimizers.
We analyze this behavior and show that the randomly update 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 masking probabilities using momentum-gradient alignment.
Experiments on LLM pre-training demonstrate that Magma is a simple drop-in replacement for adaptive optimizers providing consistent gains with negligible computational overhead.
