Turbo-Muon: Accelerating Orthogonality-Based Optimization with Pre-Conditioning

2025-12-04

AI Summary

Orthogonal optimizers (e.g., Muon) suffer from high computational overhead in the gradient orthogonalization step, particularly due to the Newton–Schulz (N-S) iteration, which requires dozens of matrix multiplications for convergence. To address this, we propose a hyperparameter-free preconditioning mechanism that accelerates N-S iteration out-of-the-box: by constructing structured preconditioning matrices, we significantly improve convergence rate—reducing iterations by just one while preserving orthogonal accuracy. Our method integrates seamlessly into existing training pipelines without modifying model architectures or optimizer backbones. Experiments across language and vision tasks demonstrate end-to-end training speedups of 5–10%, up to 2.8× faster N-S approximation, and maintained or improved model performance. The core contribution is the first integration of tuning-free preconditioning into orthogonal iterative refinement, achieving a favorable trade-off among efficiency, orthogonality precision, and deployment simplicity.

Technology Category

• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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 preconditioning procedure that accelerates Newton-Schulz convergence and reduces its computational cost. We evaluate its impact and show that the overhead of our preconditioning can be made negligible. Furthermore, the faster convergence it enables allows us to remove one iteration out of the usual five without degrading approximation quality. Our publicly available implementation achieves up to a 2.8x speedup in the Newton-Schulz approximation. We also show that this has a direct impact on end-to-end training runtime with 5-10% improvement in realistic training scenarios across two efficiency-focused tasks. On challenging language or vision tasks, we validate that our method maintains equal or superior model performance while improving runtime. Crucially, these improvements require no hyperparameter tuning and can be adopted as a simple drop-in replacement. Our code is publicly available on github.

Problem

• Accelerates Newton-Schulz convergence for gradient orthogonalization
• Reduces computational cost of orthogonality-based optimization methods
• Improves training runtime without degrading model performance

Innovation

• Preconditioning accelerates Newton-Schulz convergence
• Reduces computational cost with negligible overhead
• Enables faster training without hyperparameter tuning