Structured weights make Muon training cheaper at scale
In plain words
Training a language model takes substantial computing resources. This paper introduces Monarch, Muon, an optimizer that stores some model weights as smaller pieces and updates those pieces with matrix-aware steps.
In single-run tests on two model types, it had lower final validation loss than AdamW when both used the same piece structure. Separate measurements found lower training-step time and memory at the largest tested size, on one H100.
The research
Can language-model pretraining use matrix-aware updates, which can improve optimization, without paying the full storage and computation costs of dense weights?
The paper tests whether representing selected weights with structured factors can make those updates cheaper while retaining their optimization benefit.
Monarch–Muon applies matrix-aware updates to structured weight factors, aiming to retain optimization benefits while reducing storage and orthogonalization costs.
Monarch–Muon updates structured weight factors with matrix-aware steps. Researchers compared it with AdamW at matching factor structures in dense and mixture-of-experts pretraining on FineWeb-Edu.
Final validation loss was lower with Monarch–Muon in all four single-run comparisons. This suggests the matrix-aware updates can help within these tested structures.
The quality comparisons are single runs at specific model sizes, and the separate systems measurements do not establish the combined quality and hardware tradeoff in longer distributed training.
The convergence proof assumes exact polar updates, while the practical optimizer approximates them with finite Newton–Schulz iterations.
Analysis
Deep
- Original paper authors
- Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov, Dmitrii Vlashchuk, Vladimir Aletov
- Submitted to arXiv
- License
- CC BY 4.0