通过极分解在GPU上实现快速可微SVD
Fast Differentiable SVD on GPU via Polar Decomposition
- HSE University（高等经济大学）
机构由 AI 辅助整理，请以论文原文为准。
AI总结：
本文提出基于极分解的GPU SVD流水线，利用Newton-Schulz迭代实现2倍加速，并推导数值稳定的反向传播以获得完全可微的SVD，开源于PyTorch和JAX。
AI中文摘要：
我们提出了一种完全面向GPU的SVD（奇异值分解）流水线，其基于极分解，并受到仅依赖矩阵乘法的迭代方法（如Newton-Schulz迭代）的启发。我们证明，与标准实现相比，该方法可实现高达2倍的加速。此外，我们推导了极分解的数值稳定的反向传播，并利用它获得完全可微的SVD。我们的方法以PyTorch和JAX两种框架的开源实现形式发布：此https URL。
英文摘要：
We present a fully GPU-oriented SVD pipeline based on polar decomposition, motivated by iterative methods that rely solely on matrix multiplications, such as the Newton-Schulz iteration. We show that this approach enables up to a $2\times$ speedup compared to standard implementations. Furthermore, we derive a numerically stable backward pass for the polar decomposition and leverage it to obtain a fully differentiable SVD. Our methods are released as open-source implementations in both PyTorch and JAX: https://github.com/fallnlove/cans_svd.
