On Surprising Effectiveness of Masking Updates in Adaptive Optimizers

Abstract: Training 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.

• The paper demonstrates that random masking of parameter updates serves as a geometric regularizer, leading to improved training stability and reduced perplexity in LLM pre-training.
• It introduces Momentum-Aligned Gradient Masking (Magma), which adaptively scales block updates based on momentum cosine similarity to suppress noise-driven errors.
• Empirical results across transformer and MoE architectures show that Magma outperforms state-of-the-art optimizers, achieving significant performance gains on benchmarks.

Surprising Effectiveness of Masking Updates in Adaptive Optimizers: A Technical Essay

Introduction

The dominance of dense adaptive optimizers, such as Adam, in LLM training is primarily justified by the efficient utilization of dense gradients provided by backpropagation. However, this paradigm leaves little room for alternative strategies based on sparse or structured updates, despite substantial evidence that optimization in highly nonconvex, ill-conditioned landscapes could benefit from more regularized, noise-injected update protocols. The paper "On Surprising Effectiveness of Masking Updates in Adaptive Optimizers" (2602.15322) challenges this orthodoxy by demonstrating that the random masking of parameter updates—specifically via block-wise stochastic masking—can yield significant improvements in LLM pre-training, outperforming recent state-of-the-art dense optimizers both in terms of training stability and downstream generalization.

Random Masking as a Geometric Regularizer

The core observation is that introducing stochasticity through random block-wise masking of updates, even while retaining unbiasedness of the overall parameter shift, results in a form of implicit geometric regularization. By extending the update rule of a baseline optimizer (e.g., RMSProp) to include Bernoulli masking on parameter blocks, the method—SkipUpdate—discourages progress in directions exhibiting sharp local curvature (i.e., high diagonal Hessian entries), biasing the optimization trajectory towards flatter regions of the loss surface. This is substantiated by a rigorous Taylor expansion and expectation analysis, revealing a curvature-weighted penalty term induced by masking.

Figure 1: Pre-training performance on C4 across model scales; SkipUpdate yields substantial gains over dense optimizers despite skipping half of updates.

Empirically, SkipUpdate yields notable perplexity improvements over base optimizers, contrary to conventional expectations rooted in worst-case convergence theory. This effect is especially pronounced in transformer architectures, whose Hessians exhibit strong block-diagonal structure—suggesting that block-wise masking efficiently targets relevant geometric features for regularization.

Momentum-Aligned Gradient Masking (Magma)

Building on the advantages of uniform stochastic masking, the paper proposes a refined mechanism: Momentum-aligned gradient masking (Magma). Here, the magnitude of each block's update is adaptively scaled by the exponential moving average of the block-wise cosine similarity between momentum (first-moment estimate) and stochastic gradient, damped by a softmax temperature. This probabilistically prioritizes updates where the stochastic gradient aligns with historical momentum, effectively suppressing noise-driven, misaligned updates that undermine optimizer stability.

Magma thus augments geometric regularization with alignment-aware modulation, maximizing the benefits of stochastic masking without additional memory or computational overhead, and operates as a simple drop-in wrapper for modern adaptive optimizers.

Empirical Results: LLM Pre-Training and Mixture-of-Experts

The efficacy of Magma is validated extensively across Llama 2 pre-training on the C4 benchmark (60M–1B parameters) and sparse mixture-of-experts (MoE) transformers (Nano MoE on OpenWebText). In all cases, Magma achieves lower validation perplexities than both traditional optimizers (e.g., Adam, RMSProp, LaProp) and recent advanced schemes (Muon, SOAP, Adam+SGG, C-Adam), with performance margins increasing as model scale grows:

• With 1B parameters, Magma reduces perplexity by >19% over Adam and >9% over Muon.
• In MoE settings, Magma attains the best final performance and is robust to the complex, load-imbalanced dynamics inherent to sparsely-activated architectures.

  Figure 2: Optimization trajectories of pre-training the Nano MoE model on OpenWebText—Magma attains superior final performance and stability.

Ablation studies further confirm that masking granularity (element-wise, row, column, block), block selection (targeting attention/MLP only vs. global), and masking schedule all have secondary impact compared to the dominant role played by geometric regularization and alignment-based modulation.

Robustness to Heavy-Tailed Noise and Ill-Conditioned Curvature

In synthetic settings simulating LLM training phenomena—such as heavy-tailed stochastic gradients and ill-conditioned quadratic losses—Magma significantly outperforms Adam and other baselines under heavy-tailed regimes, both in convergence speed and the conditioning of traversed loss regions.

Figure 3: Magma on light-tailed and heavy-tailed data distributions—outperforming Adam under heavy-tailed stochastic gradients.

On heterogeneous block-structured quadratics, Magma delivers faster convergence and reduced final loss relative to AdamW when eigenspectra are highly mixed across blocks, validating the theoretical premise that targeted masking mitigates instability from poor local conditioning.

Figure 4: Magma on homogeneous and heterogeneous quadratics; enhanced efficacy observed under heterogeneous block structure.

Theoretical Analysis of Convergence with Masked Updates

A comprehensive analysis yields nonconvex convergence guarantees for Magma, revealing that stochastic masking reduces the effective smoothness constants ($\tilde{L}_t^{(b)}$) and curvature-weighted noise terms that dictate stepsize selection and noise floor in stationary optimization regimes. Alignment-scaling ($s_t^{(b)}$) selectively attenuates high-variance, high-curvature blocks, enlarging the stability envelope for constant-step optimization and enabling use of larger learning rates with greater robustness to hyperparameter selection.

Figure 5: Sensitivity analysis of learning rate on evaluation perplexity—Adam+Magma maintains stability across a much wider hyperparameter range than Adam or C-Adam.

Literature Context and Implications

Masking parameter updates shares conceptual lineage with stochastic perturbation methods including Dropout, adaptive noise injection, and meta-learning regularizers. However, Magma’s innovation lies in the amalgamation of curvature-adaptive regularization (without explicit Hessian computation) and alignment-aware update selection, enabling lightweight, structured noise injection that is uniquely compatible with the optimization geometry of transformers and foundation models. Importantly, unlike prior sign-based or deterministic masking enhancers (e.g., Cautious Optimizer), Magma exploits stochastic masking to achieve robust regularization without introducing optimizer bias or unduly restricting descent directions.

Practical and Theoretical Implications

Practically, Magma offers a trivial integration path into existing large-scale optimization pipelines, producing immediate empirical gains—most pronounced in settings with pronounced loss surface heterogeneity and intrinsic stochasticity (LLM pre-training, MoE architectures, substantial model scale). Theoretically, it calls into question standard assumptions regarding the strict necessity of dense gradient application in the presence of dense backpropagation outputs and suggests the existence of a broad class of structured stochastic update protocols beneficial for optimization stability, generalization, and resource efficiency.

Looking forward, open questions remain around unbiased masking with strong stability, further integration with advanced preconditioners, and potential for de-biasing or meta-learned masking policies tailored to architecture-specific landscape features.

Conclusion

By introducing and theoretically analyzing random and momentum-aligned masking of parameter updates, the paper provides compelling evidence that stochastic update masking, as embodied in Magma, yields substantial benefits in modern neural optimization regimes—most notably in LLM pre-training. This work provokes new directions for the design of optimizers that capitalize on structured stochasticity for robust, stable, and efficient large-scale training, advocating for a re-examination of dense update dogma in favor of geometry-aware, noise-regularized methodologies.