AMO: Adaptive Muon Orthogonalization
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2026
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866914576320167936 |
|---|---|
| author | Zhuang, Xinlin Ouyang, Panyi Li, Yichen Shi, Jiangming Chen, Yizhang Liu, Shuman Qian, Ying Liu, Weiyang Zhang, Haibo Razzak, Imran |
| author_facet | Zhuang, Xinlin Ouyang, Panyi Li, Yichen Shi, Jiangming Chen, Yizhang Liu, Shuman Qian, Ying Liu, Weiyang Zhang, Haibo Razzak, Imran |
| contents | Muon has recently emerged as a competitive alternative to AdamW for large-scale pre-training, with orthogonalization via Newton-Schulz (NS) iterations as its core operation. Existing Muon variants apply a uniform NS schedule to all parameter matrices, overlooking possible differences in orthogonalization difficulty and its impact on performance. Through a systematic empirical study, we show that this per-matrix heterogeneity is pervasive and largely determined by matrix geometry, which evolves dynamically across operator types, training stages, and network depths. As a result, uniform NS schedules can lead to uneven orthogonalization quality across the model. Motivated by these findings, we propose Adaptive Muon Orthogonalization (AMO), an observe-then-commit method that measures weight geometry by operator type early in training and then uses these signals to allocate the NS budget for the remainder of training. AMO delivers consistent improvements over uniform-schedule Muon across standard, prolonged, and continual pre-training, surpassing the strongest baseline by +0.76 on Llama3.1-1.4B and +0.51 on Qwen3-1.7B in average downstream performance of 12 evaluation tasks. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_17806 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | AMO: Adaptive Muon Orthogonalization Zhuang, Xinlin Ouyang, Panyi Li, Yichen Shi, Jiangming Chen, Yizhang Liu, Shuman Qian, Ying Liu, Weiyang Zhang, Haibo Razzak, Imran Machine Learning Muon has recently emerged as a competitive alternative to AdamW for large-scale pre-training, with orthogonalization via Newton-Schulz (NS) iterations as its core operation. Existing Muon variants apply a uniform NS schedule to all parameter matrices, overlooking possible differences in orthogonalization difficulty and its impact on performance. Through a systematic empirical study, we show that this per-matrix heterogeneity is pervasive and largely determined by matrix geometry, which evolves dynamically across operator types, training stages, and network depths. As a result, uniform NS schedules can lead to uneven orthogonalization quality across the model. Motivated by these findings, we propose Adaptive Muon Orthogonalization (AMO), an observe-then-commit method that measures weight geometry by operator type early in training and then uses these signals to allocate the NS budget for the remainder of training. AMO delivers consistent improvements over uniform-schedule Muon across standard, prolonged, and continual pre-training, surpassing the strongest baseline by +0.76 on Llama3.1-1.4B and +0.51 on Qwen3-1.7B in average downstream performance of 12 evaluation tasks. |
| title | AMO: Adaptive Muon Orthogonalization |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2605.17806 |