preprint Open access

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

  • arXiv (Cornell University)
  • Cornell University
Research footprint

At a glance

Citations
0
References
0
Comments
0
Paper overview

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 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.

Record transparency

Publication details

OpenAlex
W4417008095
Document type
preprint
Language
EN
Source
arXiv (Cornell University)
Last metadata update
Community

Comments

Log in to join the discussion.

  1. No comments yet. Start the discussion.