Fast and differentiable matrix decomposition method speeds up GPU work
Fast Differentiable SVD on GPU via Polar Decomposition
Machine Learning
Summary
Breaking down large tables of numbers into simpler parts is important in many technology areas but can be slow and hard to reverse for learning. The authors show a way to do this breakdown on GPUs using a method called polar decomposition, which mostly uses multiplication and is faster than usual methods. They also figured out how to reliably work backward through this process, making it easier to use in training computer models. Their approach runs up to twice as fast and has been made available as open-source code in popular programming tools.
What this means in practice
- •For machine learning engineers: Speed up training and improve gradient computation for models that use matrix decompositions on GPUs.
- •For graphics programmers: Accelerate computations involving matrix decompositions in real-time graphics and simulations using GPU-optimized routines.
Authors
Uliana Parkina, Askar Tsyganov, Sergei Kudriashov, Sergey Samsonov, Maxim Rakhuba
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.