MuonGS relaxes weight-norm control for LLM pretraining
Original title: MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control
On dense models it reports about 15% token-equivalent speedup versus MuonH at matched validation loss.
MuonGS removes the radial component of each Muon update with Gram-Schmidt projection without fixing the weight norm. The paper proves the weight norm grows at most proportionally to the square root of the step count and that accumulated effective learning rate scales as the square root of training length. In pretraining it reports lower validation perplexity than MuonH at long horizons.
- Experiments use 160M- and 400M-parameter dense models.
- A mixture-of-experts run uses 135M active and 1B total parameters.
- With a constant learning rate, norm growth is bounded by the square root of steps.
- For smooth scale-invariant objectives the paper states an O(T to the -1/2) stationarity bound.
From the original article
1 min read at arxiv.orgWeight decay and explicit weight-norm constraints improve optimizer performance in language model pretraining, yet a unified understanding of why such control helps and how much is needed remains incomplete.
By Athanasios Glentis · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.