Monarch-Muon: Structure-Aware Muon for Fast and Memory-Efficient Pretraining

Machine Learning

Summary

The gist is being written…

Authors

Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov, Dmitrii Vlashchuk, Vladimir Aletov, Nail Bashirov, Daniil Tikhonov, Maxim Krutikov, Aleksandr Shestakov, Aleksandr Beznosikov

Abstract

Training large language models is constrained by both wall-clock time and memory footprint, and recent matrix-aware optimizers sharpen this tradeoff: Muon can improve optimization but its dense polar update remains expensive. We propose Monarch-Muon, a structured matrix-aware optimizer for weights parameterized by Monarch factors. The resulting structure reduces weights, gradients, and optimizer state, and exposes small independent update subproblems in place of a single dense orthogonalization; we formalize this through a blockwise trust-region interpretation. In quality runs, Monarch-Muon improves over Monarch-AdamW at fixed structure: 3.06 vs. 3.09 in the dense 257M model and 2.68 vs. 2.77 in the 1B MoE model; the two-block MoE variant also matches dense AdamW while using 15% fewer active and 25.4% fewer total parameters. In 6.89B-parameter systems measurements, Monarch-Muon reduces peak memory from 68.9 to 52.9 GB, Muon optimizer state from 14.60 to 6.77 GB, and the Muon optimizer step from 1288 to 304 ms.