PaperlayerView exact revision on arXiv

Source-linked AI summary

On Surprising Effectiveness of Masking Updates in Adaptive Optimizers

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

arXiv:2602.15322v1cs.LGcs.AI

TL;DR

Although dense adaptive optimizers dominate LLM training, this paper studies randomly masked updates and introduces Magma, which aligns masking with momentum. Magma consistently outperforms adaptive optimizers and SkipUpdate with negligible overhead, challenging the assumption that dense updates are optimal.

  • Problem

    Dense adaptive optimizers dominate large-scale LLM training, while sparse update strategies remain rarely used despite their performance on nonsmooth optimization.

  • Method

    Magma randomly masks parameter updates and modulates surviving updates using cosine similarity between stochastic gradients and first-moment estimates.

  • Results

    Magma consistently outperforms adaptive optimizers and SkipUpdate, while introducing negligible computational overhead.

  • Takeaways & Limitations

    The results challenge the assumption that dense updates are inherently optimal for backpropagation-based neural network training.

  • Takeaways & Limitations

    Magma’s benefits may be specific to transformer-like loss geometry, as it does not improve over AdamW on ResNet-50 CIFAR-10.

Abstract

from arXiv · show

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.

1. Introduction

The section challenges dense adaptive optimization by showing that randomly masking parameter updates can improve training, despite discarding updates and retaining backpropagation cost. It explains this effect through curvature-dependent geometric regularization and introduces Magma, which aligns masked updates with momentum to improve stability and generalization.

  • Motivation and contribution: Randomly masking gradient updates in a Bernoulli block-wise RMSProp variant substantially improves optimization performance.Masked blocks skip their parameter updates for that iteration.
  • Empirical finding: Despite discarding half of the updates, SkipUpdate consistently outperforms dense optimizers, including Muon.The method retains unchanged gradient-computation cost because backpropagation remains dense.
  • Geometric explanation: Block-wise masking induces curvature-dependent geometric regularization that penalizes updates aligned with sharp loss directions and biases trajectories toward flatter regions.The regularization emerges implicitly from stochastic update noise rather than explicit curvature computation.
  • Magma: Magma modulates masked updates using cosine similarity between stochastic gradients and first-moment estimates, suppressing momentum-inconsistent directions.This prioritizes updates consistent with the accumulated gradient direction.
  • Magma: Across diverse pre-training scenarios, Magma improves training stability and generalization over state-of-the-art optimizers without additional computational cost.The method is presented as a simple optimizer wrapper for transformer loss landscapes.

2. Update Masking as a Regularization

This section presents update masking as a curvature-dependent geometric regularizer that discourages movement in sharp directions. It also motivates dense momentum updates and block-wise masking as effective, stable, and computationally favorable design choices.

  • Update masking: SkipUpdate masks block-wise parameter updates with independent Bernoulli variables while preserving the expected update through inverse-probability scaling.Each block survives with probability p, and scaling by 1/p makes the masked update unbiased.
  • Curvature regularization: Random masking induces a curvature-dependent term in expected loss decrease, acting as a geometric regularizer that penalizes updates along high-curvature directions.The regularization is weighted roughly by inverse survival probability and discourages alignment with sharp directions.
  • Curvature regularization: Block-wise masking yields principled second-order regularization toward flatter regions because transformer Hessians empirically exhibit pronounced block-diagonal structure.Under this geometry, dominant curvature interactions occur within blocks.
  • Dense momentum updates: Dense momentum updates provide a variance-reduced estimator of true momentum, producing more stable search directions and improved generalization over sparse momentum updates.Momentum states remain updated even when parameter updates are masked, unlike selected-coordinate subspace methods.
  • Masking granularity: 21.78 (column-wise), 21.73 (element-wise), and 21.81 (block-wise) perplexities under 130M Llama pre-training on C4 all substantially outperform the RMSProp baseline (22.64).The similar results across granularities motivate block-wise masking, which also enables efficient operation pruning by skipping entire blocks.

3. Momentum-Aligned Update Masking

Magma is a block-adaptive masking wrapper that uses momentum-gradient alignment to modulate updates, addressing heterogeneous optimization dynamics while remaining compatible with adaptive optimizers. Its damping improves stability without adding memory or computational overhead, though unbiased alternatives remain unstable.

  • Motivation: Transformer parameter heterogeneity motivates block-adaptive masking rather than SkipUpdate’s homogeneous masking across all blocks.The heterogeneity includes differences in Hessian spectra and gradient variances.
  • Alignment criterion: Magma uses block-wise momentum-gradient alignment to identify destabilizing updates and control the masking process.Negative alignment events are statistically abnormal fluctuations, while consistent gradient components tend to carry meaningful optimization signal.
  • Alignment criterion: Cosine similarity provides a scale-invariant alignment score suited to LLM training, where gradient norms vary substantially across parameter blocks and iterations.The score is computed from the gradient and first-moment estimate using a temperature parameter.
  • Update modulation: Magma dampens masked updates according to alignment, encouraging coherent optimization trajectories with larger alignment scores.The alignment score is maintained with an exponential moving average.
  • Implementation: Magma is a drop-in wrapper for adaptive optimizers that requires no additional memory or computational overhead.Practitioners can integrate it into existing training pipelines with minimal code changes and no additional resource requirements.
  • Limitations: Damping introduces bias but substantially improves training stability, whereas tested unbiased rescaling alternatives consistently produced unstable training.A stable yet unbiased masking scheme remains an open direction.

