Turbo-Muon: Accelerating Orthogonality-Based Optimization with Pre-Conditioning (2512.04632v1)

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.

• The paper shows Turbo-Muon achieves lower polar error by using AOL preconditioning to reduce Newton-Schulz iterations for faster convergence.
• It introduces a column-wise rescaling technique that preserves spectral properties, significantly reducing the computational overhead for large matrices.
• Empirical results demonstrate a 20% reduction in iteration count and 5–10% training speedup across language and vision tasks with negligible bias.

Turbo-Muon: Preconditioning for Accelerated Orthogonality-Based Optimization

Background and Motivation

Orthogonality-based optimizers, most notably Muon, have become central to large-scale deep learning due to their capacity for improved convergence stability, conditioning, and scale-invariance compared to traditional methods. Muon and its derivatives operate by projecting parameter updates onto the Stiefel manifold, ensuring updates are as close as possible to the space of orthogonal matrices. This orthogonalization, while beneficial, imposes a computational burden: state-of-the-art iterative techniques like Newton-Schulz require several rounds of dense matrix multiplications, especially challenging as model and update sizes increase.

The computational overhead—directly proportional to the number of Newton-Schulz steps required for prescribed polar error—remains the chief bottleneck toward broader adoption. While fast implementations (e.g., with custom Triton kernels) and improvements to polynomial coefficients have marginally improved runtime, further algorithmic innovation is needed for scaling to large and heterogeneous matrices common in both vision and LLM architectures.

Algorithmic Contributions: AOL Preconditioning

The core contribution of Turbo-Muon is the introduction of "Almost Orthogonal Layer" (AOL) preconditioning preceding Newton-Schulz iterations. Classic implementations normalize the input matrix with its Frobenius norm, securing convergence but failing to exploit matrix structure; this normalization tends to over-shrink the spectral norm, slowing convergence.

AOL preconditioning applies a column-wise rescaling to the input matrix, derived from the row sums of its Gram matrix, which not only ensures spectral norm constraints but also systematically improves matrix conditioning. The key formulation is:

$$\text{AOL}(X) = X \cdot \text{diag}\left(\sum_j |X^\top X|_{(i,j)}\right)^{-1/2}$$

Through this, the normalized matrix is closer to the polar factor by construction. Notably, this adjustment significantly reduces the initial polar error, particularly for large matrices drawn from real or heavy-tailed distributions (Figure 1).

Figure 1: AOL preconditioning consistently outperforms Frobenius normalization with increasing matrix size, substantially reducing initial polar error.

By caching intermediate matrix multiplications between the preconditioning and the first Newton-Schulz step, the additional computational cost introduced by AOL can be rendered negligible.

Convergence Acceleration and Runtime Trade-Offs

The practical upshot of AOL preconditioning is a marked reduction in the number of Newton-Schulz iterations needed to achieve a target polar error. Empirical results demonstrate that for random, normal, or heavy-tailed matrices, AOL preconditioning enables the safe removal of one Newton-Schulz iteration without compromising the orthogonalization quality achieved by five iterations in the baseline.

Figure 2: AOL preconditioning enables faster convergence in Newton-Schulz iterations, particularly benefiting large matrices by allowing an iteration to be dropped while improving or maintaining error bounds.

This improvement has direct consequences on runtime. With each iteration dominantly $\mathcal{O}(n^3)$ in cost, saving even a single iteration (from five to four) equates to a 20% reduction in compute. Enhanced Triton kernels that exploit matrix symmetry further compound runtime speedups. The composite effect is a step-time reduction of up to $2.8\times$ relative to classic Muon, with actual training throughput improvements of 5–10% in realistic large batch settings (see Figure 3).

Figure 3: Preconditioning using AOL drastically reduces the computation–polar error Pareto frontier compared to prior Newton-Schulz orthogonalization, lowering optimizer overhead.

Numerical Evidence and Robustness

Turbo-Muon’s numerical superiority persists across distributions (Gaussian and heavy-tailed Levy) and matrix sizes. The polar error of Turbo-Muon, following four iterations, matches or outperforms five-iteration Muon in both cases. Robustness across various data modalities, including simulated gradient distributions with heavy tails (typical in deep models), was specifically validated.

Furthermore, Turbo-Muon's induced bias—its tendency to converge to the polar factor of the preconditioned matrix rather than the original—was decomposed into irreducible bias and residual approximation error. Analysis (Figure 4) established that, in practically relevant iteration regimes, gains from faster convergence overwhelmingly dominate any additional bias, which remains negligible.

Figure 4: Decomposition of Turbo-Muon polar error showcases a higher-order bias term induced by AOL preconditioning, which remains much smaller than the convergence-induced approximation error at typical iteration counts.

Downstream Training Impacts

Turbo-Muon’s improvements extend meaningfully to end-to-end deep learning workloads:

• Language Modeling (NanoGPT): Drop-in replacement with Turbo-Muon in highly-optimized NanoGPT training yields 8–10% improvements in total training wall-clock time, with no statistically significant change in final loss or convergence profile.
• Computer Vision (CIFAR-10): Turbo-Muon yields comparable accuracy and modest runtime reductions even on short training tasks with highly optimized convolutional workflows.

These runtime improvements are especially pronounced at intermediate batch sizes and model scales, where optimizer overhead is non-trivial but full sharding is not yet in effect.

Theoretical Considerations and Gradient Dynamics

A potential concern with bias incurred by AOL preconditioning is whether it could degrade optimization trajectories. The analysis confirms that Turbo-Muon updates always provide a strict descent direction in an induced norm, even as iteration count increases and the iterative error vanishes, ensuring stability and convergence of the training process.

Moreover, preconditioning insulates against feature redundancy by down-weighting gradient directions with high structural coherence—implicitly adapting the optimization geometry and aligning it with spectral constraints (i.e., steepest descent in a rescaled norm).

Implications and Future Directions

Turbo-Muon demonstrates that careful, structure-aware preconditioning dramatically accelerates orthogonality-based optimization with negligible or positive impact on numerical accuracy and training quality. The approach is hyperparameter-free, compatible with current Muon pipelines, and enables further system-level scaling by reducing one of the few remaining algorithmic bottlenecks in LLM and vision model training.

There is a clear path open for:

• Refinement of polynomial coefficients dynamically tuned for preconditioned schemes,
• Extension of preconditioning strategies for distributed and mixed-batch settings,
• Broader paper of correlation-sensitive norms and their impact on feature learning and optimization geometry.

Conclusion

Turbo-Muon, via AOL preconditioning, presents a robust, efficient advancement in orthogonality-based optimization for deep learning. Empirical and theoretical analysis confirm that the method:

• Achieves lower polar error with fewer Newton-Schulz iterations,
• Consistently improves runtime in both language and vision training regimes,
• Maintains or even improves downstream accuracy and convergence,
• Introduces a negligible bias with strict theoretical guarantees of stable descent.

The simplicity of adoption and breadth of applicability position Turbo-Muon as a practical standard for scalable orthogonality-enforced optimization in modern neural network training workflows.