(a) In this experiment we compute 5, 10 and 20 PCA components for the Buzz dataset. We compared against sklearn which uses the Lanczos algorithm through the fast Fortran-based ARPACK numerical library. In this case, CoLA uses randomized SVD [35] with a rank 3000 approximation.
(b) In this experiment we fit a Ridge regression on the Song dataset with a regularization coefficient set to 0.1. We compared against sklearn using their fastest least-square solver lsqr with a tolerance of  $10^{-4}$ . In this case, CoLA uses CG with the same tolerance and with a maximum number of iterations set to 1K. Additionally, we ran CoLA using CPU and GPU whereas we used only CPU for sklearn as it has no GPU support. We observe how in the arguably most popular ML method, CoLA is able to beat a leading package such as sklearn.
(c) In this experiment we fit a GP with a RBF kernel on two datasets: Elevators and Kin40K. We only used up to 20K observations from Kin40K as that was the maximum number of observations that would fit the GPU memory without needing to partition the MVMs. We compare against GPyTorch which uses CG and stochastic Lanczos quadrature (SLQ) to compute and optimize the negative log-marginal likelihood (loss function). Both experiments were run on a GPU for 100 iterations using Adam as an optimizer with learning rate of 0.1 with the default values of  $\beta_{1} = 0.9$  and  $\beta_{2} = 0.999$ . Additionally, for both GPyTorch and CoLA, the CG tolerance was set to  $10^{-4}$  with a maximum number of CG iterations of 250 and 20 probes were used for SLQ. Note that both CoLA and GPyTorch have similar throughputs, for example GPyTorch runs a 100 iterations on Elevators on 43 seconds whereas CoLA runs a 100 iterations on 49 seconds. When training a GP, we solve a block of 11 linear systems (1 based on y and 10 based on random probes) where one key difference is that the CG solver for GPyTorch has a stopping criteria based on the convergence of the mean solves whereas CoLA has a stopping criteria based on the convergence of all the solves.
(d) In this experiment we run spectral clustering on the cit-HepPh dataset using an embedding size of 8 and also 8 clusters for k-means (with only 1 run of k-means after estimating the embeddings). We compare against sklearnnn using two different solvers, one based on Lanczos iterations using ARPACK and another using an Algebraic Multi-Grid solver AMG. In this case, CoLA also uses Lanczos iterations with a default tolerance of  $10^{-6}$ . We see how sklearnnn's AMG solver runs faster than CoLA's but this is mostly the algorithmic constants as they have similar asymptotical behavior (similar slopes).
(e) In this experiment we solve the Schrödinger equation to find the energy levels of the hydrogen atom on a 3-dimensional finite difference grid with up to  $N = 5\mathrm{K}$  points. In order to handle the infinite spatial extent, we compactify the domain by applying the arctan function. Under this change of coordinates, the Laplacian has a different form, and hence the matrix forming the discretized Hamiltonian is no longer symmetric. We compare against SciPy's Arnoldi implementation with 20 iterations where CoLA also uses Arnoldi with the same number of iterations. Surprisingly, CoLA's JAX jitted code has a competitive runtime when compared to SciPy's runtime using ARPACK.
(f) In this experiment we solve a minimal surface problem on a grid of maximum size of  $N = 100^2$  points. To solve this problem we have to run Netwon-Rhapson where each