On Surprising Effectiveness of Masking Updates in Adaptive Optimizers | Summary

Contents

• Summary
• 1. Introduction
• 2. Update Masking as a Regularization
• 3. Momentum-Aligned Update Masking
• 4. Experiments
  • 4.1. Pre-Training Llama
  • 4.2. Pre-Training Nano MoE
  • 4.3. Magma under Heavy-Tailed Gradient Noises
  • 4.4. Magma on Heterogeneous Quadratics
• 5. Discussion
• 6. Literature Review
• 7. Conclusion
• Impact Statement
• Appendix
• Brief Thoughts

This article explains the key points of On Surprising Effectiveness of Masking Updates in Adaptive Optimizers.

• 2026-02-17 (arXiv)
• Joo, Taejong, Xia, Wenhan, Kim, Cheolmin, Zhang, Ming, Ie, Eugene.
• Northwestern University, Google
• Paper

Read this article in Korean

Summary

• On Surprising Effectiveness of Masking Updates in Adaptive Optimizers studies whether randomly skipping parameter updates can improve LLM pre-training while retaining dense backpropagation. SkipUpdate independently masks entire parameter blocks, updates optimizer states densely, and rescales surviving updates to preserve their expectation.
• A local second-order analysis identifies an additional block-wise curvature term in the expected post-update loss. Momentum-aligned gradient masking (Magma) combines stochastic masking with smoothed, cosine-similarity-based damping, giving larger update multipliers to blocks whose gradients agree with accumulated momentum.
• On C4, RMSProp+Magma achieves validation perplexities of 28.55, 21.66, 16.16, and 13.19 for 60M, 130M, 350M, and 1B Llama 2 models, respectively. The 1B result improves on Adam at 16.35 and Muon at 14.52 by over 19% and 9%, respectively. Controlled experiments show benefits in the paper’s heavy-tailed-noise setting and on heterogeneous quadratics, but no improvement is reported for ResNet-50 on CIFAR-10.

1. Introduction

Dense backpropagation makes simultaneous adaptive updates convenient, but it does not establish that every parameter should be updated at every iteration. The paper tests this assumption with SkipUpdate, an RMSProp-based method that discards half of parameter-block updates while continuing to update moment estimates densely.

• Figure 1 shows lower C4 validation perplexity for SkipUpdate than the plotted dense baselines across 60M–1B model sizes.
• Skipping parameter updates does not reduce the cost of the backward pass or remove dense optimizer states. The experiment concerns optimization quality rather than sparse gradient computation.
• The authors interpret masking noise as geometric regularization and use momentum–gradient agreement to introduce block-adaptive damping in Magma.

SkipUpdate's curve is below the plotted Adam, Muon, Adafactor, and GaLore curves across the four model sizes, with a small separation from Muon at 1B. This establishes the motivating observation that discarding half of parameter-block updates can improve validation perplexity. It does not measure training speed or identify the mechanism behind the improvement.

2. Update Masking as a Regularization

SkipUpdate applies independent Bernoulli masks with survival probability p and uses the masked block update Δ̃ₜ⁽ᵇ⁾ = mₜ⁽ᵇ⁾Δₜ⁽ᵇ⁾/p. Conditioned on the parameters and base update, Proposition 1 expresses the expected post-update loss as the dense-update loss plus a block-wise curvature term, up to a third-order remainder.

• The additional term is (1−p)/(2p) times the sum of (Δₜ⁽ᵇ⁾)ᵀHᵦᵦ(θₜ)Δₜ⁽ᵇ⁾ over blocks. Independent masking preserves cross-block second-order terms but amplifies within-block terms.
• For positive-curvature directions, this term raises the expected loss associated with sharp updates. The authors interpret it as implicit geometric regularization, motivated by reported block structure in transformer Hessians.
• The result requires a valid local Taylor expansion and remainder bound. It does not establish that the practical training algorithm minimizes the curvature term, reaches flatter minima, or improves generalization.

The method masks parameter updates rather than the gradients used to update moment estimates. Masking granularity determines the Hessian interactions retained in the second-order term: element-wise masking contributes diagonal curvature terms, whereas block-wise masking retains within-block interactions.