4. Experiments

Experiments show that Magma improves adaptive and matrix-based optimizers across Llama 2, MoE, and controlled benchmarks, especially under complex or heavy-tailed optimization conditions. Its benefits are linked to curvature-dependent geometric regularization, but do not extend uniformly to non-transformer-like loss geometry.

  • Llama 2 pre-training: Magma consistently improves base optimizers across Llama 2 pre-training scales, while RMSProp+Magma achieves the lowest perplexity from 60M to 1B.The comparison includes Adam, Adafactor, APOLLO, LaProp, RMSProp, Muon, SOAP, SGG, and C-Adam.
  • Llama 2 pre-training: 13.81 validation perplexity: Adam+Magma at 1B outperforms Adam+SGG at 14.30 and C-Adam at 15.92.This result identifies Magma as the strongest evaluated Adam enhancement at that scale.
  • Llama 2 pre-training: Magma’s relative gains increase with model size, supporting random masking as geometric regularization in increasingly irregular and nonsmooth loss landscapes.The experiments use a single robust configuration, τ=2, applied to attention and MLP layers.
  • Nano MoE: Magma improves both Adam and Muon in Nano MoE training, with Muon+Magma achieving the best overall performance and outperforming all baselines.With Adam, convergence is slower during intermediate training but superior at the end; Magma also outperforms C-Adam.
  • Heavy-tailed noise: Under heavy-tailed noise, Magma significantly outperforms Adam and attains substantially smaller robust condition numbers, indicating confinement to better-conditioned loss regions.Adam and Magma perform similarly under normal noise; the benchmark uses a simplified linear transformer and compares light-tailed and heavy-tailed regimes.
  • Curvature and transfer: Magma converges faster and reaches lower final loss than AdamW on heterogeneous Hessians, but does not improve AdamW on ResNet-50, where test accuracy is 94.46% vs. 93.82%.On homogeneous Hessians, Magma and AdamW perform comparably, with Magma slightly faster early in training.

5. Discussion

The discussion analyzes Magma with block-wise smoothness and stochastic-gradient assumptions, showing that scaled random masking trades some descent efficiency for reduced curvature-weighted noise. Its benefits arise from selectively attenuating high-curvature or high-variance blocks, widening stability and lowering the stationary error floor.

  • Analytical setup: Magma’s analysis uses constant-step SGD to isolate stochasticity and masking, with extensions to adaptive optimizers following the same descent-lemma framework.The update is θ_t+1 = θ_t − ηS_tM_t(g_t), contrasted with vanilla SGD’s θ_t+1 = θ_t − ηg_t.
  • Descent and noise tradeoff: Magma’s scaling reduces the effective smoothness constants and curvature-weighted noise contribution, while potentially reducing descent efficiency.The discussion emphasizes that scaling is not uniformly beneficial; improvements depend on which blocks are attenuated rather than average suppression alone.
  • Curvature-dependent effects: For blocks with large curvature, scaling simultaneously enlarges the admissible stepsize range and lowers the stationary error floor.The same stability improvement increases the iterations satisfying the descent surrogate, accelerating effective progress per iteration.
  • Targeted masking: Magma enlarges the stability region by selectively suppressing blocks that dominate curvature-weighted noise and smoothness constraints.This mechanism is especially relevant to ill-conditioned, heterogeneous transformer landscapes where stability is governed by a small subset of high-curvature or high-variance blocks.
  • Targeted masking: Momentum–gradient alignment makes structured scaling more effective than uniform or random masking in large-scale LLM training by targeting the blocks that limit admissible stepsizes.The discussion links this targeted reduction to the empirically observed widening of the stable learning-rate regime.

6. Literature Review

