给穆翁优化器加了变率自适应,提速超30%。
Variance-Adaptive Muon: Pre-Orthogonalization Variance Modulation for Efficient Language Model Pretraining
- 在穆翁优化器前加入噪声比调制或方差缩放,实现变率自适应。
- 125M到12亿参数模型上,速度比优化后的基线快1.33倍。
- 只需一个额外缓冲区,适合大规模语言模型高效训练。
优化器设计在高效语言模型预训练中起关键作用,直接影响优化动态、收敛速度和固定预算下的计算成本。穆翁(Muon)通过正交化动量更新成为强效优化器,形成符号归一化的矩阵形式。但与Adam类方法不同,穆翁未显式利用梯度方差信息。受Adam的变率自适应启发,本文提出两种变率自适应穆翁变体:Muon-NSR在牛顿-舒尔正交化前应用噪声-信号比(NSR)调制;Muon-VS则采用方差缩放(VS),不引入额外超参数。两者均保持穆翁的谱归一化结构,仅需一个额外方差缓冲区。在从125M到1.2B参数的Llama风格和GPT-2预训练实验中,新方法优于调优后的穆翁基线,并与主流自适应穆翁家族方法相当。在Llama-1.2B上,Muon-VS相比调优基线实现1.33倍的步数到目标加速,以穆翁最终验证损失为目标。结果表明,变率自适应调制是提升穆翁型优化器效率的简单有效机制。
原文摘要 · Abstract (English)
Optimizer design plays a central role in efficient language model pretraining, directly affecting optimization dynamics, convergence speed, and compute cost under fixed training budgets. Muon has emerged as a strong optimizer by orthogonalizing momentum updates, yielding a matrix-valued analogue of sign-based normalization. However, unlike Adam-style methods, Muon does not explicitly incorporate gradient-variance information into its updates. Motivated by Adam's variance-adaptive interpretation, we propose Muon-NSR and Muon-VS, two variance-adaptive Muon variants for language model pretraining. Muon-NSR applies noise-to-signal ratio (NSR) modulation before Newton--Schulz orthogonalization, whereas Muon-VS uses variance scaling (VS) without introducing any additional hyperparameters beyond those of Muon. Both methods preserve Muon's spectral normalization structure while requiring only one additional variance buffer. Experiments on Llama-style and GPT-2 pretraining across model scales from 125M to 1.2B parameters show that our methods improve over well-tuned Muon baselines and remain competitive with representative adaptive Muon-family baselines. On Llama-1.2B, Muon-VS achieves a 1.33$\times$ step-to-target speedup over a well-tuned Muon baseline, with Muon's final validation loss as the target. These results indicate that variance-adaptive modulation is a simple and effective mechanism for improving Muon-style optimizers in language model pretraining.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。