MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

Tongle Wu, Huanyu Dong, Ying Sun +1

Muon has recently emerged as a promising alternative to AdamW for language model pretraining by orthogonalizing momentum matrices using Newton-Schulz iterations. Although Muon miti…