The literature review situates Magma among methods that stabilize LLM training, exploit optimization geometry, and use stochastic perturbations. It distinguishes Magma through stochastic masking that regularizes update trajectories without explicit curvature computation or extra gradient evaluations.

  • Stabilizing LLM training: Cautious Optimizer masks parameters whose gradients oppose first-moment estimates, whereas Magma adds geometric regularization that promotes flatter optimization trajectories.Cautious Optimizer seeks descent-direction updates under a small step size, but its deterministic rule lacks Magma’s geometric regularization effect.
  • Geometry-aware and trust-region methods: Curvature-aware preconditioners approximate second-order structure, while trust-region methods bias optimization toward flatter regions through parameter-perturbation robustness.Trust-region–flavored methods such as SAM and its variants incur additional gradient evaluations.
  • Geometry-aware and trust-region methods: Magma offers a lightweight alternative by stochastically masking block-wise updates to penalize sharp curvature along their directions, without explicit curvature matrices or adversarial perturbations.This positions Magma as a geometry-aware method with lower computational demands than approaches requiring additional gradient evaluations.
  • Stochastic perturbation and noise injection: Prior stochastic perturbation methods include random gradient masking, annealed Gaussian noise, Dropout, and embedding-space perturbations, but stateful optimization perturbations remain underexplored.Dropout induces a data-dependent weight-space regularizer that promotes model stability, while modern LLM work has extended perturbations to token embeddings.

7. Conclusion

Randomly masking parameter updates improves LLM pre-training by smoothing optimization trajectories through implicit curvature-dependent geometric regularization. Magma further enhances masked updates via momentum-gradient alignment, consistently outperforming state-of-the-art adaptive optimizers with negligible overhead.

  • 7. Conclusion: Randomly masking parameter updates can substantially improve LLM pre-training while smoothing optimization trajectories through curvature-dependent geometric regularization.The regularization is implicit and depends on curvature.
  • 7. Conclusion: Magma leverages momentum-gradient alignment to enhance masked updates and consistently improves on state-of-the-art adaptive optimizers with negligible overhead.Magma is presented as a masked-update method for adaptive optimization.
  • 7. Conclusion: These findings challenge the assumption that dense updates are inherently optimal for backpropagation-based neural network training.The conclusion frames masking as a challenge to prevailing optimizer-design assumptions.

A. Proofs of Claims … B.2. Nano MoE Pre-Training Benchmark Setup

The paper establishes its theoretical claims through conditional-expectation, Taylor-expansion, inequality, and telescoping arguments, then specifies C4 and Nano MoE pre-training benchmarks with model- and optimizer-training configurations.

  • A.1. Proof of Proposition 1: A.1 expands the masked-update loss using a second-order Taylor expansion conditioned on F_t, defining the gradient, Hessian blocks, and remainder term.The proof then takes conditional expectations and regroups terms.
  • A.2. Proof of Lemma 4: A.2 proves Lemma 4 by applying Assumption 2 sequentially to an arbitrary multi-block update and then taking conditional expectations.The conditional expectation yields the desired result.
  • A.3. Proof of Lemma 5: A.3 defines an event and sigmoid-based constants, applies total expectation and Cauchy–Schwarz, and invokes Assumption 3 for the final inequality.The proof notes that additional relations follow directly from the definitions.
  • A.4. Proof of Theorem 6: A.4 derives Theorem 6 from equations 4–6 and Lemma 5, then takes total expectations and telescopes the resulting inequality over t = 0 to T−1.The proof uses E[l(θ_T)] > l* before elementary algebra produces the desired bound.
  • B.1. C4 Pre-Training Benchmark Setup: B.1 follows Zhao et al. (2024) with batch size 512, sequence length 256, and learning-rate candidates from 1e-4 through 1e-2.Learning rates warm up for 10% of training, then cosine-anneal to 10% of the peak value.
  • B.1. C4 Pre-Training Benchmark Setup: B.1 trains the 60M, 130M, 350M, and 1B models for 10K, 20K, 60K, and 100K iterations, respectively, reporting final evaluation perplexity from the best-learning-rate run.The reported metric is selected after comparing the specified learning-rate grid.
  • B.2. Nano MoE Pre-Training Benchmark Setup: B.2 uses a 124M-parameter GPT2-style transformer with 8 experts per MoE layer, top-2 routing, and alternating dense and MoE layers applied with stride 2.The model is trained for 50K iterations on 8xA100 GPUs.
  • B.2. Nano MoE Pre-Training Benchmark Setup: B.2 adopts batch size 12, gradient accumulation 40, sequence length 1024, minimum learning rate 5e-6, weight decay 0.1, and grad clip norm 1.0.These values are part of the default Nano MoE configuration.

B.3. Heterogeneous Quadratic Benchmark Setup

