Search papers, labs, and topics across Lattice.
This paper introduces MALT, an enhancement to the Muon optimizer that incorporates lightweight diagonal preconditioning to address curvature anisotropy in the loss landscape. By employing two-sided diagonal preconditioners, MALT effectively reduces sensitivity to curvature while maintaining low memory and computational overhead. Experimental results demonstrate that MALT outperforms the original Muon optimizer across various GPT-2 model sizes, achieving superior performance without increasing resource demands.
MALT outperforms Muon in pretraining language models while keeping memory usage and computation time nearly identical.
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.
