The Newton-Schulz (NS) iteration has become a key technique for orthogonalization in optimizers such as Muon and for optimization on the Stiefel manifold. Despite its effectiveness, the conventional NS iteration incurs significant computational overhead due to repeated high-dimensional matrix multiplications. To overcome these limitations, we propose Iteration-Free Newton-Schulz Orthogonalization (IFNSO), a novel framework that consolidates the traditional iterative structure into a unified and Iteration-Free formulation. By analyzing the contribution of individual matrix powers, we streamline the process by removing insignificant terms and introducing a polynomial with learnable coefficients. These coefficients are optimized to ensure both superior computational efficiency and stable convergence. Extensive experiments demonstrate that IFNSO achieves superior performance compared to existing methods. Our code is available at: https://github.com/greekinRoma/Unified_Newton_Schulz_Orthogonalization.
翻译:牛顿-舒尔茨(NS)迭代已成为Muon等优化器中正交化及Stiefel流形优化的关键技术。尽管其效果显著,传统NS迭代因重复的高维矩阵乘法会产生巨大的计算开销。为克服这些局限,我们提出免迭代牛顿-舒尔茨正交化(IFNSO),这是一种将传统迭代结构整合为统一免迭代形式的新框架。通过分析各矩阵幂次的贡献,我们剔除不显著项并引入带可学习系数的多项式,从而简化流程。这些系数经过优化,以确保卓越的计算效率与稳定的收敛性。大量实验表明,IFNSO相比现有方法实现了更优的性能。代码发布于:https://github.com/greekinRoma/Unified_Newton_Schulz_Orthogonalization。