The benchmark uses two 9-dimensional quadratic problems with identical spectra and 3 × 3 block-diagonal Hessians, differing in how eigenvalues are arranged across blocks. Random block rotations and row subsampling provide stochastic loss and gradient approximations.

  • Benchmark construction: The setup contains two quadratic optimization problems in ℝ9, each with a 3 × 3 block-diagonal Hessian and identical eigenspectrum.Both losses have the form L(w) = 1/2 w⊤Hw.
  • Hessian structures: The homogeneous Hessian groups eigenvalues by scale, whereas the heterogeneous Hessian interleaves vastly different magnitudes within each 3 × 3 block.The heterogeneous arrangement is intended to mimic qualitative distinctions between autoregressive language-model landscapes and shallower architectures such as CNNs.
  • Stochastic evaluation: Each block receives an independent random rotation, and random row subsampling of X = H^1/2 yields stochastic loss and gradient approximations at each iteration.

B.4. Heavy-Tailed Gradient Noise Benchmark Setup · C. Ablation Studies · C.1. Masking Component

The study combines a controlled linear-transformer benchmark for evaluating optimizer robustness under heavy-tailed gradient noise with ablations of Magma’s masking scope. Targeted masking improves validation perplexity, with simultaneous attention-and-MLP masking performing best.

  • B.4. Heavy-Tailed Gradient Noise Benchmark Setup: Each input sequence in the controlled linear-transformer benchmark defines a distinct linear regression task.The benchmark is introduced by Ahn et al. (2024).
  • B.4. Heavy-Tailed Gradient Noise Benchmark Setup: The latent regression vector 𝒘∼N (0, 𝐼𝑑) is independently sampled for each sequence.The objective predicts 𝑦𝑛+1 from n contextual input-output pairs.
  • B.4. Heavy-Tailed Gradient Noise Benchmark Setup: Training minimizes mean squared prediction error with dimension 𝑑= 5 and context length 𝑛= 20.The setting is used to study in-context learning and transformer optimization dynamics.
  • B.4. Heavy-Tailed Gradient Noise Benchmark Setup: Heavy-tailed covariates are sampled uniformly from the unit sphere and scaled by independent variables drawn from Γ0.1,10.The light-tailed comparison samples covariates from 𝒙𝑖∼N (0, 𝐼𝑑).
  • B.4. Heavy-Tailed Gradient Noise Benchmark Setup: The heavy-tailed construction amplifies covariate tails, inducing heavy-tailed gradient noise for controlled optimizer-robustness evaluation.It is designed to test robustness under extreme stochastic fluctuations.
  • C. Ablation Studies: Ablations use a 130M-parameter Llama model trained on C4 with RMSProp+Magma, following the setup in § 4.1.This establishes the experimental context for the masking studies.
  • C.1. Masking Component: Attention-only masking reduces validation perplexity from a baseline of 22.64 to 21.92.The study compares masking specific transformer sub-modules against a global masking strategy.
  • C.1. Masking Component: Masking both attention and MLP achieves the lowest perplexity of 21.65, surpassing the most comprehensive setting (21.94).The result indicates that targeted sub-module regularization outperforms uniform masking strategies.

C.2. Masking Granularity … C.5. Sensitivity to Learning Rate

Across masking and optimizer settings, performance is robust to granularity and temperature, while sampling with damping improves perplexity. Dense momentum updates and Adam+Magma provide more stable optimization across update sparsity and learning-rate changes.

  • C.2. Masking Granularity: 21.73 to 21.81: Uniform Sampling changes validation perplexity only slightly from Element to Block masking, indicating robustness across granularities.Fine-grained masking provides a small edge over block masking.
  • C.2. Masking Granularity: Block masking is preferable for efficiency because it saves significant memory for cosine similarities with minimal accuracy loss.
  • C.2. Masking Granularity: 21.58: Uniform sampling plus damping achieves the minimum reported perplexity, whereas damping alone improves the RMSProp baseline from 22.64 to around 21.92.Combining sampling with damping consistently produces the lowest perplexity scores.
  • C.3. Sampling Ratio and Damping Temperature: p = 0.5 outperforms p = 0.25 and p = 0.75 across all tested temperatures, while results are not very sensitive to temperature.The tested temperatures are τ ∈ {0.5, 1.0, 2.0, 4.0}.
  • C.3. Sampling Ratio and Damping Temperature: τ = 2.0 is used in all experiments after the sampling-ratio and temperature analysis.
  • C.4. Sparse vs. Dense Momentum Update: At learning rate 0.001, dense momentum updates consistently converge robustly and achieve the lowest perplexity, whereas sparse updates without damping remain severely unstable.Damping stabilizes sparse updates, but their trajectory still underperforms dense baselines over 20,000 iterations.
  • C.5. Sensitivity to Learning Rate: Up to 0.05: Adam+Magma remains effective across a broader learning-rate spectrum, while Adam and C-Adam spike outside approximately 0.001–0.003 and fail to converge.The comparison is based on evaluation perplexity and learning-rate sensitivity.
Loading 2602.15322v1…