MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

TL;DR - MALT augments the Muon optimizer with lightweight two-sided diagonal preconditioners so orthogonalized momentum updates also account for loss-landscape curvature, reportedly beating Muon on GPT-2 pretraining at near-identical memory and wall-clock cost.

• Diagnoses a gap in Muon: Newton-Schulz orthogonalization of momentum matrices addresses gradient anisotropy but ignores curvature geometry, leaving it sensitive to curvature anisotropy.
• MALT applies cheap diagonal preconditioners on both sides of the momentum matrix, orthogonalizes the preconditioned momentum via Newton-Schulz, maps back to get the update direction, and uses norm grafting to set update magnitude.
• MALTER adds adaptive stepsize rescaling for robustness to stochastic gradient noise; convergence guarantees are given for MALT in the stochastic non-convex setting.
• Empirical claims come from GPT-2 Small/Medium/Large pretraining, where both methods outperform Muon with nearly the same memory footprint and wall-clock time (no specific numbers provided in the abstract).