Seonglae ChoAs the core subroutine in the Muon optimizer, the requirements changed substantially. In deep learning, GPU-friendly, high-throughput computation matters more than ~16-digit high precision, and the method should use only matrix–matrix products with minimal memory overhead. However, the existing Newton–Schulz iteration converges very slowly at the start, and the heuristic polynomials proposed by Jordan and You quickly reach error ≈ 0.3 but fail to converge at all.
Muon’s update takes the form
where is the layer’s weight matrix, is a momentum estimate of the gradient, is the learning rate, and is the semi-orthogonal matrix closest to , obtained from the SVD . The authors approximate via a composition of odd polynomials, , where each is chosen to be the minimax (worst-case, ) optimal approximation to the sign function. The key contribution is that solving this composition optimization greedily, one stage at a time, still yields the global optimum (Theorem 3.1). Each stage solves
where is the set of odd polynomials of degree at most , and is the lower/upper interval containing the singular values at stage . The next interval is updated by and . With this, the composed polynomial matches the optimal solution that minimizes the worst-case spectral-norm error in Eq. (5).
The optimal polynomial at each iteration is computed offline in float64 using the Remez algorithm based on the Equioscillation Theorem (or in closed form for ). In the online phase, the precomputed coefficients are reused in bfloat16, so this computation is performed only once regardless of how many input matrices are processed. For deep learning, they recommend , –6, and , and they normalize the input as . For bfloat16 stability, they scale each polynomial as , and in the initial iterations they use slightly suboptimal polynomials that oscillate less.
The Polar Express: Optimal Matrix Sign Methods and Their...
Computing the polar decomposition and the related matrix sign function has been a well-studied problem in numerical analysis for decades. Recently, it has emerged as an important subroutine within...
https://arxiv.org/abs/2505.16932
