Turbocharging SVD with JAX
By: Zichen Wang
Tags: Numpy, Optimization, Data Science, Python, Machine Learning

In the previous post, I wrote about the fundamentals of two commonly used dimensionality reduction approaches, singular value decomposition (SVD) and principal component analysis (PCA). I also explained their relationships using numpy. To quickly recap, the singular values (Σ) of a 0-centered matrix X (n samples × m features), equals the square root of its eigenvalues (Λ), making it possible to compute PCA using SVD, which is often more efficient.

In this post, I will explore and benchmark a few SVD algorithms and implementations in Python for squared dense matrix (X, n×n). I also had a very exciting finding: the recently developed JAX library powered by the XLA compiler is able to significantly accelerate SVD computation. Let’s get to it!

How to calculate SVD?

SVD is a common and fundamental technique in exploratory data analysis and Machine Learning, but how do we actually go about figuring out the left and right unitary matrices (U, V*) and the singular values? Through my far from exhaustive literature review, there are at least four algorithms to solve the full SVD for dense matrix. Even more algorithms are available for truncated/reduced SVD for sparse matrix, which only computes the largest k singular values for efficiency. In this post, I will focus on the dense matrices and full SVD algorithms:

Divide-and-conquer approach (GESDD): a two-stage algorithm. It first reduces the input matrix to bidiagonal form via Householder reflection, then divides the bidiagonal matrix into smaller matrices to calculate singular values and the unitary matrices U and V. Finally, the singular values of the entire input matrix is simply the those of the smaller sub-matrices. A more in-depth go-through of this algorithm is available here. This algorithm is conceptually similar to MapReduce, where big chunk of data is divided into pieces for parallel processing, and the results from each pieces can be aggregated using simple functions such as summation and concatenation. General rectangular approach (GESVD): also a two-stage algorithm where the first step is identical to GESDD. The second step can be done using iterative QR decomposition to derive singular values. More mathematical details are available here.

Both algorithms are implemented in LAPACK, a classic linear algebra library written in Fortran. The divide-and-conquer approach is documented to be much faster but takes more memory than the general rectangular approach.

SVD implementations in Python

Scipy and Numpy both contain methods for computing SVD under its respective linalg (linear algebra) submodules:

numpy.linalg.svd: “The decomposition is performed using LAPACK routine `_gesdd`”. However, by checking the source code of this function, it seems that numpy calls LAPACK through some calling macros in written in C. The purpose of such design is probably to circumvent the requirement of a Fortran compiler when building numpy.

scipy.linalg.svd: provides more parameters, such as overwrite_a , check_finite , for potential speed-ups computing SVD. It also has a parameter lapack_driver to choose between GESDD and GESVD algorithms in LAPACK.

, , for potential speed-ups computing SVD. It also has a parameter to choose between GESDD and GESVD algorithms in LAPACK. scipy.sparse.linalg.svds: is to perform truncated SVD, which computes the largest k singular vector for sparse matrices. Although this function is also able to compute full SVD, it has been shown to be much less efficient than scipy.linalg.svd.

You may wonder why is there implementations of linalg.svd in both numpy and scipy. What’s the difference between np.linalg.svd and scipy.linalg.svd ? I found the answer to this in scipy’s FAQ: scipy.linalg is a more complete wrapping of Fortran LAPACK using f2py, whereas numpy is trying to be independent of LAPACK.

Timings for numpy/scipy SVD methods as a function of matrix size n

To compare the speeds of different SVD implementations, I set up a very simple benchmark to measure the execution times of SVD implementations in numpy and scipy by varying sizes of square matrix of size n. As is shown in the figure above, the divide-and-conquer (GESDD) methods in both scipy ( scipy_svd ) and numpy ( np_svd ) are more than a magnitude faster than the QR decomposition approach (GESVD) method ( scipy_gesvd ). The scipy.linalg.svd is also much faster than its numpy counterpart, probably due to its more efficient communication with LAPACK.

JAX and XLA

Let’s take a short detour to JAX, the recently developed open-source library by Google, aiming to accelerate linear algebra computations. JAX is an “unofficial” Google project that is equipped with Autograd and XLA compiler, brought together for high-performance Machine Learning research. XLA (Accelerated Linear Algebra) is the compiler behind Tensorflow to accelerate model computation without source code changes. It works by “fusing” multiple simple operations into a single kernel. JAX implements the XLA compiler in the jit (just-in-time) function to allow explicit compilation of Python functions. Below is a quick tour of simple usages of JAX to show how primitive Python functions can be accelerated with the jit compiler.

Another notable feature of JAX is that XLA-compiled code is automatically cached and can be reused if run repeatedly to further speed up the process. This is demonstrated in the tutorial here.

Moreover, JAX also provides XLA accelerator-backed numpy and scipy functions under the submodules jax.numpy and jax.scipy , although not all numpy/scipy functions are currently supported. Fortunately, linalg.svd for both numpy and scipy has been implemented in JAX.

Benchmarking SVD functions for dense matrix

With the JAX accelerated numpy and scipy, I can compare their performances with their vanilla counterparts. I used the same timing function used before to get the following result.

First attempt in timing JAX powered SVD implementations

At first glance this looks amazing: JAX powered numpy and scipy both beat the naive numpy and scipy implementation of SVD by a lot. But then it looks astonishingly too good to be true: how could JAX powered SVD be four magnitudes faster and almost has a O(1) complexity? After checking my timing function, I found JAX may have found a way to cheat by caching the results from previous runs. To prevent JAX from doing that, I wrote another timing function to generate a completely new random dense matrix at each trial to get the following result, which is more realistic.

Timings of SVD methods with and without JAX acceleration as a function of matrix size n

The above plot shows JAX-powered SVD function in numpy ( jnp_svd ) significantly outperform all other countenders in speed, whereas JAX-powered scipy implementation ( jsp_svd ) does not exhibit any differences to its plain scipy counterpart. It’s worth looking into the implementations more the figure out the reason for that. It might have something to do with the different approaches scipy and numpy take to interact with LAPACK, which may affect their ability to be further optimized by XLA compiler.

It is also necessary to mention that all the benchmarking results are generated from a CPU runtime. I also tried to run them on single GPU and single TPU runtimes through Google Colab. And it seems that GPU doesn’t facilitate the acceleration of jax.numpy whereas TPU generate almost identical results from CPU. These observations also necessitate further investigation.