arXiv:2512.04632cs.AI2025-12

通过近正交预处理加速穆翁优化器,显著降低训练时间。

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates

  • 引入近正交预处理,优化牛顿-舒尔迭代初始状态。
  • 减少5次迭代中1次,训练时间缩短约3%。
  • 无需调参,可直接替换现有优化器,适合大模型训练者。

基于正交性的优化器(如Muon)在大规模训练中表现优异,但依赖昂贵的梯度正交化步骤。即使高效的牛顿-舒尔迭代也需数十次矩阵乘法才能收敛。本文提出一种预处理方法,在几乎无额外开销下改善牛顿-舒尔迭代的初始条件,降低初始极性误差,并可移除一次迭代(从五次减为四次)。该改进显著降低Muon的运行开销。端到端训练中,多个快速训练基准显示约3%的训练时间缩减,且在语言与视觉任务上保持参考性能。关键优势在于无需超参数调优,可作为简单即插即用替换。此外,理论分析揭示了更新机制对特征坍缩的潜在鲁棒性。代码已开源,支持optax与Hugging Face内核。

原文摘要 · Abstract (English)

Orthogonality-based optimizers, such as Muon, have recently shown strong performance across large-scale training and community-driven efficiency challenges. However, these methods rely on a costly gradient orthogonalization step. Even efficient iterative approximations such as Newton-Schulz remain expensive, typically requiring dozens of matrix multiplications to converge. We introduce a pre-conditioning procedure that improves the initialization of the Newton--Schulz iterations while incurring negligible overhead. Furthermore, our pre-conditioning reduces the initial polar error and enables the removal of one Newton-Schulz iteration (out of the five iterations usually used in practice). The resulting implementation significantly reduces Muon's overhead. At the end-to-end training level, we observe consistent runtime improvements across speed-run and standard benchmarks, including $\sim$3% reductions in training time on multiple fast training benchmarks, while matching reference performance on both language and vision tasks. Crucially, these improvements require no hyperparameter tuning and can be adopted as a simple drop-in replacement. Beyond empirical gains, we provide theoretical insight into the geometry of the update and its potential robustness against feature collapse. Our code is publicly available on github, in optax and huggingface kernels.

优化器正交化训练加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。