Fast Differentiable SVD on GPU via Polar Decomposition

2026-09-26

AI Summary

This study addresses the computational inefficiency and lack of end-to-end differentiability of singular value decomposition (SVD) on GPUs. We propose a GPU-friendly, fully differentiable SVD framework based on polar decomposition. By leveraging Newton-Schulz iterations, the method reformulates SVD into efficient matrix multiplication operations, and we further derive a numerically stable backpropagation algorithm to ensure full differentiability. As the first purely GPU-oriented polar decomposition SVD pipeline, our approach achieves up to a twofold speedup over standard implementations. High-performance open-source implementations are provided in both PyTorch and JAX, substantially accelerating SVD-related computations in deep learning applications.

Abstract

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.

Problem

• Singular Value Decomposition
• GPU acceleration
• Differentiable SVD
• Polar decomposition

Innovation

• Differentiable SVD
• Polar Decomposition
• GPU Acceleration
• Newton-Schulz Iteration
• Numerical Stability