MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

2026-08-05

AI Summary

This work addresses the sensitivity of the MuON optimizer to anisotropic curvature in the loss landscape by proposing MALT and its robust variant, MALTER. The method introduces a lightweight bilateral diagonal preconditioner to approximate the geometry of curvature, incorporates Newton–Schulz iteration for momentum orthogonalization, and employs norm grafting with adaptive step-size rescaling to control update magnitudes. These enhancements enable curvature-aware optimization with negligible additional computational or memory overhead. Empirical evaluation on GPT-2 Small, Medium, and Large pretraining tasks demonstrates that the MALT family consistently outperforms MuON, confirming its effectiveness, stability, and scalability across model sizes.

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.

Problem

• curvature anisotropy
• loss landscape
• optimization
• gradient anisotropy
• preconditioning

Innovation

• diagonal preconditioning
• curvature-aware optimization
• lightweight optimizer
• Newton-Schulz iterations
• stochastic non-convex convergence