SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Vivek Bharadwaj  $^{1,2}$ , Osman Asif Malik  $^{2}$ , Riley Murray  $^{3,2,1}$ , Laura Grigori  $^{4}$ , Aydin Buluc  $^{2,1}$ , James Demmel
<sup>1</sup>Electrical Engineering and Computer Science Department, UC Berkeley
$^{2}$ Computational Research Division, Lawrence Berkeley National Lab
$^{3}$ International Computer Science Institute
$^{4}$  Institute of Mathematics, EPFL & Lab for Simulation and Modelling, Paul Scherrer Institute
We present a data structure to randomly sample rows from the Khatri-Rao product of several matrices according to the exact distribution of its leverage scores. Our proposed sampler draws each row in time logarithmic in the height of the Khatri-Rao product and quadratic in its column count, with persistent space overhead at most the size of the input matrices. As a result, it tractably draws samples even when the matrices forming the Khatri-Rao product have tens of millions of rows each. When used to sketch the linear least squares problems arising in CANDECOMP / PARAFAC tensor decomposition, our method achieves lower asymptotic complexity per solve than recent state-of-the-art methods. Experiments on billion-scale sparse tensors validate our claims, with our algorithm achieving higher accuracy than competing methods as the decomposition rank grows.
The Khatri-Rao product (KRP, denoted by  $\odot$ ) is the column-wise Kronecker product of two matrices, and it appears in diverse applications across numerical analysis and machine learning [16]. We examine overdetermined linear least squares problems of the form  $\min_{X} \|AX - B\|_{F}$ , where the design matrix  $A = U_{1} \odot \ldots \odot U_{N}$  is the Khatri-Rao product of matrices  $U_{j} \in \mathbb{R}^{I_{j} \times R}$ . These problems appear prominently in signal processing [23], compressed sensing [31], inverse problems related to partial differential equations [5], and alternating least squares (ALS) CANDECOMP / PARAFAC (CP) tensor decomposition [13]. In this work, we focus on the case where  $A$  has moderate column count (several hundred at most). Despite this, the problem remains formidable because the height of  $A$  is  $\prod_{j=1}^{N} I_{j}$ . For row counts  $I_{j}$  in the millions, it is intractable to even materialize  $A$  explicitly.
Several recently-proposed randomized sketching algorithms can approximately solve least squares problems with Khatri-Rao product design matrices [4, 12, 15, 18, 29]. These methods apply a sketching operator  $S$  to the design and data matrices to solve the reduced least squares problem  $\min_{\tilde{X}}\left\| SA\tilde{X} -SB\right\| _F$ , where  $S$  has far fewer rows than columns. For appropriately chosen  $S$ , the residual of the downsampled system falls within a specified tolerance  $\varepsilon$  of the optimal residual with high probability  $1 - \delta$ . In this work, we constrain  $S$  to be a sampling matrix that selects and reweights a subset of rows from both  $A$  and  $B$ . When the rows are selected according to the distribution of statistical leverage scores on the design matrix  $A$ , only  $\tilde{O}\left(R / (\varepsilon \delta)\right)$  samples are required (subject to the assumptions at the end of section 2.1). The challenge, then, is to efficiently sample according to the leverage scores when  $A$  has Khatri-Rao structure.
We propose a leverage-score sampler for the Khatri-Rao product of matrices with tens of millions of rows each. After construction, our sampler draws each row in time quadratic in the column count, but logarithmic in the total row count of the Khatri-Rao product. Our core contribution is the following theorem.
Theorem 1.1 (Efficient Khatri-Rao Product Leverage Sampling). Given  $U_{1},\ldots ,U_{N}$  with  $U_{j}\in \mathbb{R}^{I_{j}\times R}$ , there exists a data structure satisfying the following:
1. The data structure has construction time  $O\left(\sum_{j=1}^{N} I_j R^2\right)$  and requires additional storage space  $O\left(\sum_{j=1}^{N} I_j R\right)$ . If a single entry in a matrix  $U_j$  changes, it can be updated in time  $O(R \log (I_j / R))$ . If the entire matrix  $U_j$  changes, it can be updated in time  $O(I_j R^2)$ .
2. The data structure produces  $J$  samples from the Khatri-Rao product  $U_{1} \odot \ldots \odot U_{N}$  according to the exact leverage score distribution on its rows in time
$$
O \left(N R ^ {3} + J \sum_ {k = 1} ^ {N} R ^ {2} \log \max  \left(I _ {k}, R\right)\right)
$$
using  $O(R^3)$  scratch space. The structure can also draw samples from the Khatri-Rao product of any subset of  $U_1, \ldots, U_N$ .
The efficient update property and ability to exclude one matrix are important in CP decomposition. When the inputs  $U_{1}, \ldots, U_{N}$  are sparse, an analogous data structure with  $O\left(R\sum_{j=1}^{N}\mathrm{nz}(U_{j})\right)$  construction time and  $O\left(\sum_{j=1}^{N}\mathrm{nz}(U_{j})\right)$  storage space exists with identical sampling time. Since the output factor matrices  $U_{1}, \ldots, U_{N}$  are typically dense, we defer the proof to Appendix A.8. Combined with error guarantees for leverage-score sampling, we achieve an algorithm for alternating least squares CP decomposition with asymptotic complexity lower than recent state-of-the-art methods (see Table 1).
Our method provides the most practical benefit on sparse input tensors, which may have dimension lengths in the tens of millions (unlike dense tensors that quickly incur intractable storage costs at large dimension lengths) [25]. On the Amazon and Reddit tensors with billions of nonzero entries, our algorithm STS-CP can achieve  $95\%$  of the fit of non-randomized ALS between 1.5x and 2.5x faster than a high-performance implementation of the state-of-the-art CP-ARLS-LEV algorithm [15]. Our algorithm is significantly more sample-efficient; on the Enron tensor, only  $\sim 65,000$  samples per solve were required to achieve the  $95\%$  accuracy threshold above a rank of 50, which could not be achieved by CP-ARLS-LEV with even 54 times as many samples.
Table 1: Asymptotic Complexity to decompose an  $N$ -dimensional  $I \times \ldots \times I$  dense tensor via CP alternating least squares. For randomized algorithms, each approximate least-squares solution has residual within  $(1 + \varepsilon)$  of the optimal value with high probability  $1 - \delta$ . Factors involving  $\log R$  and  $\log(1/\delta)$  are hidden ( $\tilde{O}$  notation). See A.1 for details.
ALGORITHM	COMPLEXITY PER ITERATION
CP-ALS [13]	N(N+I)IN-1R
CP-ARLS-LEV [15]	N(R+I)RN/(εδ)
TNS-CP [19]	N3IR3/(εδ)
GAUSSIAN TNE [17]	N2(N1.5R3.5/ε3+IR2)/ε2
STS-CP (OURS)	N(NR3log I+IR2)/(εδ)
Notation. We use  $[N]$  to denote the set  $\{1, \dots, N\}$  for a positive integer  $N$ . We use  $\tilde{O}$  notation to indicate the presence of multiplicative terms polylogarithmic in  $R$  and  $(1 / \delta)$  in runtime complexities. For the complexities of our methods, these logarithmic factors are no more than  $O(\log(R / \delta))$ . We
use Matlab notation  $A[i, :]$ ,  $A(:, i]$  to index rows, resp. columns, of matrices. For consistency, we use the convention that  $A[i, :]$  is a row vector. We use  $\cdot$  for standard matrix multiplication,  $\circledast$  as the elementwise product,  $\otimes$  to denote the Kronecker product, and  $\odot$  for the Khatri-Rao product. See Appendix A.2 for a definition of each operation. Given matrices  $A \in \mathbb{R}^{m_1 \times n}$ ,  $B \in \mathbb{R}^{m_2 \times n}$ , the  $j$ -th column of the Khatri-Rao product  $A \odot B \in \mathbb{R}^{m_1 m_2 \times n}$  is the Kronecker product  $A(:, j] \otimes B(:, j]$ .
We use angle brackets  $\langle \cdot, \dots, \cdot \rangle$  to denote a generalized inner product. For identically-sized vectors / matrices, it returns the sum of all entries in their elementwise product. For  $A, B, C \in \mathbb{R}^{m \times n}$ ,
$$
\langle A, B, C \rangle := \sum_ {i = 1, j = 1} ^ {m, n} A [ i, j ] B [ i, j ] C [ i, j ].
$$
Finally,  $M^{+}$  denotes the pseudoinverse of matrix  $M$
2.1 Sketched Linear Least Squares
A variety of random sketching operators  $S$  have been proposed to solve overdetermined least squares problems  $\min_{X} \|AX - B\|_{F}$  when  $A$  has no special structure [30, 2]. When  $A$  has Khatri-Rao product structure, prior work has focused on sampling matrices [6, 15], which have a single nonzero entry per row, operators composed of fast Fourier/ trigonometric transforms [12], or Countsketch-type operators [27, 1]. For tensor decomposition, however, the matrix  $B$  may be sparse or implicitly specified as a black-box function. When  $B$  is sparse, Countsketch-type operators still require the algorithm to iterate over all nonzero values in  $B$ . As Larsen and Kolda [15] note, operators similar to the FFT induce fill-in when applied to a sparse matrix  $B$ , destroying the benefits of sketching. Similar difficulties arise when  $B$  is implicitly specified. This motivates our decision to focus on row sampling operators, which only touch a subset of entries from  $B$ . Let  $\hat{x}_1, \dots, \hat{x}_J$  be a selection of  $J$  indices for the rows of  $A \in \mathbb{R}^{I \times R}$ , sampled i.i.d. according to a probability distribution  $q_1, \dots, q_I$ . The associated sampling matrix  $S \in \mathbb{R}^{J \times I}$  is specified by
$$
S \left[ j, i \right] = \left\{ \begin{array}{l l} \frac {1}{\sqrt {J q _ {i}}}, & \text {i f} \hat {x} _ {j} = i \\ 0, & \text {o t h e r w i s e} \end{array} \right.
$$
where the weight of each nonzero entry corrects bias induced by sampling. When the probabilities  $q_{j}$  are proportional to the leverage scores of the rows of  $A$ , strong guarantees apply to the solution of the downsampled problem.
Leverage Score Sampling. The leverage scores of a matrix assign a measure of importance to each of its rows. The leverage score of row  $i$  from matrix  $A \in \mathbb{R}^{I \times R}$  is given by
$$
\ell_ {i} = A [ i,: ] \left(A ^ {\top} A\right) ^ {+} A [ i,: ] ^ {\top} \tag {1}
$$
for  $1 \leq i \leq I$ . Leverage scores can be expressed equivalently as the squared row norms of the matrix  $Q$  in any reduced  $QR$  factorization of  $A$  [8]. The sum of all leverage scores is the rank of  $A$  [30]. Dividing the scores by their sum, we induce a probability distribution on the rows used to generate a sampling matrix  $S$ . The next theorem has appeared in several works, and we take the form given by Malik et al. [19]. For an appropriate sample count, it guarantees that the residual of the downsampled problem is close to the residual of the original problem.
Theorem 2.1 (Guarantees for Leverage Score Sampling). Given  $A \in \mathbb{R}^{I \times R}$  and  $\varepsilon, \delta \in (0,1)$ , let  $S \in \mathbb{R}^{J \times I}$  be a leverage score sampling matrix for  $A$ . Further define  $\tilde{X} = \arg \min_{X} \| SAX - SB \|_F$ . If  $J \gtrsim R$  max ( $\log(R / \delta), 1 / (\varepsilon \delta)$ ), then with probability at least  $1 - \delta$  it holds that
$$
\left\| A \tilde {X} - B \right\| _ {F} \leq (1 + \varepsilon) \min  _ {X} \| A X - B \| _ {F}.
$$
For the applications considered in this work,  $R$  ranges up to a few hundred. As  $\varepsilon$  and  $\delta$  tend to 0 with fixed  $R$ ,  $1 / (\varepsilon \delta)$  dominates  $\log (R / \delta)$ . Hence, we assume that the minimum sample count  $J$  to achieve the guarantees of the theorem is  $\Omega(R / (\varepsilon \delta))$ .
2.2 Prior Work
Khatri-Rao Product Leverage Score Sampling. Well-known sketching algorithms exist to quickly estimate the leverage scores of dense matrices [8]. These algorithms are, however, intractable for  $A = U_{1} \odot \ldots \odot U_{N}$  due to the height of the Khatri-Rao product. Cheng et al. [6] instead approximate each score as a product of leverage scores associated with each matrix  $U_{j}$ . Larsen and Kolda [15] propose CP-ARLS-LEV, which uses a similar approximation and combines random sampling with a deterministic selection of high-probability indices. Both methods were presented in the context of CP decomposition. To sample from the Khatri-Rao product of  $N$  matrices, both require  $O(R^{N} / (\varepsilon \delta))$  samples to achieve the  $(\varepsilon, \delta)$  guarantee on the residual of each least squares solution. These methods are simple to implement and perform well when the Khatri-Rao product has column count up to 20-30. On the other hand, they suffer from high sample complexity as  $R$  and  $N$  increase. The TNS-CP algorithm by Malik et al. [19] samples from the exact leverage score distribution, thus requiring only  $O(R / (\varepsilon \delta))$  samples per least squares solve. Unfortunately, it requires time  $O\left(\sum_{j=1}^{N} I_{j} R^{2}\right)$  to draw each sample.
Comparison to Woodruff and Zandieh. The most comparable results to ours appear in work by Woodruff and Zandieh [29], who detail an algorithm for approximate ridge leverage-score sampling for the Khatri-Rao product in near input-sparsity time. Their work relies on a prior oblivious method by Ahle et al. [1], which sketches a Khatri-Rao product using a sequence of Countsketch / OSNAP operators arranged in a tree. Used in isolation to solve a linear least squares problem, the tree sketch construction time scales as  $O\left(\frac{1}{\varepsilon}\sum_{j=1}^{N}\mathrm{nnz}(U_j)\right)$  and requires an embedding dimension quadratic in  $R$  to achieve the  $(\varepsilon, \delta)$  solution-quality guarantee. Woodruff and Zandieh use a collection of these tree sketches, each with carefully-controlled approximation error, to design an algorithm with linear runtime dependence on the column count  $R$ . On the other hand, the method exhibits  $O(N^7)$  scaling in the number of matrices involved, has  $O(\varepsilon^{-4})$  scaling in terms of the desired accuracy, and relies on a sufficiently high ridge regularization parameter. Our data structure instead requires construction time quadratic in  $R$ . In exchange, we use distinct methods to design an efficiently-updatable sampler with runtime linear in both  $N$  and  $\varepsilon^{-1}$ . These properties are attractive when the column count  $R$  is below several thousand and when error as low as  $\epsilon \approx 10^{-3}$  is needed in the context of an iterative solver (see Figure 5). Moreover, the term  $O(R^2\sum_{j=1}^{N}I_j)$  in our construction complexity arises from symmetric rank-  $k$  updates, a highly-optimized BLAS3 kernel on modern CPU and GPU architectures. Appendix A.3 provides a more detailed comparison between the two approaches.
Kronecker Regression. Kronecker regression is a distinct (but closely related) problem to the one we consider. There,  $A = U_{1} \otimes \ldots \otimes U_{N}$  and the matrices  $U_{i}$  have potentially distinct column counts  $R_{1}, \ldots, R_{N}$ . While the product distribution of leverage scores from  $U_{1}, \ldots, U_{N}$  provides only an approximation to the leverage score distribution of the Khatri-Rao product [6, 15], it provides the exact leverage distribution for the Kronecker product. Multiple works [7, 9] combine this property with other techniques, such as dynamically-updatable tree-sketches [21], to produce accurate and updatable Kronecker sketching methods. None of these results apply directly in our case due to the distinct properties of Kronecker and Khatri-Rao products.
3 An Efficient Khatri-Rao Leverage Sampler
Without loss of generality, we will prove part 2 of Theorem 1.1 for the case where  $A = U_{1} \odot \ldots \odot U_{N}$ ; the case that excludes a single matrix follows by reindexing matrices  $U_{k}$ . We further assume that  $A$  is a nonzero matrix, though it may be rank-deficient. Similar to prior sampling works [18, 29], our algorithm will draw one sample from the Khatri-Rao product by sampling a row from each of  $U_{1}, U_{2}, \ldots$  in sequence and computing their Hadamard product, with the draw from  $U_{j}$  conditioned on prior draws from  $U_{1}, \ldots, U_{j-1}$ .
Let us index each row of  $A$  by a tuple  $(i_1, \dots, i_N) \in [I_1] \times \dots \times [I_N]$ . Equation (1) gives
$$
\ell_ {i _ {1}, \dots , i _ {N}} = A \left[ \left(i _ {1}, \dots , i _ {N}\right),: \right] \left(A ^ {\top} A\right) ^ {+} A \left[ \left(i _ {1}, \dots , i _ {N}\right),: \right] ^ {\top}. \tag {2}
$$
For  $1 \leq k \leq N$ , define  $G_{k} := U_{k}^{\top}U_{k} \in \mathbb{R}^{R \times R}$  and  $G := \left(\bigotimes_{k=1}^{N} G_{k}\right) \in \mathbb{R}^{R \times R}$ ; it is a well-known fact that  $G = A^{\top}A$  [13]. For a single row sample from  $A$ , let  $\hat{s}_{1}, \dots, \hat{s}_{N}$  be random variables for the
draws from multi-index set  $[I_1] \times \ldots \times [I_N]$  according to the leverage score distribution. Assume, for some  $k$ , that we have already sampled an index from each of  $[I_1], \ldots, [I_{k-1}]$ , and that the first  $k-1$  random variables take values  $\hat{s}_1 = s_1, \ldots, \hat{s}_{k-1} = s_{k-1}$ . We abbreviate the latter condition as  $\hat{s}_{<k} = s_{<k}$ . To sample from  $I_k$ , we seek the distribution of  $\hat{s}_k$  conditioned on  $\hat{s}_1, \ldots, \hat{s}_{k-1}$ . Define  $h_{<k}$  as the transposed elementwise product of rows already sampled:
$$
h _ {<   k} := \bigotimes_ {i = 1} ^ {k - 1} U _ {i} [ s _ {i},: ] ^ {\top}. \tag {3}
$$
Also define  $G_{>k}$  as
$$
G _ {> k} := G ^ {+} * \bigotimes_ {i = k + 1} ^ {N} G _ {i}. \tag {4}
$$
Then the following theorem provides the conditional distribution of  $\hat{s}_k$ .
Theorem 3.1 (Malik 2022, [18], Adapted). For any  $s_k \in [I_k]$ ,
$$
\begin{array}{l} p \left(\hat {s} _ {k} = s _ {k} \mid \hat {s} _ {<   k} = s _ {<   k}\right) = C ^ {- 1} \left\langle h _ {<   k} h _ {<   k} ^ {\top}, U _ {k} [ s _ {k},: ] ^ {\top} U _ {k} [ s _ {k},: ], G _ {> k} \right\rangle \tag {5} \\ := q _ {h <   k, U _ {k}, G > k} \left[ s _ {k} \right] \\ \end{array}
$$
where  $C = \langle h_{<k}h_{<k}^{\top}, U_k^{\top}U_k, G_{>k} \rangle$  is nonzero.
We include the derivation of Theorem 3.1 from Equation (2) in Appendix A.4. Computing all entries of the probability vector  $q_{h_{<k},U_k,G_{>k}}$  would cost  $O(I_jR^2)$  per sample, too costly when  $U_j$  has millions of rows. It is likewise intractable (in preprocessing time and space complexity) to precompute probabilities for every possible conditional distribution on the rows of  $U_j$ , since the conditioning random variable has  $\prod_{k=1}^{j-1} I_k$  potential values. Our key innovation is a data structure to sample from a discrete distribution of the form  $q_{h_{<k},U_k,G_{>k}}$  without materializing all of its entries or incurring superlinear cost in either  $N$  or  $\varepsilon^{-1}$ . We introduce this data structure in the next section and will apply it twice in succession to get the complexity in Theorem 1.1.
3.1 Efficient Sampling from  $q_{h,U,Y}$
We introduce a slight change of notation in this section to simplify the problem and generalize our sampling lemma. Let  $h \in \mathbb{R}^R$  be a vector and let  $Y \in \mathbb{R}^{R \times R}$  be a positive semidefinite (p.s.d.) matrix, respectively. Our task is to sample  $J$  rows from a matrix  $U \in \mathbb{R}^{I \times R}$  according to the distribution
$$
q _ {h, U, Y} [ s ] := C ^ {- 1} \left\langle h h ^ {\top}, U ^ {\top} [ s,: ] U [ s,: ], Y \right\rangle \tag {6}
$$
provided the normalizing constant  $C = \langle hh^{\top}, U^{\top}U, Y \rangle$ , is nonzero. We impose that all  $J$  rows are drawn with the same matrices  $Y$  and  $U$ , but potentially distinct vectors  $h$ . The following lemma establishes that an efficient sampler for this problem exists.
Lemma 3.2 (Efficient Row Sampler). Given matrices  $U \in \mathbb{R}^{I \times R}$ ,  $Y \in \mathbb{R}^{R \times R}$  with  $Y$  p.s.d., there exists a data structure parameterized by positive integer  $F$  that satisfies the following:
1. The structure has construction time  $O\left( IR^2 \right)$  and storage requirement  $O\left(R^{2} \lceil I / F \rceil\right)$ . If  $I < F$ , the storage requirement drops to  $O(1)$ .
2. After construction, the data structure can produce a sample according to the distribution  $q_{h,U,Y}$  in time  $O(R^2\log [I / F] + FR^2)$  for any vector  $h$ .
3. If  $Y$  is a rank-1 matrix, the time per sample drops to  $O(R^2\log \lceil I / F\rceil +FR)$ .
This data structure relies on an adaptation of a classic binary-tree inversion sampling technique [22]. Consider a partition of the interval  $[0,1]$  into  $I$  bins, the  $i$ -th having width  $q_{h,U,Y}[i]$ . We sample  $d \sim \mathrm{Uniform}[0,1]$  and return the index of the containing bin. We locate the bin index through a binary search terminated when at most  $F$  bins remain in the search space, which are then scanned in linear time. Here,  $F$  is a tuning parameter that we will use to control sampling complexity and space usage.
For  $a > b$ , assume that  $\bigoplus_{i=a}^{b} (\ldots)$  produces a vector / matrix filled with ones.
We can regard the binary search as a walk down a full, complete binary tree  $T_{I,F}$  with  $\lceil I / F \rceil$  leaves, the nodes of which store contiguous, disjoint segments  $S(v) = \{S_0(v)..S_1(v)\} \subseteq [I]$  of size at most  $F$ . The segment of each internal node is the union of segments held by its children, and the root node holds  $\{1,\dots,I\}$ . Suppose that the binary search reaches node  $v$  with left child  $L(v)$  and maintains the interval [low, high]  $\subseteq [0,1]$  as the remaining search space to explore. Then the search branches left in the tree iff  $d < \mathrm{low} + \sum_{i\in S(L(v))}q_{h,U,Y}[i]$ .
This branching condition can be evaluated effi
ciently if appropriate information is stored at each node of the segment tree. Excluding the offset "low", the branching threshold takes the form
$$
\sum_ {i \in S (v)} q _ {h, U, Y} [ i ] = C ^ {- 1} \left\langle h h ^ {\top}, \sum_ {i \in S (v)} U [ i,: ] ^ {\top} U [ i,: ], Y \right\rangle := C ^ {- 1} \left\langle h h ^ {\top}, G ^ {v}, Y \right\rangle . \tag {7}
$$
Here, we call each matrix  $G^v \in \mathbb{R}^{R \times R}$  a partial Gram matrix. In time  $O(IR^2)$  and space  $O(R^2 [I / F])$ , we can compute and cache  $G^v$  for each node of the tree to construct our data structure. Each subsequent binary search costs  $O(R^2)$  time to evaluate Equation (7) at each of  $\log [I / F]$  internal nodes and  $O(FR^2)$  to evaluate  $q_{h,U,Y}$  at the  $F$  indices held by each leaf, giving point 2 of the lemma. This cost at each leaf node reduces to  $O(FR)$  in case  $Y$  is rank-1, giving point 3. A complete proof of this lemma appears in Appendix A.5.
3.2 Sampling from the Khatri-Rao Product
We face difficulties if we directly apply Lemma 3.2 to sample from the conditional distribution in Theorem 3.1. Because  $G_{>k}$  is not rank-1 in general, we must use point 2 of the lemma, where no selection of the parameter  $F$  allows us to simultaneously satisfy the space and runtime constraints of Theorem 1.1. Selecting  $F = R$  results in cost  $O(R^3)$  per sample (violating the runtime requirement in point 2), whereas  $F = 1$  results in a superlinear storage requirement  $O(IR^2)$  (violating the space requirement in point 1, and becoming prohibitively expensive for  $I \geq 10^6$ ). To avoid these extremes, we break the sampling procedure into two stages. The first stage selects a 1-dimensional subspace spanned by an eigenvector of  $G_{>k}$ , while the second samples according to Theorem 3.1 after projecting the relevant vectors onto the selected subspace. Lemma 3.2 can be used for both stages, and the second stage benefits from point 3 to achieve better time and space complexity.
Below, we abbreviate  $q = q_{h_{<k},U_k,G_{>k}}$  and  $h = h_{<k}$ . When sampling from  $I_k$ , observe that  $G_{>k}$  is the same for all samples. We compute a symmetric eigendecomposition  $G_{>k} = V\Lambda V^{\top}$ , where each column of  $V$  is an eigenvector of  $G_{>k}$  and  $\Lambda = \mathrm{diag}((\lambda_u)_{u=1}^R)$  contains the eigenvalues along the diagonal. This allows us to rewrite entries of  $q$  as
$$
q [ s ] = C ^ {- 1} \sum_ {u = 1} ^ {R} \lambda_ {u} \left\langle h h ^ {\top}, U _ {k} [ s,: ] ^ {\top} U _ {k} [ s,: ], V [ :, u ] V [ :, u ] ^ {\top} \right\rangle . \tag {8}
$$