• For 130M Llama pre-training, uniform masking yields perplexities of 21.78 for column-wise, 21.73 for element-wise, and 21.81 for block-wise masking, compared with 22.64 for RMSProp.
• The authors adopt block-wise masking because the measured differences are small and skipping whole blocks permits operation pruning.
• The dense-versus-sparse state-update ablation shows instability with sparse momentum updates. Unlike memory-saving subspace methods, SkipUpdate retains dense auxiliary states.

3. Momentum-Aligned Update Masking

Magma computes a scale-invariant alignment score from the cosine similarity between each block’s current gradient and first-moment estimate. Algorithm 1 uses s̃ₜ⁽ᵇ⁾ = sigmoid(cossim(μₜ⁽ᵇ⁾, gₜ⁽ᵇ⁾)/τ), smooths it as sₜ⁽ᵇ⁾ = 0.9sₜ₋₁⁽ᵇ⁾ + 0.1s̃ₜ⁽ᵇ⁾, and multiplies the base optimizer update by this score and an independent Bernoulli(0.5) mask.

• The default mask remains uniformly sampled; alignment controls update magnitude rather than survival probability.
• Unlike SkipUpdate’s factor of 2 at p = 0.5, Magma’s damping introduces bias. The authors report unstable training with tested unbiased alternatives that use alignment-dependent survival probabilities and inverse-probability rescaling.
• Cautious Optimizer, MGUP, and RPROP also use gradient consistency, but do not implement Magma’s structured stochastic masking rule.
• The wrapper requires no additional backward pass or explicit curvature calculation. Its alignment computations and score smoothing still entail operations and state; the PDF provides no wall-clock or memory measurements substantiating its zero-overhead claim.

4. Experiments

The experiments combine language-model pre-training with controlled optimization benchmarks. Llama 2 on C4 tests multiple model scales, Nano MoE on OpenWebText tests compatibility with sparse expert routing, and simplified regression and quadratic tasks examine noise and curvature effects.

4.1. Pre-Training Llama

The C4 benchmark follows Zhao et al. (2024), training Llama 2 models of 60M, 130M, 350M, and 1B parameters with batch size 512 and maximum sequence length 256. The learning-rate grid is {1e-4, 5e-4, 1e-3, 5e-3, 1e-2}; each run uses 10% warm-up followed by cosine decay to 10% of the peak learning rate.

• Training lasts 10K, 20K, 60K, and 100K iterations for the four model sizes, respectively. Reported perplexity is the final evaluation result from the run with the best searched learning rate.
• Magma uses τ = 2 and is applied exclusively to attention and MLP layers.
• Table 1 distinguishes the authors’ results from marked results taken from Li et al. (2025); not all comparisons are newly executed runs.

Table 1 shows improvements from adding Magma to Adam, LaProp, and RMSProp at every scale with a reported baseline result. RMSProp+Magma obtains the lowest perplexity in every column, reaching 13.19 at 1B, where unwrapped RMSProp diverges within the searched learning-rate grid.

• Adam improves from 30.79, 24.77, 18.42, and 16.35 to 29.09, 22.08, 16.41, and 13.71 across the four scales.
• LaProp+Magma reaches 29.05, 22.16, 16.37, and 13.82, compared with 29.98, 23.07, 18.56, and 16.38 for LaProp.
• At 1B, the table reports Adam+SGG at 14.30, APOLLO+SGG at 13.95, Muon at 14.52, and C-Adam at 15.92.
• Table 1 gives Adam+Magma at 1B as 13.71, while the surrounding prose gives 13.81. This review uses the table value. SOAP is mentioned in the comparison prose but has no row in Table 1.

RMSProp+Magma has the lowest reported perplexity in every model-size column, including 13.19 at 1B, where unwrapped RMSProp diverges within the search grid. Adam and LaProp also improve with Magma at every reported scale. The table combines the authors' runs with marked results from Li et al. (2025), and its Adam+Magma value of 13.71 at 1B conflicts with the prose value of 13.81.

4.2. Pre-Training Nano MoE

