SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 1 Arnoldi iteration
1: Inputs:  $\mathbf{A}$ $\mathbf{q}_0 = \nu_0 / \| \nu_0\|$  where possibly  $\nu_{0}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$  maximum number of iterations  $T$  and tolerance  $\epsilon \in (0,1)$    
2: for  $j = 0$  to  $T - 1$  do   
3:  $\pmb{\nu}_{j + 1}\gets \mathbf{A}\mathbf{q}_j$    
4: for  $i = 0$  to  $j$  do   
5:  $h_{i,j} = \mathbf{q}_i^* (\mathbf{A}\mathbf{q}_j)$    
6:  $\pmb{\nu}_{j + 1}\leftarrow \pmb{\nu}_{j + 1} - h_{i,j}\mathbf{q}_i$    
7: end for   
8:  $h_{j + 1,j} = \| \pmb{\nu}_{j + 1}\|$    
9: if  $h_{j + 1,j} <   \epsilon$  then   
10: stop   
11: else   
12:  $\mathbf{q}_{j + 1} = \pmb{\nu}_{j + 1} / h_{j + 1,j}$    
13: end if   
14: end for   
15: return H,  $\mathbf{Q} = (\mathbf{q}_0|\dots |\mathbf{q}_{T - 1}|\mathbf{q}_T)$
Algorithm 2 Householder Arnoldi iteration
1: Inputs: A,  $\nu_0 \neq 0$  where possibly  $\nu_0 \sim \mathcal{N}(0, \mathbf{I})$ , and maximum number of iterations  $T$ .  
2: for  $j = 0$  to  $T$  do  
3:  $\mathbf{u}_j = \text{GET\_HOUSEHOLDER\_VEC}(\nu_j, j)$   
4:  $\mathbf{R}_j = \mathbf{I} - 2\mathbf{u}_j\mathbf{u}_j^*$   
5:  $\mathbf{h}_j = \mathbf{R}_j\nu_j$   
6:  $\mathbf{q}_j = \mathbf{R}_0 \cdots \mathbf{R}_j\mathbf{e}_{j+1}$   
7: if  $j < T$  then  
8:  $\nu_{j+1} = \mathbf{R}_j \cdots \mathbf{R}_0(\mathbf{A}\mathbf{q}_j)$   
9: end if  
10: end for  
11: return H, Q = ( $\mathbf{q}_0\| \ldots |\mathbf{q}_T$ )  
12: function GETHOUSEHOLDER\_VEC(w, k)  
13:  $u_i = 0$  for  $i < k$  and  $u_i = w_i$  for  $i > k$ .  
14:  $u_k = w_k - \| \mathbf{w}\|$   
15: return u  
16: end function
Figure 7: Different versions of the same algorithm, but the Householder variant being more numerically robust.
2. Kin40K. The full name of this UCI dataset is Statlog (Shuttle) Data Set. This dataset contains information about NASA shuttle flights and we used a subset that consists of  $N = 40\mathrm{K}$  observations and has  $D = 8$  dimensions.
3. Buzz. The full name of this UCI dataset is Buzz in social media. This dataset consists of examples of buzz events from Twitter and Tom's Hardware. We used a subset consisting of  $N = 430\mathrm{K}$  observations and has  $D = 77$  dimensions.
4. Song. The full name of this UCI dataset is YearPredictionMSD. This dataset consists of  $N = 386.5\mathrm{K}$  observations and it has  $D = 90$  audio features such as 12 timbre average features and 78 timbre covariance features.
5. cit-HepPh. This dataset is based on arXiv's HEP-PH (high energy physics phenomenology) citation graph and can be found here: https://snap.stanford.edu/data/cit-HepPh.html. The dataset covers all the citations from January 1993 to April 2003 of  $|V| = 34,549$  papers, ultimately containing  $|E| = 421$ , 578 directed edges. The notion of relationship that we used in our spectral clustering experiment creates a connection between two papers when at least one cites another (undirected symmetric graph). Therefore the dataset that we used has the same number of nodes but instead  $|E| = 841$ , 798 undirected edges.
D.2 Compositional experiments
This section pertains to the experiments of Section 3.2 displayed in Figure 1. We now elaborate on each of Figure 1's panels.
(a) The multi-task GP problem exploits the structure of the following Kronecker operator  $\mathbf{K}_T\otimes \mathbf{K}_X$ , where  $\mathbf{K}_T$  is a kernel matrix containing the correlation between the tasks and  $\mathbf{K}_X$  is a RBF kernel on the data. For this experiment, we used a synthetic Gaussian dataset where the train data  $\mathbf{x}_i\sim \mathcal{N}(\mathbf{0},\mathbf{I}_D)$  which has dimension  $D = 33$ ,  $N = 1\mathrm{K}$  and we used  $T = 11$  tasks (where the tasks basically set the size of  $\mathbf{K}_T$ ). We used conjugate gradients (CG) as the iterative method, where we set the hyperparameters to a tolerance of  $10^{-6}$  and to a maximum number of iterations to  $1\mathrm{K}$ . We used the exact same hyperparameters for CoLA.
(b) For the bi-poisson problem we set up the maximum grid to be  $N = 1000^2$ . Since this PDE problem involves solving a symmetric linear system, we used CG as the iterative method with a tolerance of  $10^{-11}$  and a maximum number of iterations of  $10\mathrm{K}$ . The previous
parameters also apply for CoLA. We note that PDE problems are usually solved to higher tolerances as the numerical error compounds as we advance the PDE.
(c) For the EMLP experiment we consider solving the equivariance constraints to find the equivariant linear layers of a graph neural network with 5 nodes. To solve this problem, we need to find the nullspace of a large structured constraint matrix. We use the uniformly channel heuristic from [16] which distributes the  $N$  channels across tensors of different orders. We consider our approach which exploits the block diagonal structure, separating the nullspaces into blocks, as opposed to the direct iterative approach exploiting only the fast MVMs of the constraint matrix. We use a tolerance of  $10^{-5}$ .
D.3 Sum structure experiments
This section pertains to the experiments of Section 3.3 contained in Figure 2. We now elaborate on each of Figure 2's panels.
(a) In this experiment we computed the first principal component of the  $B u zz$  dataset. For the iterative method we used power iteration with a maximum number of iterations of 300 and a stop tolerance of  $10^{-7}$ . CoLA used SVRG also with the same stop tolerance and maximum number of iterations. Additionally, we set SVRG's batch size to  $10\mathrm{K}$  and the learning rate to 0.0008. We note that a single power iteration roughly contains  $43 / 2 = 21.5$  times more MVMs than a single iteration of SVRG. In this particular case, the length of the sum is given by the number of observations and therefore SVRG uses  $430 / 10 = 43$  times less elements per iteration, where 10 comes from the  $10\mathrm{K}$  batch size. Finally, the 2 is explained by noting that SVRG incurs in a full sum update on every epoch.
(b) In this experiment we trained a GP by estimating the covariance RBF kernel with  $J = 1\mathrm{K}$  random Fourier features (RFFs). The hyperparameters for the RBF kernel are the following: length scale  $(\ell = 0.1)$ , output scale  $(a = 1)$  and likelihood noise  $(\sigma^2 = 0.1)$ . Moreover, we used CG as the iterative solver with a tolerance of  $10^{-8}$  and 100 as the maximum number of iterations (the convergence took much less iterations than the max). For SVRG we used the same tolerance but set the maximum number of iterations to 10K, a batch size of 100 and learning rate of 0.004. We note that a single CG iteration roughly contains  $10/2 = 5$  times more MVMs than a single iteration of SVRG. In this particular case, the length of the sum is given by the number of RFFs and therefore SVRG uses  $1000 / 100 = 10$  times less elements per iteration, where 100 comes from the batch size.
(c) In this experiment we implemented the Neural-IVP method from Finzi et al. [17]. We consider the time evolution of a wave equation in two spatial dimensions. At each integrator step, a linear system  $\mathbf{M}(\theta)\dot{\theta} = F(\theta)$  must be solved to find  $\dot{\theta}$ , for a  $d = 12\mathrm{K} \times 12\mathrm{K}$  dimensional matrix. While Finzi et al. [17] use conjugate gradients to solve the linear system, we demonstrate the advantages of using SVRG, as  $\mathbf{M}(\theta) = \frac{1}{m}\sum_{i=1}^{m}M_i(\theta)$  is a sum over the evaluation at  $m = 50\mathrm{K}$  distinct sample locations within the domain. In this experiment we use a batch size of 500 for SVRG, and employ rank 250 randomized Nyström preconditioning for both SVRG and the iterative CG baseline.
D.4 Hardware speed-up comparisons
This section pertains to the experiments of Figure 3. For all these experiments we computed the runtime reduction as a fraction between the time that it takes CoLA to run some linear algebra operation and PyTorch using the same hardware. As an example, assume that PyTorch takes 200 seconds to compute a solve using a CPU and 100 seconds to compute the same solve but now using a GPU. Moreover, assume that CoLA's iterative algorithm takes 100 seconds to compute the same solve on a CPU and 40 seconds on a GPU. Thus, the runtime reduction would be  $\frac{100}{200} = 0.5\%$  for the CPU column whereas  $\frac{40}{100} = 0.4$  for the GPU column.
1. Solves. In this experiment we calculated the  $\%$  runtime reduction when running torch.linalg.solve on the Trefethen  $N = 20K$  matrix market sparse operator. In this experiment, CG was run with a tolerance of  $10^{-11}$  and a maximum number of iterations equal to the operator size.
2. Eigenvalue estimation. In this experiment we calculated the  $\%$  runtime reduction when running torch.linalg.eigh on the mhd4800b  $N = 4.8K$  matrix market sparse operator. In this experiment, Lanczos was run with a tolerance of  $10^{-9}$  and a maximum number of iterations equal to 100.
3. Log determinant computation. In this experiment we calculated the  $\%$  runtime reduction when running torch.linalg.logdet on the bcstsk18  $N = 11.9K$  matrix market sparse operator. In this experiment, the stochastic Lanczos quadrature was run using 30 Lanczos probe estimates and 25 samples.
D.5 Applications
This section pertains to the experiments of Section 4 displayed in Figure 4. We now elaborate on each of Figure 4's panels.
inner step involves a linear solve of an non-symmetric operator. We compare against SciPy's GMRES implementation as well as JAX's integrated version of SciPy. The main difference between the two is that SciPy calls the fast and highly-optimized ARPACK library whereas SciPy (JAX) has its only Python implementation of GMRES which only uses JAX's primitives (equally as it is done in CoLA). The tolerance for this experiment was 5e-3. We see how CoLA's GMRES implementation is competitive with SciPy (JAX) but it still does not beat ARPACK mostly due to the faster runtime of using a lower level GMRES implementation.