The Newton-Schulz (NS) iteration has gained increasing interest for its role in the Muon optimizer and the Stiefel manifold. However, the conventional NS iteration suffers from inefficiency and instability. Although various improvements have been introduced to NS iteration, they fail to deviate from the conventional iterative paradigm, which could increase computation burden largely due to the matrix products along the long dimension repeatedly.

UNSO: A Unified Approach

To address this, a new study introduces a unified framework, named Unified Newton-Schulz Orthogonalization (UNSO). This approach consolidates the iterative structure, avoiding a polynomial expansion. Instead, it evaluates the role of each matrix power, removes the insignificant terms, and provides a recommended polynomial with learnable coefficients. These learnable coefficients are then optimized, and achieve an outstanding performance with stable convergence.

The code of our method is available on GitHub: https://github.com/greekinRoma/Unified_Newton_Schulz_Orthogonalization.