Nano MoE pre-training uses a 124M-parameter GPT2-style transformer on OpenWebText, with 8 experts per MoE layer, top-2 routing, and alternating dense and MoE layers. Figure 2 shows lower final evaluation loss when Magma is added to AdamW or Muon, with Muon+Magma achieving the best plotted endpoint.

• Training follows Wolfe (2024) for 50K iterations on 8xA100 GPUs, with batch size 12, gradient accumulation 40, sequence length 1024, minimum learning rate 5e-6, weight decay 0.1, and gradient clipping norm 1.0.
• AdamW+Magma progresses more slowly during intermediate training but finishes below AdamW and C-AdamW. These names follow the figure legend, whereas the accompanying prose uses Adam and C-Adam.
• The Muon result demonstrates compatibility with matrix-based preconditioning. It does not establish that preconditioning and masking act through independent mechanisms.

Muon+Magma finishes with the lowest evaluation loss over 50K Nano MoE iterations. AdamW+Magma initially progresses more slowly than AdamW but reaches a lower endpoint, so its advantage is not uniform throughout training. C-AdamW shows visible loss spikes and finishes above the Magma-enhanced variants.

4.3. Magma under Heavy-Tailed Gradient Noises

The noise benchmark follows Ahn et al. (2024), training a simplified linear transformer to predict the final label in a sequence representing a random linear regression task. With dimension d = 5 and context length n = 20, it compares Gaussian covariates against unit-sphere covariates scaled by an independent √Γ₀.₁,₁₀ variable, where Γ denotes a Gamma distribution with shape 0.1 and scale 10.

• Figure 3 shows similar final losses for Adam and Magma under Gaussian covariates, but lower loss for Magma in the setting labeled heavy-tailed by the paper.
• That label should not be taken as evidence of mathematically heavy-tailed covariates: the specified Gamma-based radial distribution has light tails. The experiment nevertheless tests a different, highly variable covariate distribution and its induced optimization behavior.
• The robust condition number is the maximum Hessian eigenvalue divided by its median eigenvalue. Its much smaller value along Magma’s trajectory in the second setting associates lower loss with better-conditioned visited regions, without isolating the contribution of masking or damping.

The Gaussian-covariate panels show similar late-training losses and broadly comparable conditioning. In the panels labeled heavy-tailed, Magma attains lower loss while its robust Hessian condition number remains far below Adam's rising trajectory. The paired diagnostics associate improved optimization with better-conditioned visited regions, but do not establish a causal decomposition of masking and damping or validate the mathematical tail classification.

4.4. Magma on Heterogeneous Quadratics

The quadratic benchmark holds the Hessian eigenspectrum fixed while changing its arrangement across three parameter blocks in ℝ⁹. Both objectives have eigenvalues {1, 2, 3, 99, 100, 101, 4998, 4999, 5000}, independent random rotations within each 3 × 3 block, and stochastic gradients generated by row subsampling from X = H¹ᐟ².

• The homogeneous blocks group similar scales: {1, 2, 3}, {99, 100, 101}, and {4998, 4999, 5000}. The heterogeneous blocks mix scales: {1, 99, 4998}, {2, 100, 4999}, and {3, 101, 5000}.
• Figure 4 shows comparable best final performance on the homogeneous objective and faster convergence to lower loss for AdamW+Magma on the heterogeneous objective.
• The plotted learning rates are 0.003, 0.01, 0.03, and 0.1. Several Magma configurations make little progress, so the figure does not show an advantage at every learning rate.

The lower panels of Figure 4 show no substantial increase in gradient–momentum agreement under Magma despite its better heterogeneous-quadratic optimization. The plotted diagnostic averages the positive-inner-product indicator, rather than the cosine-similarity score used for damping; the authors interpret the result as evidence that controlling applied updates matters more than changing alignment statistics.

• The paper reports no improvement over AdamW for ResNet-50 on CIFAR-10 after 100 epochs with tuned configurations, giving test accuracies of 94.46% versus 93.82%.
• This negative result limits the demonstrated benefit to particular settings; it does not establish that curvature alone explains differences between architectures.
• The nine-dimensional quadratic construction isolates block geometry, but remains a qualitative abstraction of transformer training.

