Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

UNSO: Unified Newton Schulz Orthogonalization

Chen Hu, Qianxi Zhao, Yuming Li, Mingyu Zhou, Xiyin Li

TL;DR

This work introduces Unified Newton-Schulz Orthogonalization (UNSO), a non-iterative framework that replaces repeated NS steps with a single learned polynomial to compute a matrix's orthogonal factor efficiently. By transforming inputs to minimize $A^\top A$ and applying a unified polynomial on the singular values, UNSO reduces long-dimension multiplications and uses learnable coefficients to stabilize convergence toward the polar factor $UV^\top$. The method leverages an exponential parameterization of term powers $n_k$ and derives closed-form coefficient constraints, achieving competitive accuracy with lower computational cost than traditional NS-based methods. Empirical results show improved efficiency, favorable ablations, and robust curve-based and practical-OP evaluations, highlighting UNSO's potential for large-scale orthogonalization tasks in RM/Riemannian optimization and related neural-network applications.

Abstract

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. To address this, we consolidate the iterative structure into a unified framework, named Unified Newton-Schulz Orthogonalization (UNSO). To do so, we could avoid a polynomial expansion. Instead, we evaluate the role of each matrix power, remove the insignificant terms, and provide 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: https://github.com/greekinRoma/Unified_Newton_Schulz_Orthogonalization.

UNSO: Unified Newton Schulz Orthogonalization

TL;DR

This work introduces Unified Newton-Schulz Orthogonalization (UNSO), a non-iterative framework that replaces repeated NS steps with a single learned polynomial to compute a matrix's orthogonal factor efficiently. By transforming inputs to minimize and applying a unified polynomial on the singular values, UNSO reduces long-dimension multiplications and uses learnable coefficients to stabilize convergence toward the polar factor . The method leverages an exponential parameterization of term powers and derives closed-form coefficient constraints, achieving competitive accuracy with lower computational cost than traditional NS-based methods. Empirical results show improved efficiency, favorable ablations, and robust curve-based and practical-OP evaluations, highlighting UNSO's potential for large-scale orthogonalization tasks in RM/Riemannian optimization and related neural-network applications.

Abstract

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. To address this, we consolidate the iterative structure into a unified framework, named Unified Newton-Schulz Orthogonalization (UNSO). To do so, we could avoid a polynomial expansion. Instead, we evaluate the role of each matrix power, remove the insignificant terms, and provide 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: https://github.com/greekinRoma/Unified_Newton_Schulz_Orthogonalization.
Paper Structure (14 sections, 26 equations, 6 figures, 1 table, 1 algorithm)

This paper contains 14 sections, 26 equations, 6 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Method
  3. Overall Structure
  4. Unified Polynomial
  5. Term Selection
  6. Coefficient Optimization
  7. Matrix Iteration
  8. Experiment
  9. Experimental Setup
  10. Ablation Study
  11. Comparative Experiments
  12. Curve-Based Comparative Evaluation
  13. Comparative Evaluation of Practical Matrix Operations
  14. Conclusion

Figures (6)

  • Figure 1: The existing NS iteration methods.
  • Figure 2: The overall of the unified NS iteration structure.
  • Figure 3: Different parameter $n_k$ pattern.
  • Figure 4: Family of Curves with Increasing $k$. (a) Extreme point values of $x*$ and $y*$; (b) The normalized $y$.
  • Figure 5: The optimized curve with the different $N$.
  • ...and 1 more figures