MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

• URL: http://arxiv.org/abs/2608.05088v1
• Date: Wed, 05 Aug 2026 17:26:47 GMT
• Title: MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning
• Abstract summary: We propose MALT to reduce the sensitivity of Muon to curvature geometry of the loss landscape.<n>MALT uses two-sided diagonal preconditioners with low memory and overhead to capture the curvature geometry of the loss landscape.
• Score: 14.373662883104808
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Muon has recently emerged as a promising alternative to AdamW for language model pretraining by orthogonalizing momentum matrices using Newton-Schulz iterations. Although Muon mitigates gradient anisotropy, it does not explicitly account for the curvature geometry of the loss landscape and may therefore remain sensitive to curvature anisotropy. We bridge this gap by proposing MALT (Muon Augmented by Lightweight Two-sided Preconditioning), which uses lightweight diagonal preconditioners to reduce the sensitivity of Muon to curvature anisotropy. Specifically, MALT uses two-sided diagonal preconditioners with low memory and computational overhead to approximately capture the curvature geometry of the loss landscape. It orthogonalizes the preconditioned momentum using Newton-Schulz iterations and maps the result back to define the update direction, while norm grafting controls the update magnitude. To improve the robustness of MALT to stochastic gradient noise, we further propose MALTER (MALT with Adaptive stEpsize Rescaling). Convergence guarantees are provided for MALT in the stochastic non-convex setting. Experiments on GPT-2 Small, Medium, and Large pretraining show that the proposed methods outperform Muon while maintaining nearly the same memory footprint and wall-clock time.

Related papers

• Beyond the Matrix Sign: Quadratic Spectral Descent [6.589542614722917]
  We show that curvature can change both the singular values and the singular directions of the optimal update.<n>Experiments on GPT pre-training show that QSD consistently improves validation loss over Muon and recent Muon variants.
  arXiv  Detail & Related papers  (2026-09-07T15:07:24Z)
• Why Muon Outperforms Adam: A Curvature Perspective [49.85900116602357]
  Muon improves training efficiency over Adam in large language-model training by about two times.<n>Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective.
  arXiv  Detail & Related papers  (2026-06-03T09:40:30Z)
• Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training [56.323119575322146]
  Extra-Merge consistently outperforms standard merging baselines in experiments across GPT-2 and LLaMA families.<n>It yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon citepjordan2024muon.
  arXiv  Detail & Related papers  (2026-05-26T02:48:34Z)
• AMUSE: Anytime Muon with Stable Gradient Evaluation [25.647698512832648]
  We study Muon through the river-valley loss landscape, where useful training progress occurs along a flat, low-curvature bulk subspace.<n>We propose Anytime MUon with Stable gradient Evaluation (AMUSE), which integrates Muon's rapid bulk progress with the stabilizing effect of Schedule-Free averaging.
  arXiv  Detail & Related papers  (2026-05-21T12:55:52Z)
• Distance-Aware Muon: Adaptive Step Scaling for Normalized Optimization [47.817303172845506]
  Muon and related normalizeds decouple the choice of update direction from the choice of step scale.<n>We develop three complementary algorithms for Muon tuning and develop the normalized adaptive scaling parameter.
  arXiv  Detail & Related papers  (2026-05-18T18:19:26Z)
• Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence [15.754779716816174]
  We show that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest.<n>We recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor.
  arXiv  Detail & Related papers  (2026-05-13T06:54:01Z)
• Muon is Not That Special: Random or Inverted Spectra Work Just as Well [50.969177887027115]
  We demonstrate that precise geometric structure is not key factor affecting performance.<n>We introduce Freon, a family of geometrics based on Schatten (quasi-size)norms.<n>We also introduce Kaon, an absurd that replaces singular values with random noise.
  arXiv  Detail & Related papers  (2026-05-11T19:42:48Z)
• Muown: Row-Norm Control for Muon Optimization [41.541871767555484]
  Muon is a strong competitor to AdamW for language model pre-training.<n>Muown improves perplexity over Muon, SOAP, AdamW, and Lion.
  arXiv  Detail & Related papers  (2026-05-11T16:26:06Z)
• Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning [72.62839712454196]
  Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improveimation.<n>We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned.<n>We propose textbfMousse (textbfMuon textbfOptimization textbfUtilizing textbfShampoo's textbfStruct
  arXiv  Detail & Related papers  (2026-03-10T14:03:49Z)
• AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates [5.049533819651459]
  We propose a new adaptive update, AdaGO, which combines a norm-based update with aGrad-type step.<n>AdaGO preserves the orthogonality of the update, which can be interpreted as a spectral descent, while adapting the stepsizes to the optimization landscape by scaling the direction with accumulated past gradients.
  arXiv  Detail & Related papers  (2025-09-03T03:42:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.