Fast Differentiable SVD on GPU via Polar Decomposition
math.NA, cs.LG, cs.NA
Submitted: 2026-09-26
Updated: 2026-09-26
Code: https://github.com/fallnlove/cans_svd
Project page: https://j-towns.github.io/papers/svd-derivative.pdf
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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.
Terminology
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.
Sources
- Accelerating Newton-Schulz Iteration for Orthogonalization via Chebyshev-type Polynomials
- Pointer Sentinel Mixture Models
- Llama 2: Open Foundation and Fine-Tuned Chat Models
- Differentiable SVD based on Moore-Penrose Pseudoinverse for Inverse Imaging Problems
Related papers
- Do physics-informed neural networks (PINNs) need to be deep? Shallow PINNs using the Levenberg-Marquardt algorithm
- A Neural-preconditioned Poisson Solver for Mixed Dirichlet and Neumann Boundary Conditions
- Operator learning meets inverse problems: A probabilistic perspective
- Second-order consistency for learning chaotic dynamics via randomized Jacobian matching
- Windowed thinning and query complexity for the bouncy particle and Zigzag samplers
- Data-efficient Kernel Methods for Learning Hamiltonian Systems