SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Define matrix  $W\in \mathbb{R}^{I_k\times R}$  elementwise by
$$
W [ t, u ] := \langle h h ^ {\top}, U _ {k} [ t,: ] ^ {\top} U _ {k} [ t,: ], V [ :, u ] V [ :, u ] ^ {\top} \rangle
$$
and observe that all of its entries are nonnegative. Since  $\lambda_{u} \geq 0$  for all  $u$  ( $G_{>k}$  is p.s.d.), we can write  $q$  as a mixture of probability distributions given by the normalized columns of  $W$ :
$$
q = \sum_ {u = 1} ^ {R} w [ u ] \frac {W [ , u ]}{\| W [ , u ] \| _ {1}},
$$
where the vector  $w$  of nonnegative weights is given by  $w[u] = (C^{-1}\lambda_u\| W[:,u]\| _1)$ . Rewriting  $q$  in this form gives us the two stage sampling procedure: first sample a component  $u$  of the mixture according to the weight vector  $w$ , then sample an index in  $\{1..I_k\}$  according to the probability vector
Figure 1: A segment tree  $T_{8,2}$  and probability distribution  $\{q_1,\dots,q_8\}$  on  $[1,\dots,8]$ .
defined by  $W[(:,u] / \| W[(:,u]\| _1$ . Let  $\hat{u}_k$  be a random variable distributed according to the probability mass vector  $w$ . We have, for  $C$  taken from Theorem 3.1,
$$
\begin{array}{l} p \left(\hat {u} _ {k} = u _ {k}\right) = C ^ {- 1} \lambda_ {u _ {k}} \sum_ {t = 1} ^ {I _ {k}} W [ t, u _ {k} ] \tag {9} \\ = C ^ {- 1} \lambda_ {u _ {k}} \left\langle h h ^ {\top}, V [,: u _ {k} ] V [,: u _ {k} ] ^ {\top}, G _ {k} \right\rangle \\ = q _ {h, \sqrt {\Lambda} V ^ {\top}, G _ {k}} [ u _ {k} ]. \\ \end{array}
$$
Hence, we can use point 2 of Lemma 3.2 to sample a value for  $\hat{u}_k$  efficiently. Because  $\sqrt{\Lambda} V^{\top}$  has only  $R$  rows with  $R\sim 10^{2}$ , we can choose tuning parameter  $F = 1$  to achieve lower time per sample while incurring a modest  $O(R^3)$  space overhead. Now, introduce a random variable  $\hat{t}_k$  with distribution conditioned on  $\hat{u}_k = u_k$  given by
$$
p \left(\hat {t} _ {k} = t _ {k} \mid \hat {u} _ {k} = u _ {k}\right) := W \left[ t _ {k}, u _ {k} \right] / \| W [ \therefore , u _ {k} ] \| _ {1}. \tag {10}
$$
This distribution is well-defined, since we suppose that  $\hat{u}_k = u_k$  occurs with nonzero probability  $e[u_k]$ , which implies that  $\| W[.,u_k]\| _1\neq 0$ . Our remaining task is to efficiently sample from the distribution above. Below, we abbreviate  $\tilde{h} = V[.,u_k]\otimes h$  and derive
$$
\begin{array}{l} p (\hat {t} _ {k} = t _ {k} \mid \hat {u} _ {k} = u _ {k}) = \frac {\langle h h ^ {\top} , U _ {k} [ t _ {k} , : ] ^ {\top} U _ {k} [ t _ {k} , : ] , V [ : , u _ {k} ] V [ : , u _ {k} ] ^ {\top} \rangle}{\| W [ : , u _ {k} ] \| _ {1}} \\ = \frac {\left\langle \tilde {h} \tilde {h} ^ {\top} , U _ {k} \left[ t _ {k} , : \right] ^ {\top} U _ {k} \left[ t _ {k} , : \right] , [ 1 ] \right\rangle}{\| W [ : , u _ {k} ] \| _ {1}} \tag {11} \\ = q _ {\tilde {h}, U _ {k}, [ 1 ]} \left[ t _ {k} \right]. \\ \end{array}
$$
Based on the last line of Equation (11), we apply Lemma 3.2 again to build an efficient data structure to sample a row of  $U_{k}$ . Since  $Y = [1]$  is a rank-1 matrix, we can use point 3 of the lemma and select a larger parameter value  $F = R$  to reduce space usage. The sampling time for this stage becomes  $O(R^2\log \lceil I_j / R\rceil)$ .
To summarize, Algorithms 1 and 2 give the construction and sampling procedures for our data structure. They rely on the "BuildSampler" and "RowSample" procedures from Algorithms 3 and 4 in Appendix A.5, which relate to the data structure in Lemma 3.2. In the construction phase, we build  $N$  data structures from Lemma 3.2 for the distribution in Equation (11). Construction costs  $O\left(\sum_{j=1}^{N} I_j R^2\right)$ , and if any matrix  $U_j$  changes, we can rebuild  $Z_j$  in isolation. Because  $F = R$ , the space required for  $Z_j$  is  $O(I_j R)$ . In the sampling phase, the procedure in Algorithm 2 accepts an optional index  $j$  of a matrix to exclude from the Khatri-Rao product. The procedure begins by computing the symmetric eigendecomposition of each matrix  $G_{>k}$ . The eigendecomposition is computed only once per binary tree structure, and its computation cost is amortized over all  $J$  samples. It then creates data structures  $E_k$  for each of the distributions specified by Equation (9). These data structures (along with those from the construction phase) are used to draw  $\hat{u}_k$  and  $\hat{t}_k$  in succession. The random variables  $\hat{t}_k$  follow the distribution in Theorem 3.1 conditioned on prior draws, so the multi-index  $(\hat{t}_k)_{k \neq j}$  follows the leverage score distribution on  $A$ , as desired. Appendix A.6 proves the complexity claims in the theorem and provides further details about the algorithms.
3.3 Application to Tensor Decomposition
A tensor is a multidimensional array, and the CP decomposition represents a tensor as a sum of outer products [13]. See Appendix A.9 for an overview. To approximately decompose tensor  $\mathcal{T} \in \mathbb{R}^{I_1 \times \dots \times I_N}$ , the popular alternating least squares (ALS) algorithm begins with randomly initialized factor matrices  $U_j$ ,  $U_j \in \mathbb{R}^{I_j \times R}$  for  $1 \leq j \leq N$ . We call the column count  $R$  the rank of the decomposition. Each round of ALS solves  $N$  overdetermined least squares problems in sequence, each optimizing a single factor matrix while holding the others constant. The  $j$ -th least squares problem occurs in the update
$$
U _ {j} := \underset {X} {\arg \min } \left\| U _ {\neq j} \cdot X ^ {\top} - \operatorname {m a t} (\mathcal {T}, j) ^ {\top} \right\| _ {F}
$$
where  $U_{\neq j} = U_N \odot \ldots \odot U_{j+1} \odot U_{j-1} \odot \ldots \odot U_1$  is the Khatri-Rao product of all matrices excluding  $U_j$  and  $\mathrm{mat}(\cdot)$  denotes the mode- $j$  matricization of tensor  $\mathcal{T}$ . Here, we reverse the order of matrices
in the Khatri-Rao product to match the ordering of rows in the matricized tensor (see Appendix A.9 for an explicit formula for the matricization). These problems are ideal candidates for randomized sketching [4, 12, 15], and applying the data structure in Theorem 1.1 gives us the STS-CP algorithm.
Corollary 3.3 (STS-CP). Suppose  $\mathcal{T}$  is dense, and suppose we solve each least squares problem in ALS with a randomized sketching algorithm. A leverage score sampling approach as defined in section 2 guarantees that with  $\tilde{O}(R / (\varepsilon \delta))$  samples per solve, the residual of each sketched least squares problem is within  $(1 + \varepsilon)$  of the optimal residual with probability  $(1 - \delta)$ . The efficient sampler from Theorem 1.1 brings the complexity of ALS to
$$
\tilde {O} \left(\frac {\# i t}{\varepsilon \delta} \cdot \sum_ {j = 1} ^ {N} \left(N R ^ {3} \log I _ {j} + I _ {j} R ^ {2}\right)\right)
$$
where "it" is the number of ALS iterations, and with any term  $\log I_j$  replaced by  $\log R$  if  $I_j < R$ .
The proof appears in Appendix A.9 and combines Theorem 1.1 with Theorem 2.1. STS-CP also works for sparse tensors and likely provides a greater advantage here than the dense case, as sparse tensors tend to have much larger mode size [25]. The complexity for sparse tensors depends heavily on the sparsity structure and is difficult to predict. Nevertheless, we expect a significant speedup based on prior works that use sketching to accelerate CP decomposition [6, 15].
Experiments were conducted on CPU nodes of NERSC Perlmutter, an HPE Cray EX supercomputer, and our code is available at https://github.com/vbharadwaj-bk/ fast_tensor_ leverage.git. On tensor decomposition experiments, we compare our algorithms against the random and hybrid versions of CP-ARLS-LEV proposed by Larsen and Kolda [15]. These algorithms outperform uniform sampling and row-norm-squared sampling, achieving excellent accuracy and runtime relative to exact ALS. In contrast to TNS-CP and the Gaussian tensor network embedding proposed by Ma and Solomonik (see Table 1), CP-ARLS-LEV is one of the few algorithms that can practically
decompose sparse tensors with mode sizes in the millions. In the worst case, CP-ARLS-LEV requires  $\tilde{O}(R^{N - 1} / (\varepsilon \delta))$  samples per solve for an  $N$ -dimensional tensor to achieve solution guarantees like those in Theorem 2.1, compared to  $\tilde{O}(R / (\varepsilon \delta))$  samples required by STS-CP. Appendices A.10, A.11, and A.13 provide configuration details and additional results.
4.1 Runtime Benchmark
Figure 2 shows the time to construct our sampler and draw 50,000 samples from the Khatri-Rao product of i.i.d. Gaussian initialized factor matrices. We quantify the runtime impacts of varying  $N$ ,  $R$ , and  $I$ . The asymptotic behavior in Theorem 1.1 is reflected in our performance measurements, with the exception of the plot that varies  $R$ . Here, construction becomes disproportionately cheaper than sampling due to cache-efficient BLAS3 calls during construction. Even when the full Khatri-Rao product has  $\approx 3.78 \times 10^{22}$  rows (for  $I = 2^{25}$ ,  $N = 3$ ,  $R = 32$ ), we require only 0.31 seconds on average for sampling (top plot, rightmost points).
Algorithm 1 ConstructKRPSampler  $(U_{1},\dots,U_{N})$
1: for  $j = 1..N$  do
2:  $Z_{j}\coloneqq \mathrm{BuildSampler}(U_{j},F = R,[1])$
3:  $G_{j}\coloneqq U_{j}^{\top}U_{j}$
Algorithm 2 KRPSample(j, J)
1:  $G \coloneqq \bigstar_{k \neq j} G_k$
2: for  $k \neq j$  do
3:  $G_{>k} \coloneqq G^{+} * *_{k+j+1}^{N} G_{k}$
4: Decompose  $G_{>k} = V_k \Lambda_k V_k^\top$
5:  $E_{k}:= \text{BuildSampler}(\sqrt{\Lambda_{k}}\cdot V_{k}^{\top},F = 1,G_{k})$
6: for  $d = 1..J$  do
7:  $h = [1,\dots,1]^{\prime}$
8: for  $k \neq j$  do
9:  $\hat{u}_k\coloneqq \mathrm{RowSample}(E_k,h)$
0:  $\hat{t}_k := \mathrm{RowSample}(Z_k, h \circledast (V_k[:, \hat{u}_k]))$
11:  $h* = U_k[\hat{t}_k,:]$
12:  $s_d = (\hat{t}_k)_{k\neq j}$
13: return  $s_1, \ldots, s_J$
Figure 2: Average time (5 trials) to construct our proposed sampler and draw  $J = 50,000$  samples from  $U_{1} \odot \ldots \odot U_{N}$ , with  $U_{j} \in \mathbb{R}^{I \times R} \forall j$ . Error bars indicate 3 standard deviations.
Figure 3: Distortion and residual error (50 trials) for varying  $R$  and  $N$  on least squares,  $I = 2^{16}$ ,  $J = 5000$ . "X" marks indicate outliers 1.5 times the interquartile range beyond the median, stars indicate means.
4.2 Least Squares Accuracy Comparison
We now test our sampler on least squares problems of the form  $\min_x\| Ax - b\|$ , where  $A = U_{1}\odot \ldots \odot U_N$  with  $U_{j}\in \mathbb{R}^{I\times R}$  for all  $j$ . We initialize all matrices  $U_{j}$  entrywise i.i.d. from a standard normal distribution and randomly multiply  $1\%$  of all entries by 10. We choose  $b$  as a Kronecker product  $c_{1}\otimes \ldots \otimes c_{N}$ , with each vector  $c_{j}\in \mathbb{R}^{I}$  also initialized entrywise from a Gaussian distribution. We assume this form for  $b$  to tractably compute the exact solution to the linear least squares problem and evaluate the accuracy of our randomized methods. We do not give our algorithms access to the Kronecker form of  $b$ ; they are only permitted on-demand, black-box access to its entries.
For each problem instance, define the distortion of our sampling matrix  $S$  with respect to the column space of  $A$  as
$$
D (S, A) = \frac {\kappa (S Q) - 1}{\kappa (S Q) + 1} \tag {12}
$$
where  $Q$  is an orthonormal basis for the column space of  $A$  and  $\kappa(SQ)$  is the condition number of  $SQ$ . A higher-quality sketch  $S$  exhibits lower distortion, which quantifies the preservation of distances from the column space of  $A$  to the column space of  $SA$  [20]. For details about computing  $\kappa(SQ)$  efficiently when  $A$  is a Khatri-Rao product, see Appendix A.12. Next, define  $\varepsilon = \frac{\text{residual}_{\text{approx}}}{\text{residual}_{\text{opt}}} - 1$ , where  $\text{residual}_{\text{approx}}$  is the residual of the randomized least squares algorithm.  $\varepsilon$  is nonnegative and (similar to its role in Theorem 2.1) quantifies the quality of the randomized algorithm's solution.
For varying  $N$  and  $R$ , Figure 3 shows the average values of  $D$  and  $\varepsilon$  achieved by our algorithm against the leverage product approximation used by Larsen and Kolda [15]. Our sampler exhibits nearly constant distortion  $D$  for fixed rank  $R$  and varying  $N$ , and it achieves  $\varepsilon \approx 10^{-2}$  even when  $N = 9$ . The product approximation increases both the distortion and residual error by at least an order of magnitude.
4.3 Sparse Tensor Decomposition
We next apply STS-CP to decompose several large sparse tensors from the FROSTT collection [25] (see Appendix A.11 for more details on the experimental configuration). Our accuracy metric is the tensor fit. Letting  $\tilde{\mathcal{T}}$  be our low-rank CP approximation, the fit with respect to ground-truth tensor  $\mathcal{T}$  is  $\mathrm{fit}(\tilde{\mathcal{T}},\mathcal{T}) = 1 - \left\| \tilde{\mathcal{T}} -\mathcal{T}\right\|_{F} / \| \mathcal{T}\|_{F}$ .
Table 4 in Appendix A.13.1 compares the runtime per round of our algorithm against the Tensorly Python package [14] and Matlab Tensor Toolbox [3], with dramatic speedup over both. As Figure 4 shows, the fit achieved by CP-ARLS-LEV compared to STS-CP degrades as the rank increases for
Figure 4: Average fits (8 trials) achieved by randomized  $(J = 2^{16})$  and exact ALS for sparse tensor CP decomposition. Error bars indicate 3 standard deviations. See Appendix A.11 for details.
Figure 5: Average  $\varepsilon$  (5 runs) for randomized least squares solves in 10 ALS rounds,  $R = 50$ .
Figure 6: Fit vs. time, Reddit tensor,  $R = 100$ . Thick lines are averages 4 trial interpolations.
fixed sample count. By contrast, STS-CP improves the fit consistently, with a significant improvement at rank 125 over CP-ARLS-LEV. Timings for both algorithms are available in Appendix A.13.5. Figure 5 explains the higher fit achieved by our sampler on the Uber and Amazon tensors. In the first 10 rounds of ALS, we compute the exact solution to each least squares problem before updating the factor matrix with a randomized algorithm's solution. Figure 5 plots  $\varepsilon$  as ALS progresses for hybrid CP-ARLS-LEV and STS-CP. The latter consistently achieves lower residual per solve. We further observe that CP-ARLS-LEV exhibits an oscillating error pattern with period matching the number of modes  $N$ .
To assess the tradeoff between sampling time and accuracy, we compare the fit as a function of ALS update time for STS-CP and random CP-ARLS-LEV in Figure 6 (time to compute the fit excluded). On the Reddit tensor with  $R = 100$ , we compared CP-ARLS-LEV with  $J = 2^{16}$  against CP-ARLS-LEV with progressively larger sample count. Even with  $2^{18}$  samples per randomized least squares solve, CP-ARLS-LEV cannot achieve the maximum fit of STS-CP. Furthermore, STS-CP makes progress more quickly than CP-ARLS-LEV. See Appendix A.13.4 for similar plots for other datasets.
5 Discussion and Future Work
Our method for exact Khatri-Rao leverage score sampling enjoys strong theoretical guarantees and practical performance benefits. Especially for massive tensors such as Amazon and Reddit, our randomized algorithm's guarantees translate to faster progress to solution and higher final accuracies. The segment tree approach described here can be applied to sample from tensor networks besides the Khatri-Rao product. In particular, modifications to Lemma 3.2 permit efficient leverage sampling from a contraction of 3D tensor cores in ALS tensor train decomposition. We leave the generalization of our fast sampling technique as future work.
Acknowledgements, Funding, and Disclaimers
We thank the referees for valuable feedback which helped improve the paper.
V. Bharadwaj was supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC0022158. O. A. Malik and A. Buluc were supported by the Office of Science of the DOE under Award Number DE-AC02-05CH11231. L. Grigori was supported by the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation program through grant agreement 810367. R. Murray was supported by Laboratory Directed Research and Development (LDRD) funding from Berkeley Lab, provided by the Director, Office of Science, of the U.S. DOE under Contract No. DE-AC02-05CH11231. R. Murray was also funded by an NSF Collaborative Research Framework under NSF Grant Nos. 2004235 and 2004763. This research used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility, under Contract No. DE-AC02-05CH11231 using NERSC award ASCR-ERCAP0024170.
This report was prepared as an account of work sponsored by an agency of the United States Government. Neither the United States Government nor any agency thereof, nor any of their employees, makes any warranty, express or implied, or assumes any legal liability or responsibility for the accuracy, completeness, or usefulness of any information, apparatus, product, or process disclosed, or represents that its use would not infringe privately owned rights. Reference herein to any specific commercial product, process, or service by trade name, trademark, manufacturer, or otherwise does not necessarily constitute or imply its endorsement, recommendation, or favoring by the United States Government or any agency thereof. The views and opinions of authors expressed herein do not necessarily state or reflect those of the United States Government or any agency thereof.
[1] Thomas D. Ahle, Michael Kapralov, Jakob B. T. Knudsen, Rasmus Pagh, Ameya Velingker, David P. Woodruff, and Amir Zandieh. Oblivious Sketching of High-Degree Polynomial Kernels, pages 141-160. Society for Industrial and Applied Mathematics, 2020. doi: 10.1137/1.9781611975994.9.
[2] Nir Ailon and Bernard Chazelle. The fast Johnson-Lindenstrauss transform and approximate nearest neighbors. SIAM Journal on computing, 39(1):302-322, 2009.
[3] Brett W. Bader and Tamara G. Kolda. Efficient matlab computations with sparse and factored tensors. SIAM Journal on Scientific Computing, 30(1):205-231, 2008. doi: 10.1137/060676489.
[4] Casey Battaglino, Grey Ballard, and Tamara G. Kolda. A practical randomized CP tensor decomposition. SIAM Journal on Matrix Analysis and Applications, 39(2):876-901, 2018. doi: 10.1137/17M1112303.
[5] Ke Chen, Qin Li, Kit Newton, and Stephen J. Wright. Structured random sketching for PDE inverse problems. SIAM Journal on Matrix Analysis and Applications, 41(4):1742-1770, 2020. doi: 10.1137/20M1310497.
[6] Dehua Cheng, Richard Peng, Yan Liu, and Ioakeim Perros. SPALS: Fast alternating least squares via implicit leverage scores sampling. In Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
[7] Huaian Diao, Rajesh Jayaram, Zhao Song, Wen Sun, and David Woodruff. Optimal sketching for Kronecker product regression and low rank approximation. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[8] Petros Drineas, Malik Magdon-Ismail, Michael W. Mahoney, and David P. Woodruff. Fast approximation of matrix coherence and statistical leverage. J. Mach. Learn. Res., 13(1): 3475-3506, dec 2012. ISSN 1532-4435.
[9] Matthew Fahrbach, Gang Fu, and Mehrdad Ghadiri. Subquadratic Kronecker regression with applications to tensor decomposition. In Advances in Neural Information Processing Systems, volume 35, pages 28776-28789. Curran Associates, Inc., 2022.
[26] Linghao Song, Yuze Chi, Atefeh Sohrabizadeh, Young-kyu Choi, Jason Lau, and Jason Cong. Sextans: A streaming accelerator for general-purpose sparse-matrix dense-matrix multiplication. In Proceedings of the 2022 ACM/SIGDA International Symposium on Field-Programmable Gate Arrays, FPGA '22, page 65–77, New York, NY, USA, 2022. Association for Computing Machinery. ISBN 9781450391498. doi: 10.1145/3490422.3502357.
[27] Yining Wang, Hsiao-Yu Tung, Alexander J Smola, and Anima Anandkumar. Fast and guaranteed tensor decomposition via sketching. Advances in neural information processing systems, 28, 2015.
[28] Sasindu Wijeratne, Ta-Yang Wang, Rajgopal Kannan, and Viktor Prasanna. Accelerating sparse MTTKRP for tensor decomposition on FPGA. In Proceedings of the 2023 ACM/SIGDA International Symposium on Field Programmable Gate Arrays, FPGA '23, page 259-269, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9781450394178. doi: 10.1145/3543622.3573179.
[29] David Woodruff and Amir Zandieh. Leverage score sampling for tensor product matrices in input sparsity time. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 23933-23964. PMLR, 17-23 Jul 2022.
[30] David P Woodruff et al. Sketching as a tool for numerical linear algebra. Foundations and Trends® in Theoretical Computer Science, 10(1-2):1-157, 2014.
[31] Yao Yu, Athina P. Petropulu, and H. Vincent Poor. MIMO radar using compressive sampling. IEEE Journal of Selected Topics in Signal Processing, 4(1):146-163, February 2010. doi: 10.1109/JSTSP.2009.2038973.
A.1 Details about Table 1
CP-ALS [13] is the standard, non-randomized alternating least squares method given by Algorithm 6 in Appendix A.9. The least squares problems in the algorithm are solved by exact methods. CP-ARLS-LEV is the algorithm proposed by Larsen and Kolda [15] that samples rows from the Khatri-Rao product according to a product distribution of leverage scores on each factor matrix. The per-iteration runtimes for both algorithms are re-derived in Appendix C.3 of the work by Malik [18] from their original sources. Malik [18] proposed the CP-ALS-ES algorithm (not listed in the table), which is superseded by the TNS-CP algorithm [19]. We report the complexity from Table 1 of the latter work. The algorithm by Ma and Solomonik [17] is based on a general method to sketch tensor networks. Our reported complexity is listed in Table 1 for Algorithm 1 in their work.
Table 1 does not list the one-time initialization costs for any of the methods. All methods require at least  $O(NIR)$  time to randomly initialize factor matrices, and CP-ALS requires no further setup. CP-ARLS-LEV, TNS-CP, and STS-CP all require  $O(NIR^2)$  initialization time. CP-ARLS-LEV uses the initialization phase to compute the initial leverage scores of all factor matrices. TNS-CP uses the initialization step to compute and cache Gram matrices of all factors  $U_j$ . STS-CP must build the efficient sampling data structure described in Theorem 1.1. The algorithm from Ma and Solomonik requires an initialization cost of  $O(I^N m)$ , where  $m$  is a sketch size parameter that is  $O(NR / \varepsilon^2)$  to achieve the  $(\varepsilon, \delta)$  accuracy guarantee for each least squares solve.
A.2 Definitions of Matrix Products
Table 2 defines the standard matrix product  $\cdot$ , Hadamard product  $*$ , Kronecker product  $\otimes$ , and Khatri-Rao product  $\odot$ , as well as the dimensions of their operands.
Table 2: Matrix product definitions.
OPERATION	SIZE OF A	SIZE OF B	SIZE OF C	DEFINITION
C = A · B	(m, k)	(k, n)	(m, n)	C[i, j] = ∑a=1kA [i, a] B [a, j]
C = A ⊕ B	(m, n)	(m, n)	(m, n)	C [i, j] = A [i, j] B [i, j]
C = A ⊗ B	(m1, n1)	(m2, n2)	(m2m1, n2n1)	C [(i2, i1), (j2, j1)] = A [i1, j1] B [i2, j2]
C = A ⊙ B	(m1, n)	(m2, n)	(m2m1, n)	C [(i2, i1), j] = A [i1, j] B [i2, j]
A.3 Further Comparison to Prior Work
In this section, we provide a more detailed comparison of our sampling algorithm with the one proposed by Woodruff and Zandieh [29]. Their work introduces a ridge leverage-score sampling algorithm for Khatri-Rao products with the attractive property that the sketch can be formed in input-sparsity time. For constant failure probability  $\delta$ , the runtime to produce a  $(1 \pm \epsilon)$ $\ell_2$ -subspace embedding for  $A = U_{1} \odot \ldots \odot U_{N}$  is given in Appendix B of their work (proof of Theorem 2.7). Adapted to our notation, their runtime is
$$
O \left(\log^ {4} R \log N \sum_ {i = 1} ^ {N} \mathrm {n n z} (U _ {i}) + \frac {N ^ {7} s _ {\lambda} ^ {2} R}{\varepsilon^ {4}} \log^ {5} R \log N\right)
$$