# MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

arXiv: https://arxiv.org/abs/2608.05088 · HTML: https://mlsanity.com/paper/2608.05088/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 74 references that are in the map.
- First posted on arXiv 5 Aug 2026.

## Builds on (most cited first, 30 of 51)

- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities](https://mlsanity.com/paper/2507.06261/) (2025)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models](https://mlsanity.com/paper/2512.02556/) (2025)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [Carbon Emissions and Large Neural Network Training](https://mlsanity.com/paper/2104.10350/) (2021)
- [Large Language Models: A Survey](https://mlsanity.com/paper/2402.06196/) (2024)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [An Investigation into Neural Net Optimization via Hessian Eigenvalue Density](https://mlsanity.com/paper/1901.10159/) (2019)
- [Empirical Analysis of the Hessian of Over-Parametrized Neural Networks](https://mlsanity.com/paper/1706.04454/) (2017)
- [Identifying and attacking the saddle point problem in high-dimensional non-convex optimization](https://mlsanity.com/paper/1406.2572/) (2014)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [Eigenvalues of the Hessian in Deep Learning: Singularity and Beyond](https://mlsanity.com/paper/1611.07476/) (2016)
- [Practical Gauss-Newton Optimisation for Deep Learning](https://mlsanity.com/paper/1706.03662/) (2017)
- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Simple and Scalable Strategies to Continually Pre-train Large Language Models](https://mlsanity.com/paper/2403.08763/) (2024)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm](https://mlsanity.com/paper/2505.16932/) (2025)
- [NorMuon: Making Muon more efficient and scalable](https://mlsanity.com/paper/2510.05491/) (2025)
- [Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective](https://mlsanity.com/paper/2410.05192/) (2024)
- [AdaMuon: Adaptive Muon Optimizer](https://mlsanity.com/paper/2507.11005/) (2025)
- [PolarGrad: A Class of Matrix-Gradient Optimizers from a Unifying Preconditioning Perspective](https://mlsanity.com/paper/2505.21799/) (2025)
- [Disentangling Adaptive Gradient Methods from Learning Rates](https://mlsanity.com/paper/2002.11803/) (2020)

## Built on by (most cited first, 1 of 1)

- [Second-Order Muon Done Right: A Principled Marriage of Spectral Geometry and Curvature](https://mlsanity.com/paper/2608.09763/) (2026)

## Abstract (arXiv)

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.