With the eigenspectrum held fixed, the best-performing runs reach similar final losses on homogeneous quadratics, while AdamW+Magma reaches lower loss on heterogeneous quadratics. Some low-learning-rate Magma runs make little progress. The lower panels compare the frequency of positive gradient–momentum inner products at selected learning rates and show no substantial increase under Magma; they do not plot the damping cosine-similarity score.

5. Discussion

The convergence discussion considers constant-learning-rate SGD rather than the practical adaptive-optimizer wrapper. Under stated block-wise smoothness and stochastic-gradient second-moment assumptions, Lemma 4 proposes a descent bound, Lemma 5 introduces effective descent and noise-coupling factors, and Theorem 6 states an average squared-gradient bound.

• The effective block smoothness is L̃ₜ⁽ᵇ⁾ = ρₜ⁽ᵇ⁾L⁽ᵇ⁾/p, where ρ measures the second-moment reduction caused by damping.
• The stated theorem contains an optimization term decreasing as 1/T, a noise-descent coupling term, and a curvature-weighted noise term, subject to its learning-rate restriction.
• Damping can reduce useful descent as well as destabilizing second moments. The proposed favorable regime requires attenuation of blocks that dominate curvature-weighted noise or smoothness constraints.
• The proofs require qualification: A.2 asserts a simultaneous multi-block smoothness bound that does not follow from coordinate-wise smoothness alone, and A.3 uses alignment with the true gradient to bound damping computed from smoothed momentum alignment.

The analysis offers a proposed account of stability rather than a verified guarantee for the practical wrapper. Its masking operator includes inverse-probability rescaling, whereas Algorithm 1’s Magma update directly multiplies the base update by the damped Bernoulli mask.

• Even apart from the proof gaps, the SGD analysis is not a complete convergence guarantee for RMSProp+Magma or Muon+Magma.
• The stochastic-gradient assumption bounds second moments and does not cover arbitrary infinite-variance heavy-tailed noise.
• The wider favorable learning-rate range in Figure A3 is empirical evidence; the analysis does not establish that momentum alignment reliably identifies high-curvature blocks in LLM training.

6. Literature Review

The literature review relates Magma to update stabilization, geometry-aware optimization, and stochastic regularization. Its design perturbs parameter updates through structured random masks and alignment-dependent damping, without computing an explicit curvature matrix or an additional adversarial gradient.

• Cautious Optimizer masks updates according to gradient–momentum sign disagreement. Other stabilization approaches include gradient clipping with momentum resets, initialization changes, and architectural interventions.
• Curvature-aware preconditioners approximate second-order structure, while sharpness-aware methods such as SAM use additional gradient evaluations. Magma avoids those computations, although its local curvature interpretation does not establish an equivalent flatness-seeking objective.
• The review also covers gradient masking in meta-learning and federated learning, Gaussian noise injection, Dropout, and noisy embeddings. Magma places its perturbations on optimizer updates rather than hidden activations or token representations.

7. Conclusion

The conclusion argues that dense parameter updates need not be optimal merely because dense gradients are available. The demonstrated contribution is an optimizer wrapper that improves reported Llama and Nano MoE pre-training metrics, accompanied by a local curvature expansion and controlled optimization experiments.

Impact Statement

The Impact Statement describes the work as advancing machine learning and states that the authors do not identify societal consequences requiring specific discussion. It contains no separate assessment of deployment risks or resource impacts.

Appendix

