自适应调整正交化强度,让大模型训练更高效。
AMO: Adaptive Muon Orthogonalization

- 根据参数矩阵类型动态分配正交化计算资源
- 在12项任务上提升0.76~0.51分,优于基线
- 适合大规模语言模型持续预训练场景
Muon近期成为大规模预训练中AdamW的有力替代,其核心是通过牛顿-舒尔茨(NS)迭代实现参数矩阵正交化。现有变体对所有参数矩阵采用统一的NS调度,忽略了不同矩阵正交化难度的差异及其对性能的影响。通过系统性实证研究,我们发现这种异质性普遍存在,主要由矩阵几何结构决定,且随操作类型、训练阶段和网络深度动态变化。因此,统一调度会导致模型内正交化质量不均。为此,我们提出自适应Muon正交化(AMO),一种先观察后决策的方法:早期根据操作类型测量权重几何特征,并据此分配后续训练中的NS计算预算。AMO在标准、延长及持续预训练任务中均表现稳健,相比统一调度的Muon,在Llama3.1-1.4B上平均下游性能提升+0.76,在Qwen3-1.7B上提升+0.51,覆盖12个评估任务。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。