Fast Differentiable SVD on GPU via Polar Decomposition

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: this https URL.

Numerical Analysis (math.NA); Machine Learning (cs.LG)

arXiv:2609.32505 [math.NA]

(or arXiv:2609.32505v1 [math.NA] for this version)

https://doi.org/10.48550/arXiv.2609.32505 Focus to learn more arXiv-issued DOI via DataCite (pending registration)

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.