• A. Proofs of Claims contains A.1. Proof of Proposition 1, A.2. Proof of Lemma 4, A.3. Proof of Lemma 5, and A.4. Proof of Theorem 6. A.1 uses a Taylor expansion and independent mask moments; A.2 asserts a multi-block smoothness bound; A.3 decomposes true-gradient alignment events; and A.4 telescopes the proposed descent inequalities. The latter results need additional justification because coordinate-wise smoothness alone does not imply A.2’s bound, and A.3 does not justify its alignment-dependent lower bound for the practical smoothed damping score.
• B. Experimental Details documents the four benchmark setups. B.1. C4 Pre-Training Benchmark Setup specifies batch size 512, sequence length 256, the learning-rate grid {1e-4, 5e-4, 1e-3, 5e-3, 1e-2}, 10% warm-up, cosine decay to 10% of the peak, and 10K/20K/60K/100K iterations for 60M/130M/350M/1B models. Selection uses final evaluation perplexity from the run with the best searched learning rate.
• B.2. Nano MoE Pre-Training Benchmark Setup follows Wolfe (2024), specifying a 124M GPT2-style model with 8 experts, top-2 routing, MoE stride 2, an auxiliary load-balancing loss, and Switch Transformer-style initialization. Its 50K-iteration run uses 8xA100 GPUs, batch size 12, accumulation 40, sequence length 1024, minimum learning rate 5e-6, weight decay 0.1, and gradient clipping norm 1.0.
• B.3. Heterogeneous Quadratic Benchmark Setup fixes the eigenvalues {1, 2, 3, 99, 100, 101, 4998, 4999, 5000} and three 3 × 3 Hessian blocks, changing their grouping by scale. Independent random block rotations and row subsampling of X = H¹ᐟ² produce the stochastic quadratic benchmark used to compare homogeneous and heterogeneous curvature.
• B.4. Heavy-Tailed Gradient Noise Benchmark Setup samples a separate latent regression vector w from N(0, I_d) for each sequence and minimizes mean squared prediction error for the final label. It sets d = 5 and n = 20, comparing N(0, I_d) covariates with unit-sphere samples multiplied by √Γ₀.₁,₁₀. The latter is the paper’s heavy-tailed setting, although the specified covariate distribution is not mathematically heavy-tailed.
• C. Ablation Studies primarily uses a 130M Llama model on C4 with RMSProp+Magma. C.1. Masking Component reports perplexities of 22.64 for the baseline, 21.92 for attention-only masking, 21.65 for attention+MLP masking, and 21.94 for masking all components. These results favor the targeted attention+MLP configuration over global masking in this experiment.
• C.2. Masking Granularity compares element, row, column, and block masks with several sampling and damping combinations. Uniform sampling gives 21.73, 21.76, 21.78, and 21.81; damping alone gives 21.97, 21.95, 21.91, and 21.92; and uniform sampling+damping gives 21.58, 21.62, 21.61, and 21.65. Alignment-based sampling gives 21.77, 21.78, 21.75, and 21.78, while alignment-based sampling+damping gives 21.63, 21.60, 21.61, and 21.65, providing no consistent advantage over uniform sampling.
• C.3. Sampling Ratio and Damping Temperature tests p ∈ {0.25, 0.5, 0.75} and τ ∈ {0.5, 1.0, 2.0, 4.0}. Figure A1 shows p = 0.5 outperforming the other survival probabilities at every tested temperature, with relatively modest temperature sensitivity. The authors adopt τ = 2.0, although it is not the lowest plotted perplexity for p = 0.5.
• C.4. Sparse vs. Dense Momentum Update compares four configurations at learning rate 0.001 over 20,000 iterations. Figure A2 shows similar strong convergence for dense momentum updates with and without damping, severe instability for sparse updates without damping, and partial stabilization with damping. The damped sparse variant still finishes above the dense variants in training perplexity.
• C.5. Sensitivity to Learning Rate reports that Adam+Magma remains effective at learning rates up to 0.05, whereas Adam and C-Adam have narrower favorable regions described in the prose as approximately 0.001–0.003. Figure A3 supports a wider favorable learning-rate range at the tested points, not the elimination of learning-rate tuning across tasks.

Brief Thoughts

The strongest evidence is the consistent C4 improvement across model sizes, the successful Muon combination on Nano MoE, and the ablations separating sampling, damping, and dense state updates. Together, these results justify testing update masking as an optimizer modification, even though it discards some already-computed updates.

The mechanism remains less secure than the empirical result: a positive curvature term in expected post-update loss does not prove flat-minimum selection, and the convergence proofs contain gaps between their assumptions and the practical damping rule. The PDF also reports no uncertainty estimates for the C4 results, no scale beyond 1B, no measured wall-clock overhead, and a negative CNN result. The 13.71/13.81 reporting discrepancy and the Gamma-based benchmark’s heavy-tailed label warrant caution when interpreting the numerical and robustness claims.