Title: Fast Exact Leverage Score Sampling from Khatri-Rao Products with Applications to Tensor Decomposition
Venue: NeurIPS 2023
Accepted: yes
Keywords: Tensor Decomposition

Leverage Scores

Randomized Linear Algebra

Sketching

Khatri-Rao Product

Sparse Tensors
Abstract: 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 decomposition, our method achieves lower asymptotic complexity per solve than recent state-of-the-art methods. Experiments on billion-scale sparse tensors and synthetic data validate our theoretical claims, with our algorithm achieving higher accuracy than competing methods as the decomposition rank grows.

Fast Exact Leverage Score Sampling from Khatri-Rao Products with Applications to Tensor Decomposition

Abstract

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 KhatriRao 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.

1 Introduction

The Khatri-Rao product (KRP, denoted by ⊙ ) 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 ⊙ ... ⊙ U N is the Khatri-Rao product of matrices U j ∈ R I j × 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 ∏ N j =1 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 ˜ X ∥ ∥ ∥ SA ˜ X -SB ∥ ∥ ∥ F , where S has far fewer rows than columns. For appropriately chosen S , the residual of the downsampled system falls within a specified tolerance ε of the optimal residual with high probability 1 -δ . 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 ˜ O ( R/ ( εδ )) 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 , ..., U N with U j ∈ R I j × R , there exists a data structure satisfying the following:

The data structure has construction time O ( ∑ N j =1 I j R 2 ) and requires additional storage space O ( ∑ N j =1 I j R ) . 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 ) .

The data structure produces J samples from the Khatri-Rao product U 1 ⊙ ... ⊙ U N according to the exact leverage score distribution on its rows in time

using O ( R 3 ) scratch space. The structure can also draw samples from the Khatri-Rao product of any subset of U 1 , ..., U N .

The efficient update property and ability to exclude one matrix are important in CP decomposition. When the inputs U 1 , ..., U N are sparse, an analogous data structure with O ( R ∑ N j =1 nnz ( U j ) ) construction time and O ( ∑ N j =1 nnz ( U j ) ) storage space exists with identical sampling time. Since the output factor matrices U 1 , ..., 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 ∼ 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.

2 Preliminaries and Related Work

Notation. We use [ N ] to denote the set { 1 , ..., N } for a positive integer N . We use ˜ O notation to indicate the presence of multiplicative terms polylogarithmic in R and (1 /δ ) in runtime complexities. For the complexities of our methods, these logarithmic factors are no more than O (log( R/δ )) . 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 · for standard matrix multiplication, ⊛ as the elementwise product, ⊗ to denote the Kronecker product, and ⊙ for the Khatri-Rao product. See Appendix A.2 for a definition of each operation. Given matrices A ∈ R m 1 × n , B ∈ R m 2 × n , the j -th column of the Khatri-Rao product A ⊙ B ∈ R m 1 m 2 × n is the Kronecker product A [: , j ] ⊗ B [: , j ] .

We use angle brackets ⟨· , ..., ·⟩ 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 ∈ R m × n ,

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 Countsketchtype 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 ˆ x 1 , ..., ˆ x J be a selection of J indices for the rows of A ∈ R I × R , sampled i.i.d. according to a probability distribution q 1 , ..., q I . The associated sampling matrix S ∈ R J × I is specified by

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 ∈ R I × R is given by

for 1 ≤ i ≤ 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 ∈ R I × R and ε, δ ∈ (0 , 1) , let S ∈ R J × I be a leverage score sampling matrix for A . Further define ˜ X = arg min X ∥ SAX -SB ∥ F . If J ≳ R max(log ( R/δ ) , 1 / ( εδ )) , then with probability at least 1 -δ it holds that

For the applications considered in this work, R ranges up to a few hundred. As ε and δ tend to 0 with fixed R , 1 / ( εδ ) dominates log( R/δ ) . Hence, we assume that the minimum sample count J to achieve the guarantees of the theorem is Ω( R/ ( εδ )) .

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 ⊙ ... ⊙ 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 / ( εδ )) samples to achieve the ( ε, δ ) 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/ ( εδ )) samples per least squares solve. Unfortunately, it requires time O ( ∑ N j =1 I j R 2 ) 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 ( 1 ε ∑ N j =1 nnz ( U j ) ) and requires an embedding dimension quadratic in

R to achieve the ( ε, δ ) 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 ( ε -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 ε -1 . These properties are attractive when the column count R is below several thousand and when error as low as ϵ ≈ 10 -3 is needed in the context of an iterative solver (see Figure 5). Moreover, the term O ( R 2 ∑ N j =1 I j ) in our construction complexity arises from symmetric rankk 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 ⊗ ... ⊗ U N and the matrices U i have potentially distinct column counts R 1 , ..., R N . While the product distribution of leverage scores from U 1 , ..., 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 ⊙ ... ⊙ 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 , .... in sequence and computing their Hadamard product, with the draw from U j conditioned on prior draws from U 1 , ..., U j -1 .

Let us index each row of A by a tuple ( i 1 , ..., i N ) ∈ [ I 1 ] × ... × [ I N ] . Equation (1) gives

For 1 ≤ k ≤ N , define G k := U ⊤ k U k ∈ R R × R and G := ( ⊛ N k =1 G k ) ∈ R R × R ; it is a well-known fact that G = A ⊤ A [13]. For a single row sample from A , let ˆ s 1 , ..., ˆ s N be random variables for the

draws from multi-index set [ I 1 ] × ... × [ I N ] according to the leverage score distribution. Assume, for some k , that we have already sampled an index from each of [ I 1 ] , ..., [ I k -1 ] , and that the first k -1 random variables take values ˆ s 1 = s 1 , ..., ˆ s k -1 = s k -1 . We abbreviate the latter condition as ˆ s <k = s <k . To sample from I k , we seek the distribution of ˆ s k conditioned on ˆ s 1 , ... ˆ s k -1 . Define h <k as the transposed elementwise product 1 of rows already sampled:

Then the following theorem provides the conditional distribution of ˆ s k .

Theorem 3.1 (Malik 2022, [18], Adapted) . For any s k ∈ [ I k ] ,

where C = ⟨ h <k h ⊤ <k , U ⊤ k U k , G >k ⟩ 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 j R 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 ∏ j -1 k =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 ε -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 ∈ R R be a vector and let Y ∈ R R × R be a positive semidefinite (p.s.d.) matrix, respectively. Our task is to sample J rows from a matrix U ∈ R I × R according to the distribution

provided the normalizing constant C = ⟨ hh ⊤ , U ⊤ U, Y ⟩ , 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 ∈ R I × R , Y ∈ R R × R with Y p.s.d., there exists a data structure parameterized by positive integer F that satisfies the following:

The structure has construction time O ( IR 2 ) and storage requirement O ( R 2 ⌈ I/F ⌉ ) . If I < F , the storage requirement drops to O (1) .

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 .

If Y is a rank-1 matrix, the time per sample drops to O ( R 2 log ⌈ I/F ⌉ + 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 ∼ 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.

Also define G >k as

We can regard the binary search as a walk down a full, complete binary tree T I,F with ⌈ I/F ⌉ leaves, the nodes of which store contiguous, disjoint segments S ( v ) = { S 0 ( v ) ..S 1 ( v ) } ⊆ [ 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 , ..., I } . Suppose that the binary search reaches node v with left child L ( v ) and maintains the interval [ low , high ] ⊆ [0 , 1] as the remaining search space to explore. Then the search branches left in the tree iff d < low + ∑ i ∈ S ( L ( v )) q h,U,Y [ i ] .

[1

..

8]

[1

..

4]

[1

,

2]

[3

,

4]

[5

..

8]

[5

,

6]

[7

,

8]

q

1

q

2

q

3

q

4

q

5

q

6

q

7

q

8

ciently if appropriate information is stored at each node of the segment tree. Excluding the offset

This branching condition can be evaluated effi'low', the branching threshold takes the form

Here, we call each matrix G v ∈ R R × 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 ≥ 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 Λ V ⊤ , where each column of V is an eigenvector of G >k and Λ = diag(( λ u ) R u =1 ) contains the eigenvalues along the diagonal. This allows us to rewrite entries of q as

Define matrix W ∈ R I k × R elementwise by

and observe that all of its entries are nonnegative. Since λ u ≥ 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 :

where the vector w of nonnegative weights is given by w [ u ] = ( C -1 λ 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

defined by W [: , u ] / ∥ W [: , u ] ∥ 1 . Let ˆ u k be a random variable distributed according to the probability mass vector w . We have, for C taken from Theorem 3.1,

Hence, we can use point 2 of Lemma 3.2 to sample a value for ˆ u k efficiently. Because √ Λ V ⊤ has only R rows with R ∼ 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 ˆ t k with distribution conditioned on ˆ u k = u k given by

̸

This distribution is well-defined, since we suppose that ˆ u k = u k occurs with nonzero probability e [ u k ] , which implies that ∥ W [: , u k ] ∥ 1 = 0 . Our remaining task is to efficiently sample from the distribution above. Below, we abbreviate ˜ h = V [: , u k ] ⊛ h and derive

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 ⌈ I j /R ⌉ ) .

̸

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 ( ∑ N j =1 I j R 2 ) , 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 ˆ u k and ˆ t k in succession. The random variables ˆ t k follow the distribution in Theorem 3.1 conditioned on prior draws, so the multi-index ( ˆ t k ) k = 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 T ∈ R I 1 × ... × I N , the popular alternating least squares (ALS) algorithm begins with randomly initialized factor matrices U j , U j ∈ R I j × R for 1 ≤ j ≤ 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

̸

where U = j = U N ⊙ ... ⊙ U j +1 ⊙ U j -1 ⊙ ... ⊙ U 1 is the Khatri-Rao product of all matrices excluding U j and mat ( · ) denotes the modej matricization of tensor 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 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 ˜ O ( R/ ( εδ )) samples per solve, the residual of each sketched least squares problem is within (1 + ε ) of the optimal residual with probability (1 -δ ) . The efficient sampler from Theorem 1.1 brings the complexity of ALS to

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].

4 Experiments

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

Algorithm 1 ConstructKRPSampler( U 1 , ..., U N )

1: for j = 1 ..N do

2: Z j := BuildSampler ( U j , F = R, [1])

3: G j := U ⊤ j U j

Algorithm 2 KRPSample( j , J )

̸

1: G := ⊛ k = j G k

̸

2: for k = j do

3: G >k := G + ⊛⊛ N k = j +1 G k

4: Decompose G >k = V k Λ k V ⊤ k √

5: E k := BuildSampler ( Λ k · V ⊤ k , F = 1 , G k )

6: for d = 1 ..J do

7: h = [1 , ..., 1]

⊤

̸

8: for k = j do

9: ˆ u k := RowSample ( E k , h )

10: ˆ t k := RowSample ( Z k , h ⊛ ( V k [: , ˆ u k ]))

11: h ∗ = U k [ ˆ t k , : ]

̸

12: s d = ( ˆ t k ) k = j

13: return s 1 , ..., s J

decompose sparse tensors with mode sizes in the millions. In the worst case, CP-ARLS-LEV requires ˜ O ( R N -1 / ( εδ )) samples per solve for an N -dimensional tensor to achieve solution guarantees like those in Theorem 2.1, compared to ˜ O ( R/ ( εδ )) 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 ≈ 3 . 78 × 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).

10 2

10 3

10 4

10 5

10 6

10 7

I

10

3

10

2

10

1

10 0

Time (s)

R=32,N=3

16

32

64

128

R

0

1

2

Time (s)

I = 2 22 , N=3

2

4

6

8

N

0.2

0.4

0.6

I = 2 22 , R=32

Construction

Sampling

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 ⊙ ... ⊙ U N with U j ∈ R I × 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 ⊗ ... ⊗ c N , with each vector c j ∈ 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

where Q is an orthonormal basis for the column space of A and κ ( 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 κ ( SQ ) efficiently when A is a Khatri-Rao product, see Appendix A.12. Next, define ε = residualapprox residualopt -1 , where residualapprox is the residual of the randomized least squares algorithm. ε 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 ε 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 ε ≈ 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 ˜ T be our low-rank CP approximation, the fit with respect to ground-truth tensor T

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

4

5

6

7

8

9

10

N

10

6

10

5

10

4

10

3

10

2

10

1

10 0

10 1

Residual Error

16

32

64

128

R

N=6

10

1

10 0

10 1

10 2

Distortion D(S, A)

N=6

Product Approximation

25

50

75

100

125

Target Rank

0.20

0.22

0.24

Fit

Uber (~3.3e6 nz)

25

50

75

100

125

Target Rank

0.05

0.10

0.15

0.20

Fit

Enron* (~5.4e7 nz)

25

50

75

100

125

Target Rank

0.05

0.06

0.07

0.08

Fit

NELL-2* (~7.7e7 nz)

25

50

75

100

125

Target Rank

0.34

0.36

0.38

0.40

Fit

Amazon (~1.8e9 nz)

25

50

75

100

125

Target Rank

0.06

0.08

0.10

Fit

Reddit* (~4.7e9 nz)

STS-CP (ours)

Exact Solve

0

10

20

30

LSTSQ Problem Number

10

3

10

2

Amazon

0

10

20

30

40

LSTSQ Problem Number

10

3

10

2

Uber

STS-CP (ours)

0

1000

2000

3000

4000

5000

6000

7000

Time (s)

0.080

0.085

0.090

0.095

0.100

Fit

STS-CP (ours), J=65,536

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 ε 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.

References

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.

Nir Ailon and Bernard Chazelle. The fast Johnson-Lindenstrauss transform and approximate nearest neighbors. SIAM Journal on computing , 39(1):302-322, 2009.

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.

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.

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.

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.

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.

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.

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.

Azzam Haidar, Tingxing Dong, Stanimire Tomov, Piotr Luszczek, and Jack Dongarra. Framework for batched and GPU-resident factorization algorithms to block Householder transformations. In ISC High Performance , Frankfurt, Germany, 07-2015 2015. Springer, Springer.

N. Halko, P. G. Martinsson, and J. A. Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM Review , 53(2):217-288, 2011. doi: 10.1137/090771806.

Ruhui Jin, Tamara G Kolda, and Rachel Ward. Faster Johnson-Lindenstrauss transforms via Kronecker products. Information and Inference: A Journal of the IMA , 10(4):1533-1562, October 2020. ISSN 2049-8772. doi: 10.1093/imaiai/iaaa028.

Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Review , 51(3):455-500, August 2009. ISSN 0036-1445. doi: 10.1137/07070111X. Publisher: Society for Industrial and Applied Mathematics.

Jean Kossaifi, Yannis Panagakis, Anima Anandkumar, and Maja Pantic. Tensorly: Tensor learning in python. Journal of Machine Learning Research (JMLR) , 20(26), 2019.

Brett W. Larsen and Tamara G. Kolda. Practical leverage-based sampling for low-rank tensor decomposition. SIAM J. Matrix Analysis and Applications , 43(3):1488-1517, August 2022. doi: 10.1137/21M1441754.

Shuangzhe Liu and Götz Trenkler. Hadamard, Khatri-Rao, Kronecker and other matrix products. International Journal of Information and Systems Sciences , 4(1):160-177, 2008.

Linjian Ma and Edgar Solomonik. Cost-efficient gaussian tensor network embeddings for tensor-structured inputs. In Advances in Neural Information Processing Systems , volume 35, pages 38980-38993. Curran Associates, Inc., 2022.

Osman Asif Malik. More efficient sampling for tensor decomposition with worst-case guarantees. In Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 14887-14917. PMLR, 17-23 Jul 2022.

Osman Asif Malik, Vivek Bharadwaj, and Riley Murray. Sampling-based decomposition algorithms for arbitrary tensor networks, October 2022. arXiv:2210.03828 [cs, math].

Riley Murray, James Demmel, Michael W. Mahoney, N. Benjamin Erichson, Maksim Melnichenko, Osman Asif Malik, Laura Grigori, Piotr Luszczek, Michał Derezi´ nski, Miles E. Lopes, Tianyu Liang, Hengrui Luo, and Jack Dongarra. Randomized numerical linear algebra : A perspective on the field with an eye to software, 2023.

Aravind Reddy, Zhao Song, and Lichen Zhang. Dynamic tensor product regression. In Advances in Neural Information Processing Systems , volume 35, pages 4791-4804. Curran Associates, Inc., 2022.

Feras A. Saad, Cameron E. Freer, Martin C. Rinard, and Vikash K. Mansinghka. Optimal approximate sampling from discrete probability distributions. Proceedings of the ACM on Programming Languages , 4(POPL):1-31, January 2020. ISSN 2475-1421. doi: 10.1145/ 3371104.

N.D. Sidiropoulos and R.S. Budampati. Khatri-Rao space-time codes. IEEE Transactions on Signal Processing , 50(10):2396-2407, 2002. doi: 10.1109/TSP.2002.803341.

Shaden Smith and George Karypis. SPLATT: The Surprisingly ParalleL spArse Tensor Toolkit. https://github.com/ShadenSmith/splatt , 2016.

Shaden Smith, Jee W. Choi, Jiajia Li, Richard Vuduc, Jongsoo Park, Xing Liu, and George Karypis. FROSTT: The formidable repository of open sparse tensors and tools, 2017. URL http://frostt.io/ .

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.

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.

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.

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.

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.

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 Appendix

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/ε 2 ) to achieve the ( ε, δ ) accuracy guarantee for each least squares solve.

A.2 Definitions of Matrix Products

Table 2 defines the standard matrix product · , Hadamard product ⊛ , Kronecker product ⊗ , and Khatri-Rao product ⊙ , as well as the dimensions of their operands.

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 δ , the runtime to produce a (1 ± ϵ ) ℓ 2 -subspace embedding for A = U 1 ⊙ ... ⊙ U N is given in Appendix B of their work (proof of Theorem 2.7). Adapted to our notation, their runtime is

where s λ = ∑ R i =1 λ i λ i + λ , λ 1 , ..., λ R are the eigenvalues of the Gram matrix G of matrix A , and λ ≥ 0 is a regularization parameter. For comparison, our runtime for constant failure probability δ is

Woodruff and Zandieh's method provides a significant advantage for large column count R or high regularization parameter λ . As a result, it is well-suited to the problem of regularized low-rank approximation when the column count R is given by the number of data points in a dataset. On the

other hand, the algorithm has poor dependence on the matrix count N and error parameter ε . For tensor decomposition, R is typically no larger than a few hundred, while high accuracy ( ϵ ≈ 10 -3 ) is required for certain tensors to achieve a fit competitive with non-randomized methods (see section 4.3, Figures 4 and 5). When λ is small, we have s λ ≈ R . Here, Woodruff and Zandieh's runtime has an O ( R 3 ) dependence similar to ours. When R ≤ log 4 R log N , our sampler has faster construction time as well.

Finally, we note that our sampling data structure can be constructed using highly cache-efficient, parallelizable symmetric rankR updates (BLAS3 operation dSYRK ). As a result, the quadratic dependence on R in our algorithm can be mitigated by dense linear algebra accelerators, such as GPUs or TPUs.

A.4 Proof of Theorem 3.1

Theorem 3.1 appeared in a modified form as Lemma 10 in the work by Malik [18]. This original version used the definition ˜ G >k = Φ ⊛⊛ N a = k +1 G k in place of G >k defined in Equation (4), where Φ was a sketched approximation of G + . Woodruff and Zandieh [29] exhibit a version of the theorem with similar modifications. We prove the version stated in our work below.

Proof of Theorem 3.1. We rely on the assumption that the Khatri-Rao product A is a nonzero matrix (but it may be rank-deficient). We begin by simplifying the expression for the leverage score of a row of A corresponding to multi-index ( i 1 , ..., i N ) . Beginning with Equation (2), we derive

We proceed to the main proof of the theorem. To compute p (ˆ s k = s k | ˆ s <k = s <k ) , we marginalize over random variables ˆ s k +1 ... ˆ s N . Recalling the definition of h <k from Equation (3), we have

The first line above follows by marginalizing over ˆ s k +1 , ..., ˆ s N . The second line follows because the joint random variable (ˆ s 1 , ..., ˆ s N ) follows the distribution of statistical leverage scores on the rows of

A . We now plug in Equation (13) to get

We now compute the normalization constant C for the distribution by summing the last line of Equation (15) over all possible values for ˆ s k :

̸

For k = 1 , we have h <k = [1 , ..., 1] ⊤ , so C = ⟨ G k , G >k ⟩ . Then C is the sum of all leverage scores, which is known to be the rank of A [30]. Since A was assumed nonzero, C = 0 . For k > 1 , assume that the conditioning event ˆ s <k = s <k occurs with nonzero probability. This is a reasonable assumption, since our sampling algorithm will never select prior values ˆ s 1 , ..., ˆ s k -1 that have 0 probability of occurrence. Let ˜ C be the normalization constant for the conditional distribution on ˆ s k -1 . Then we have

Since ˜ C > 0 , we must have C > 0 .

A.5 Proof of Lemma 3.2

We detail the construction procedure, sampling procedure, and correctness of our proposed data structure. Recall that T I,F denotes the collection of nodes in a full, complete binary tree with ⌈ I/F ⌉ leaves. Each leaf v ∈ T I,F holds a segment S ( v ) = { S 0 ( v ) ..S 1 ( v ) } ⊆ { 1 ..I } , with | S ( v ) | ≤ F and S ( u ) ∩ S ( v ) = ∅ for distinct leaves u, v . For each internal node v , S ( v ) = S ( L ( v )) ∪ S ( R ( v )) , where L ( v ) and R ( v ) denote the left and right children of node v . The root node r satisfies S ( r ) = { 1 ..I } .

Construction: Algorithm 3 gives the procedure to build the data structure. We initialize a segment tree T I,F and compute G v for all leaf nodes v ∈ T I,F as a sum of outer products of rows from U (lines 1-3). Starting at the level above the leaves, we then compute G v for each internal node as

the sum of G L ( v ) and G R ( v ) , the partial Gram matrices of its two children. Runtime O ( IR 2 ) is required to compute I outer products across all iterations of the loop on line 3. Our segment tree has ⌈ I/F ⌉ -1 internal nodes, and the addition in line 6 contributes runtime O ( R 2 ) for each internal node. This adds complexity O ( R 2 ( ⌈ I/F ⌉ -1)) ≤ O ( IR 2 ) , for total construction time O ( IR 2 ) .

To analyze the space complexity, observe that we store a matrix G v ∈ R R × R at all 2 ⌈ I/F ⌉ -1 nodes of the segment tree, for asymptotic space usage O ( ⌈ I/F ⌉ R 2 ) . We can cut the space usage in half by only storing G v when v is either the root or a left child in our tree, since the sampling procedure in Algorithm 4 only accesses the partial Gram matrix stored by left children. We can cut the space usage in half again by only storing the upper triangle of each symmetric matrix G v . Finally, in the special case that I < F , the segment tree has depth 1 and the initial binary search can be eliminated entirely. As a result, the data structure has O (1) space overhead, since we can avoid storing any partial Gram matrices G v . This proves the complexity claims in point 1 of Lemma 3.2.

Sampling: Algorithm 4 gives the procedure to draw a sample from our proposed data structure. It is easy to verify that the normalization constant C for q h,U,Y is ⟨ hh ⊤ , G root ( T I,F ) , Y ⟩ , since G root ( T I,F ) = U ⊤ U . Lines 8 and 9 initialize a pair of templated procedures ˜ m and ˜ q , each of which accepts a node from the segment tree. The former is used to compute the branching threshold at each internal node, and the latter returns the probability vector q h,U,Y [ S 0 ( v ) : S 1 ( v )] for the segment { S 0 ( v ) ..S 1 ( v ) } maintained by a leaf node. To see this last fact, observe for i ∈ [ I ] that

The loop on line 12 performs the binary search using the two templated procedures. Line 18 uses the procedure ˜ q to scan through at most F bin endpoints after the binary search finishes early.

The depth of segment tree T I,F is log ⌈ I/F ⌉ . As a result, the runtime of the sampling procedure is dominated by log ⌈ I/F ⌉ evaluations of ˜ m and a single evaluation of ˜ q during the binary search. Each execution of procedure ˜ m requires time O ( R 2 ) , relying on the partial Gram matrices G v computed during the construction phase. When Y is a general p.s.d. matrix, the runtime of ˜ q is O ( FR 2 ) . This complexity is dominated by the matrix multiplication W · ( hh ⊤ ⊛ Y ) on line 5. In this case, the runtime of the 'RowSampler" procedure to draw one sample is O ( R 2 log ⌈ I/F ⌉ + FR 2 ) , satisfying the complexity claims in point 2 of the lemma.

Now supppose Y is a rank-1 matrix with Y = uu ⊤ for some vector u . We have hh ⊤ ⊛ Y = ( h ⊛ u )( h ⊛ u ) ⊤ . This gives

where the square is elementwise. The runtime of the procedure ˜ q is now dominated by a matrix-vector multiplication that costs time O ( FR ) . In this case, we have per-sample complexity O ( R 2 log ⌈ I/F ⌉ + FR ) , matching the complexity claim in point 3 of the lemma.

Correctness: Recall that the inversion sampling procedure partitions the interval [0 , 1] into I bins, the i -th bin having width q h,U,Y [ i ] . The goal of our procedure is to find the bin that contains the uniform random draw d . Since procedure ˜ m correctly returns the branching threshold (up to the offset 'low') given by Equation (7), the loop on line 12 correctly implements a binary search on the list of bin endpoints specified by the vector q h,U,Y . At the end of the loop, c is a leaf node that

Algorithm 4 Row Sampling Procedure

maintains a collection S ( c ) of bins, one of which contains the random draw d . Since the procedure ˜ q correctly returns probabilities q h,U,Y [ i ] for i ∈ S ( c ) for leaf node c , (see Equation (18)), line 18 finds the bin that contains the random draw d . The correctness of the procedure follows from the correctness of inversion sampling [22].

A.6 Cohesive Proof of Theorem 1.1

In this proof, we fully explain Algorithms 1 and 2 in the context of the sampling procedure outlined in section 3.2. We verify the complexity claims first and then prove correctness.

Construction and Update: For each matrix U j , Algorithm 1 builds an efficient row sampling data structure Z j as specified by Lemma 3.2. We let the p.s.d. matrix Y that parameterizes each sampler be a matrix of ones, and we set F = R . From Lemma 3.2, the time to construct sampler Z j is O ( I j R 2 ) . The space used by sampler Z j is O ( ⌈ I j /F ⌉ R 2 ) = O ( I j R ) , since F = R . In case I j < R , we use the special case described in Appendix A.5 to get a space overhead O (1) , avoiding a term O ( R 2 ) in the space complexity.

Summing the time and space complexities over all j proves part 1 of the theorem. To update the data structure if matrix U j changes, we only need to rebuild sampler Z j for a cost of O ( I j R 2 ) . The construction phase also computes and stores the Gram matrix G j for each matrix U j . We defer the update procedure in case a single entry of matrix U j changes to Appendix A.7.

̸

Sampling: For all indices k (except possibly j ), lines 1-5 from Algorithm 2 compute G >k and its eigendecomposition. Only a single pass over the Gram matrices G k is needed, so these steps cost O ( R 3 ) for each index k . Line 5 builds an efficient row sampler E k for the matrix of scaled eigenvectors √ Λ k · V k . For sampler k , we set Y = G k with cutoff parameter F = 1 . From Lemma 3.2, the construction cost is O ( R 3 ) for each index k , and the space required by each sampler is O ( R 3 ) . Summing these quantities over all k = j gives asymptotic runtime O ( NR 3 ) for lines 2-5.

̸

The loop spanning lines 6-12 draws J row indices from the Khatri-Rao product U = j . For each sample, we maintain a 'history vector' h to write the variables h <k from Equation (3). For each index k = j , we draw random variable ˆ u k using the row sampler E k . This random draw indexes a scaled eigenvector of G >k . We then use the history vector h multiplied by the eigenvector to sample a row index ˆ t k using data structure Z k . The history vector h is updated, and we proceed to draw the next index ˆ t k .

̸

As written, lines 2-5 also incur scratch space usage O ( NR 3 ) . The scratch space can be reduced to O ( R 3 ) by exchanging the order of loops on line 6 and line 8 and allocating J separate history vectors h , once for each draw. Under this reordering, we perform all J draws for each variable ˆ u k and ˆ t k before moving to ˆ u k +1 and ˆ t k +1 . In this case, only a single data structure E k is required at each iteration of the outer loop, and we can avoid building all the structures in advance on line 5. We keep the algorithm in the form written for simplicity, but we implemented the memory-saving approach in our code.

̸

From Lemma 3.2, lines 9 and 10 cost O ( R 2 log R ) and O ( R 2 log ⌈ I k /R ⌉ ) , respectively. Line 11 costs O ( R ) and contributes a lower-order term. Summing over all k = j , the runtime to draw a single sample is

̸

̸

̸

Adding the runtime for all J samples to the runtime of the loop spanning lines 2-6 gives runtime O ( NR 3 + J ∑ k = j R 2 log max ( I k , R ) ) , and the complexity claims have been proven.

Correctness: We show correctness for the case where j = -1 and we sample from the Khatri-Rao product of all matrices U k , since the proof for any other value of j requires a simple reindexing of matrices. To show that our sampler is correct, it is enough to prove the condition that for 1 ≤ k ≤ N ,

since, by Theorem 3.1, p (ˆ s k = s k | ˆ s <k = s <k ) = q h <k ,U k ,G >k [ s k ] . This would imply that the joint random variable ( ˆ t 1 , ..., ˆ t N ) has the same probability distribution as (ˆ s 1 , ..., ˆ s N ) , which by definition follows the leverage score distribution on U 1 ⊙ ... ⊙ U N . To prove the condition in Equation (19), we apply Equations (9) and (11) derived earlier:

A.7 Efficient Single-Element Updates

Applications such as CP decomposition typically change all entries of a single matrix U j between iterations, incurring an update cost O ( I j R 2 ) for our data structure from Theorem 1.1. In case only a single element of U j changes, our data structure can be updated in time O ( R log I j ) .

Proof. Algorithm 5 gives the procedure when the update U j [ r, c ] := ˆ u is performed. The matrices G v refer to the partial Gram matrices maintained by each node v of the segment trees in our data structure, and the matrix ˜ U j refers to the matrix U j before the update operation.

Let T I j ,R be the segment tree corresponding to matrix U j in the data structure, and let v ∈ T I j ,R be the leaf whose segment contains r . Lines 3-5 of the algorithm update the row and column indexed by c in the partial Gram matrix held by the leaf node.

The only other nodes requiring an update are ancestors of v , each holding a partial Gram matrix that is the sum of its two children. Starting from the direct parent A ( v ) , the loop on line 6 performs these ancestor updates. The addition on line 8 only requires time O ( R ) , since only row and column c change between the old value of G v and its updated version. Thus, the runtime of this procedure is O ( R log ⌈ I j /R ⌉ ) from multiplying the cost to update a single node by the depth of the tree.

Algorithm 5 UpdateSampler( j, r, c, ˆ u )

1: Let u = ˜ U j [ r, c ]

2: Locate v such that r ∈ S ( v )

3: Update G v [ c, :] += (ˆ u -u ) ˜ U j [ r, :]

4: Update G v [: , c ] += (ˆ u -u ) ˜ U j [ r, :] ⊤

5: Update G v [ c, c ] += (ˆ u -u ) 2

̸

6: while v = root ( T I j ,R ) do

7: v prev := v , v := A ( v )

8: Update G v := G sibling ( v prev ) + G v prev

A.8 Extension to Sparse Input Matrices

Our data structure is designed to sample from Khatri-Rao products U 1 ⊙ ... ⊙ U N where the input matrices U 1 , ..., U N are dense, a typical situation in tensor decomposition. Slight modifications to the construction procedure permit our data structure to handle sparse matrices efficiently as well. The following corollary states the result as a modification to Theorem 1.1.

Corollary A.1 (Sparse Input Modification) . When input matrices U 1 , ..., U N are sparse, point 1 of Theorem 1.1 can be modified so that the proposed data structure has O ( R ∑ N j =1 nnz ( U j ) ) construction time and O ( ∑ N j =1 nnz ( U j ) ) storage space. The sampling time and scratch space usage in point 2 of Theorem 1.1 does not change. The single-element update time in point 1 is likewise unchanged.

Proof. We will modify the data structure in Lemma 3.2. The changes to its construction and storage costs will propagate to our Khatri-Rao product sampler, which maintains one of these data structures for each input matrix.

Let us restrict ourselves to the case F = R,Y = [1] in relation to the data structure in Lemma 3.2. These choices for F and Y are used in the construction phase given by Algorithm 1. The proof in Appendix A.5 constrains each leaf v of a segment tree T I,F to hold a contiguous segment S ( v ) ⊆ [ I ] of cardinality at most F . Instead, choose each segment S ( v ) = { S 0 ( v ) ..S 1 ( v ) } so that U [ S 0 ( v ) : S 1 ( v ) , :] has at most R 2 nonzeros, and the leaf count of the tree is at most ⌈ nnz ( U ) /R 2 ⌉ +1 for input matrix U ∈ R I × R . Assuming the nonzeros of U are sorted in row-major order, we can construct such a partition of [ I ] into segments in time O ( nnz ( U )) by iterating in order through the nonzero rows and adding each of them to a 'current' segment. We shift to a new segment when the current segment cannot hold any more nonzeros.

This completes the modification to the data structure in Lemma 3.2, and we now analyze its updated time / space complexity.

Updated Construction / Update Complexity of Lemma 3.2, F = R,Y = [1] : Algorithm 3 constructs the partial Gram matrix for each leaf node v in the segment tree. Each nonzero in the segment U [ S 0 ( v ) : S 1 ( v ) , :] contributes time O ( R ) during line 3 of Algorithm 3 to update a single row and column of G v . Summed over all leaves, the cost of line 3 is O ( nnz ( U ) R ) . The remainder of the construction procedure updates the partial Gram matrices of all internal nodes. Since there are at most O ( ⌈ nnz ( U ) /R 2 ⌉ ) internal nodes and the addition on line 6 costs O ( R 2 ) per node, the remaining steps of the construction procedure cost O ( nnz ( U )) , a lower-order term. The construction time is therefore O ( nnz ( U ) R ) .

Since we store a single partial Gram matrix of size R 2 at each of O ( ⌈ nnz ( U ) /R 2 ⌉ ) internal nodes, the space complexity of our modified data structure is O ( nnz ( U )) .

Finally, the data structure update time in case a single element of U is modified does not change from Theorem 1.1. Since the depth of the segment tree ⌈ nnz ( U ) /R 2 ⌉ +1 is upper-bounded by ⌈ I/R ⌉ +1 , the runtime of the update procedure in Algorithm 5 stays the same.

Updated Sampling Complexity of Lemma 3.2, F = R,Y = [1] : The procedure 'RowSample' in Algorithm 4 now conducts a traversal of a tree of depth O ( ⌈ nnz ( U ) /R 2 ⌉ ) . As a result, we

can still upper-bound the number of calls to procedure ˜ m as ⌈ I/F ⌉ . The runtime of procedure ˜ m is unchanged. The runtime of procedure ˜ q for leaf node c is dominated by the matrix-vector multiplication U [ S 0 ( c ) : S 1 ( c ) , :] · h . This runtime is O ( nnz ( U [ S 0 ( c ) : S 1 ( c ) , :])) ≤ O ( R 2 ) . Putting these facts together, the sampling complexity of the data structure in Lemma 3.2 does not change under our proposed modifications for F = R,Y = [1] .

Updated Construction Complexity of Theorem 1.1 : Algorithm 1 now requires O ( R ∑ N j =1 nnz ( U j ) ) construction time and O ( ∑ N j =1 nnz ( U j ) ) storage space, summing the costs for the updated structure from Lemma 3.2 over all matrices U 1 , ..., U N . The sampling complexity of these data structures is unaffected by the modifications, which completes the proof of the corollary.

A.9 Alternating Least Squares CP Decomposition

CP Decomposition. CP decomposition represents an N -dimensional tensor ˜ T ∈ R I 1 × ... × I n as a weighted sum of generalized outer products. Formally, let U 1 , ..., U N with U j ∈ R I j × R be factor matrices with each column having unit norm, and let σ ∈ R R be a nonnegative coefficient vector. We call R the rank of the decomposition. The tensor ˜ T that the decomposition represents is given elementwise by

which is a generalized inner product between σ ⊤ and rows U j [ i j , :] for 1 ≤ j ≤ N . Given an input tensor T and a target rank R , the goal of approximate CP decomposition is to find a rankR representation ˜ T that minimizes the Frobenius norm ∥ ∥ ∥ T -˜ T ∥ ∥ ∥ F .

̸

Definition of Matricization. The matricization mat ( T , j ) flattens tensor T ∈ R I 1 × ... × I N into a matrix and isolates mode j along the row axis of the output. The output of matricization has dimensions I j × ∏ k = j I k . We take the formal definition below from a survey by Kolda and Bader [13]. The tensor entry T [ i 1 , ..., i N ] is equal to the matricization entry mat ( T , j ) [ i N , u ] , where

̸

̸

̸

Details about Alternating Least Squares. Let U 1 , ..., U N be factor matrices of a low-rank CP decomposition, U k ∈ R I k × R . We use U = j to denote ⊙ k =1 k = N,k = j U k . Note the inversion of order here to match indexing in the definition of matricization above. Algorithm 6 gives the non-randomized alternating least squares algorithm CP-ALS that produces a decomposition of target rank R given input tensor T ∈ R I 1 × ... × I N in general format. The random initialization on line 1 of the algorithm can be implemented by drawing each entry of the factor matrices U j according to a standard normal distribution, or via a randomized range finder [11]. The vector σ stores the generalized singular values of the decomposition. At iteration j within a round, ALS holds all factor matrices except U j constant and solves a linear-least squares problem on line 6 for a new value for U j . In between least squares solves, the algorithm renormalizes the columns of each matrix U j to unit norm and stores their original norms in the vector σ . Appendix A.11 contains more details about the randomized range finder and the convergence criteria used to halt iteration.

̸

We obtain a randomized algorithm for sparse tensor CP decomposition by replacing the exact least squares solve on line 6 with a randomized method according to Theorem 2.1. Below, we prove Corollary 3.3, which derives the complexity of the randomized CP decomposition algorithm.

̸

Proof of Corollary 3.3. The design matrix U = j for optimization problem j within a round of ALS has dimensions ∏ k = j I k × R . The observation matrix mat ( T , j ) ⊤ has dimensions ∏ k = j I k × I j . To achieve error threshold 1 + ε with probability 1 -δ on each solve, we draw J = ˜ O ( R/ ( εδ )) rows from both the design and observation matrices and solve the downsampled problem (Theorem 2.1).

̸

̸

Algorithm 6 CP-ALS( T , R )

1: Initialize U j ∈ R I j × R randomly for 1 ≤ j ≤ N .

2: Renormalize U j [: , i ] / = ∥ U j [: , i ] ∥ 2 , 1 ≤ j ≤ N, 1 ≤ i ≤ R .

3: Initialize σ ∈ R R to [1] .

4: while not converged do

5: for j = 1 ...N do

6: U j := arg min X ∥ ∥ U = j · X ⊤ -mat ( T , j ) ⊤ ∥ ∥ F

̸

7: σ [ i ] = ∥ U j [: , i ] ∥ 2 , 1 ≤ i ≤ R

9: return [ σ ; U 1 , ..., U N ] .

8: Renormalize U j [: , i ] / = ∥ U j [: , i ] ∥ 2 , 1 ≤ i ≤ R .

These rows are sampled according to the leverage score distribution on the rows of U = j , for which we use the data structure in Theorem 1.1. After a one-time initialization cost O ( ∑ N j =1 I j R 2 ) ) before the ALS iteration begins, the complexity to draw J samples (assuming I j ≥ R ) is

̸

̸

̸

The cost to assemble the corresponding subset of the observation matrix is O ( JI j ) = ˜ O ( RI j / ( εδ )) . The cost to solve the downsampled least squares problem is O ( JR 2 ) = ˜ O ( I j R 2 / ( εδ )) , which dominates the cost of forming the subset of the observation matrix. Finally, we require additional time O ( I j R 2 ) to update the sampling data structure (Theorem 1.1 part 1). Adding these terms together and summing over 1 ≤ j ≤ N gives

̸

Rounding N -1 to N and multiplying by the number of iterations gives the desired complexity. When I j < R for any j , the complexity changes in Theorem 1.1 propagate to the equation above. The column renormalization on line 8 of the CP-ALS algorithm contributes additional time O ( ∑ N j =1 I j R ) per round, a lower-order term.

A.10 Experimental Platform and Sampler Parallelism

We provide two implementations of our sampler. The first is a slow reference implementation written entirely in Python, which closely mimics our pseudocode and can be used to test correctness. The second is an efficient implementation written in C++, parallelized in shared memory with OpenMP and Intel Thread Building Blocks.

Each Perlmutter CPU node (our experimental platform) is equipped with two sockets, each containing an AMD EPYC 7763 processor with 64 cores. All benchmarks were conducted with our efficient C++ implementation using 128 OpenMP threads. We link our code against Intel Thread Building blocks to call a multithreaded sort function when decomposing sparse tensors. We use OpenBLAS 0.3.21 to handle linear algebra with OpenMP parallelism enabled, but our code links against any linear algebra library implementing the CBLAS and LAPACKE interfaces.

Our proposed data structure samples from the exact distribution of leverage scores of the Khatri-Rao product, thereby enjoying better sample efficiency than alternative approaches such as CP-ARLS-LEV [15]. The cost to draw each sample, however, is O ( R 2 log H ) , where H is the number of rows in the Khatri-Rao product. Methods such as row-norm-squared sampling or CP-ARLS-LEV can draw each

sample in time O (log H ) after appropriate preprocessing. Therefore, efficient parallelization of our sampling procedure is required for competitive performance, and we present two strategies below.

Asynchronous Thread Parallelism : The KRPSampleDraw procedure in Algorithm 2 can be called by multiple threads concurrently without data races. The simplest parallelization strategy divides the J samples equally among the threads in a team, each of which makes calls to KRPSampleDraw asynchronously. This strategy works well on a CPU, but is less attractive on a SIMT processor like a GPU where instruction streams cannot diverge without significant performance penalties.

Synchronous Batch Parallelism As an alternative to the asynchronous strategy, suppose for the moment that all leaves have the same depth in each segment tree. Then for every sample, STSample makes a sequence of calls to ˜ m , each updating the current node by branching left or right in the tree. The length of this sequence is the depth of the tree, and it is followed by a single call to the function ˜ q . Observe that procedure ˜ m in Algorithm 4 can be computed with a matrix-vector multiplication followed by a dot product. The procedure ˜ q of Algorithm 4 requires the same two operations if F = 1 or Y = [1] . Thus, we can create a batched version of our sampling procedure that makes a fixed length sequence of calls to batched gemv and dot routines. All processors march in lock-step down the levels of each segment tree, each tracking the branching paths of a distinct set of samples. The MAGMAlinear algebra library provides a batched version of gemv [10], while a batched dot product can be implemented with an ad hoc kernel. MAGMA also offers a batched version of the symmetric rankk update routine syrk , which is helpful to parallelize row sampler construction (Algorithm 3). When all leaves in the tree are not at the same level, the the bottom level of the tree can be handled with a special sequence of instructions making the required additional calls to ˜ m .

Our CPU code follows the batch synchronous design pattern. To avoid dependency on GPU-based MAGMA routines in our CPU prototype, portions of the code that should be batched BLAS calls are standard BLAS calls wrapped in a for loop. These sections can be easily replaced when the appropriate batched routines are available.

A.11 Sparse Tensor CP Experimental Configuration

Table 3 lists the nonzero counts and dimensions of sparse tensors in our experiments [25]. We took the log of all values in the Enron, NELL-2, and Reddit-2015 tensors. Consistent with established practice, this operation damps the effect of a few high magnitude tensor entries on the fit metric [15].

̸

The factor matrices for the Uber, Amazon, NELL-2, and Reddit experiments were initialized with i.i.d. entries from the standard normal distribution. As suggested by Larsen and Kolda [15], the Enron tensor's factors were initialized with a randomized range finder [11]. The range finder algorithm initializes each factor matrix U j as mat ( T , j ) S , a sketch applied to the modej matricization of T with S ∈ R ∏ k = j I k × R . Larsen and Kolda chose S as a sparse sampling matrix to select a random subset of fibers along each mode. We instead used an i.i.d. Gaussian sketching matrix that was not materialized explicitly. Instead, we exploited the sparsity of T and noted that at most nnz ( T ) columns of mat ( T , j ) were nonzero. Thus, we computed at most nnz ( T ) rows of the random sketching matrix S , which were lazily generated and discarded during the matrix multiplication without incurring excessive memory overhead.

ALS was run for a maximum of 40 rounds on all tensors except for Reddit, which was run for 80 rounds. The exact fit was computed every 5 rounds (defined as 1 epoch), and we used an early

stopping condition to terminate runs before the maximum round count. The algorithm was terminated at epoch T if the maximum fit in the last 3 epochs did not exceed the maximum fit from epoch 1 through epoch T -3 by tolerance 10 -4 .

Hybrid CP-ARLS-LEV deterministically includes rows from the Khatri-Rao product whose probabilities exceed a threshold τ . The ostensible goal of this procedure is to improve diversity in sample selection, as CP-ARLS-LEV may suffer from many repeat draws of high probability rows. We replicated the conditions proposed in the original work by selecting τ = 1 /J [15].

Individual trials of non-randomized (exact) ALS on the Amazon and Reddit tensors required several hours on a single Perlmutter CPU node. To speed up our experiments, accuracy measurements for exact ALS in Figure 3 were carried out using multi-node SPLATT, The Surprisingly ParalleL spArse Tensor Toolkit [24], on four Perlmutter CPU nodes. The fits computed by SPLATT agree with those computed by our own non-randomized ALS implementation. As a result, Figure 3 verifies that our randomized algorithm STS-CP produces tensor decompositions with accuracy comparable to those by highly-optimized, state-of-the-art CP decomposition software. We leave a distributed-memory implementation of our randomized algorithms to future work.

A.12 Efficient Computation of Sketch Distortion

The definition of σ in this section is different from its definition in the rest of this work. The condition number κ of a matrix M is defined as

where σ min ( M ) and σ max ( M ) denote the minimum and maximum nonzero singular values of M . Let A be a Khatri-Rao product of N matrices U 1 , ..., U N with ∏ N j =1 I j rows, R columns, and rank r ≤ R . Let A = Q Σ V ⊤ be its reduced singular value decomposition with Q ∈ R ∏ j I j × r , Σ ∈ R r × r , and V ∈ R r × R . Finally, let S ∈ R J × ∏ j I j be a leverage score sampling matrix for A . Our goal is to compute κ ( SQ ) without fully materializing either A or its QR decomposition. We derive

The matrix SA ∈ R J × R is efficiently computable using our leverage score sampling data structure. We require time O ( JR 2 ) to multiply it by V Σ -1 and compute the singular value decomposition of the product to get the condition number. Next observe that A ⊤ A = V Σ 2 V ⊤ , so we can recover V and Σ -1 by eigendecomposition of A ⊤ A ∈ R R × R in time O ( R 3 ) . Finally, recall the formula

used at the beginning of Section 3 that enables computation of A ⊤ A in time O ( ∑ N j =1 I j R 2 ) without materializing the full Khatri-Rao product. Excluding the time to form SA (which is given by Theorem 1.1), κ ( SQ ) is computable in time

Plugging κ ( SQ ) into Equation (12) gives an efficient method to compute the distortion.

A.13 Supplementary Results

A.13.1 Comparison Against Standard CP Decomposition Packages

Table 4 compares the runtime per ALS round for our algorithm against existing common software packages for sparse tensor CP decomposition. We compared our algorithm against Tensorly version 0.81 [14] and Matlab Tensor Toolbox version 3.5 [3]. We compared our algorithm against both non-randomized ALS and a version of CP-ARLS-LEV in Tensor Toolbox.

As demonstrated by Table 4, our implementation exhibits more than 1000x speedup over Tensorly and 295x over Tensor Toolbox (non-randomized) for the NELL-2 tensor. STS-CP enjoys a dramatic speedup over Tensorly because the latter explicitly materializes the Khatri-Rao product, which is prohibitively expensive given the large tensor dimensions (see Table 3).

STS-CP consistently exhibits at least 2.5x speedup over the version of CP-ARLS-LEV in Tensor Toolbox, with more than 10x speedup on the Amazon tensor. To ensure a fair comparison with CPARLS-LEV, we wrote an improved implementation in C++ that was used for all other experiments.

A.13.2 Probability Distribution Comparison

Figure 7 provides confirmation on a small test problem that our sampler works as expected. For the Khatri-Rao product of three matrices A = U 1 ⊙ U 2 ⊙ U 3 , it plots the true distribution of leverage scores against a normalized histogram of 50,000 draws from the data structure in Theorem 1.1. We choose U 1 , U 2 , U 3 ∈ R 8 × 8 initialized i.i.d. from a standard normal distribution with 1% of all entries multiplied by 10. We observe excellent agreement between the histogram and the true distribution.

0

100

200

300

400

500

0.000

0.005

0.010

0.015

0.020

0.025

Density

True Leverage Score Distribution

Histogram of Draws from Our Sampler

Row Index from KRP

A.13.3 Fits Achieved for J = 2 16

Table 5 gives the fits achieved for sparse tensor decomposition for varying rank and algorithm (presented graphically in Figure 4). Uncertainties are one standard deviation across 8 runs of ALS.

A.13.4 Fit as a Function of Time

Figures 8a and 8b shows the fit as a function of time for the Amazon Reviews and NELL2 tensors. The hybrid version of CP-ARLS-LEV was used for comparison in both experiments. As in section 4.3, thick lines are averages of the running max fit across 4 ALS trials, shown by the thin dotted lines. For Amazon, the STS-CP algorithm makes faster progress than CP-ARLS-LEV at all tested sample counts.

For the NELL-2 tensor, STS-CP makes slower progress than CP-ARLS-LEV for sample counts up to J = 163 , 840 . On average, these trials with CP-ARLS-LEV do not achieve the same final fit as STS-CP. CP-ARLS-LEV finally achieves a comparable fit to STS-CP when the former uses J = 196 , 608 samples, compared to J = 65 , 536 for our method.

0

200

400

600

800

Time (s)

0.365

0.370

0.375

0.380

0.385

0.390

Fit

STS-CP (ours), J=65,536

(a) Amazon

0

20

40

60

80

100

120

Time (s)

0.0600

0.0625

0.0650

0.0675

0.0700

0.0725

0.0750

0.0775

0.0800

Fit

STS-CP (ours), J=65,536

(b) NELL-2

A.13.5 Speedup of STS-CP and Practical Usage Guide

Timing Comparisons. For each tensor, we now compare hybrid CP-ARLS-LEV and STS-CP on the time required to achieve a fixed fraction of the fit achieved by non-randomized ALS. For each tensor and rank in the set { 25 , 50 , 75 , 100 , 125 } , we ran both algorithms using a range of sample counts. We tested STS-CP on values of J from the set { 2 15 x | 1 ≤ x ≤ 4 } for all tensors. CPARLS-LEV required a sample count that varied significantly between datasets to hit the required thresholds, and we report the sample counts that we tested in Table 6. Because CP-ARLS-LEV has poorer sample complexity than STS-CP, we tested a wider range of sample counts for the former algorithm.

For each configuration of tensor, target rank R , sampling algorithm, and sample count J , we ran 4 trials using the configuration and stopping criteria in Appendix A.11. The result of each trial was a set of ( time , fit ) pairs. For each configuration, we linearly interpolated the pairs for each trial and averaged the resulting continuous functions over all trials. The result for each configuration was a function f T ,R,A,J : R + → [0 , 1] . The value f T ,R,A,J ( t ) is the average fit at time t achieved by algorithm A to decompose tensor T with target rank R using J samples per least squares solve. Finally, let

be the speedup of STS-CP to over CP-ARLS-LEV (hybrid) to achieve a threshold fit P on tensor T with target rank R . We let the threshold P for each tensor T be a fixed fraction of the fit achieved by non-randomized ALS (see Table 5).

Figure 9 reports the speedup of STS-CP over hybrid CP-ARLS-LEV for P = 0 . 95 on all tensors except Enron. For large tensors with over one billion nonzeros, we report a significant speedup anywhere from 1.4x to 2.0x for all tested ranks. For smaller tensors with less than 100 million nonzeros, the lower cost of each least squares solve lessens the impact of the expensive, more accurate sample selection phase of STS-CP. Despite this, STS-CP performs comparably to CP-ARLS-LEV at most ranks, with significant slowdown only at rank 25 on the smallest tensor Uber.

25

50

75

100

125

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Speedup over LK to attain 95% of Exact

Amazon

25

50

75

100

125

Reddit

(a) Large Tensors

25

50

75

100

125

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Speedup over LK to attain 95% of Exact

Uber

25

50

75

100

125

NELL-2

(b) Small Tensors

On the Enron tensor, hybrid CP-ARLS-LEV could not achieve the 95% accuracy threshold for any rank above 25 for the sample counts tested in Table 6. STS-CP achieved the threshold accuracy for all ranks tested . Instead, Figure 10 reports the speedup to achieve 85% of the fit of non-randomized ALS on the Enron. Beyond rank 25, our method consistently exhibits more than 2x speedup to reach the threshold.

Guide to Sampler Selection. Based on the performance comparisons in this section, we offer the following guide to CP decomposition algorithm selection. Our experiments demonstrate that STS-CP offers the most benefit on sparse tensors with billions of nonzeros (Amazon and Reddit) at high target decomposition rank . Here, the runtime of our more expensive sampling procedure is offset by reductions in the least squares solve time. For smaller tensors, our sampler may still offer significant performance benefits (Enron). In other cases (Uber, NELL-2), CP-ARLS-LEV exhibits better performance, but by small margins for rank beyond 50.

25

50

75

100

125

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup over LK to attain 85% of Exact

Enron

STS-CP reduces the cost of each least squares solve through a sample selection process that relies on dense linear algebra primitives (see Algorithms 3 and 4). Because these operations can be expressed as standard BLAS calls and can be carried out in parallel (see Appendix A.10, we hypothesize that STS-CP is favorable when GPUs or other dense linear algebra accelerators are available.

Because our target tensor is sparse, the least squares solve during each ALS iteration requires a sparse matricized-tensor times Khatri-Rao product (spMTTKRP) operation. After sampling, this primitive can reduced to sparse-matrix dense-matrix multiplication (SpMM). Development of accelerators for these primitives is an active area of research [28, 26]. When such accelerators are available, the lower cost of the spMTTKRP operation reduces the relative benefit provided by the STS-CP sample selection method. We hypothesize that CP-ARLS-LEV, with its faster sample selection process but lower sample efficiency, may retain its benefit in this case. We leave verification of these two hypotheses as future work.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies fast, exact leverage score sampling for Khatri-Rao product\nmatrices (given access to its factor matrices) and an application to CP tensor\ndecomposition via alternating least squares with row sampling. The main\ntheoretical contribution is a data structure for sampling from the leverage\nscore distribution of a Khatri-Rao product by iteratively sampling from the\ncorrect conditional marginal distribution as it constructs the row sample (one\nfactor at a time). To apply this most effectively to CP decompositions, the\nauthors propose a two-step version of this sampler that first samples a rank-1\neigenspace of the partial Gram factor matrices. The authors provide a good\ncomparison of their work to the leverage score-based CP decomposition algorithm\nof Larsen and Kolda (Journal of Matrix Analysis and Applications, 2022).\n\nSTRENGTHS:\n- The sparse factor matrix version of Theorem 1.1 complements the paper.\n- The comparison with Woodruff-Zandieh [27] is valuable and discusses the\n  shortcomings of the theory-heavy results (i.e., hidden second-order constants\n  that blow up).\n- The paper is very nicely organized.\n- The two-stage eigenspace sampling idea in Section 3.2 is a solid research\n  contribution.\n\nWEAKNESSES:\n- Theorem 3.1 is one of the main theory components of the paper, but largely\n  comes from [Malik, ICML 2022].\n- Section 3.1 could benefit from a formal description of what the leaf nodes\n  represent. In particular, what is $S_{0}(v)$ and how does this partitioning\n  get decided (i.e., where to draw the cut points)? It seems reasonable that the\n  result is correct, just not immediately implementable by the reader.\n- In Corollary 3.3, it would be better to use the original expression $R=O(R\n  \\max\\{\\log(R/\\delta), 1/(\\varepsilon \\delta)\\})$ samples so as to not cause\n  confusion by the implicit case assumption.\n- Re experiments: It would be good to compare the running time of this\n  sampling-based method with both ALS-ARLS-LEV and an out-of-the-box ALS\n  implementation (e.g., MATLAB or Tensorly). Analyzing the fit of the\n  decomposition in isolation doesn't tell the full story.\n\nQUESTIONS:\n**Questions**\n- Re Theorem 1.1: If a single entry of $U_j$ changes, is there a faster update\n  time?\n- Should $I$ in Table 1 be $I_j$ to represent the $j$-th factor update?\n- In Line 81, we have $S \\in \\mathbb{R}^{J \\times I}$, so shouldn't we have $S[j,i]$ in\n  the following line?\n- In Line 95, what is the tilde in the lower bound for $J$? Related: Line 99\n  should be $\\Omega(R/(\\varepsilon \\delta))$ since this is a lower bound.\n- In the discussion about Kronecker regression, Cheng et al. [4] and\n  Larsen-Kolda [12] sample from the leverage score distribution of the\n  Kronecker product for their CP decomposition algorithms (i.e., the product of\n  the leverage scores). This connection seems worth pointing out.\n- In Line 337, shouldn't the fixed sample count increase with the rank? This is\n  needed for the leverage score sampling guarantees.\n\n**Typos and suggestions**\n- [line 14] \"denoted\" --> \"denoted by\"\n- [line 43] Suggestion: Can remove the parenthesis around the summand in the\n  second part of Theorem 1.1. The sentence \"The structure can also draw\n  samples...\" can probably be generalized to say \"... for any subset of factor\n  matrices\"?\n- [line 49] You say that \"our applications deal with dense inputs,\" but this is\n  immediately followed by mention of the sparse Amazon tensor experiments. The\n  writing here could be improved.\n- [line 55] Suggestion: Add a citation for CP-ARLS-LEV after you first mention\n  it.\n- [line 58] Table 1 could benefit from $\\tilde{O}$ notation to account for\n  missing constant factors. The description of what the complexity is for could\n  be made more explicit too: \"Complexity of factor matrix $U_j$ for\n  $N$-dimensional dense CP decomposition...\"\n- [line 64] suggestion: $j$'th --> $j$-th, same for later occurrences\n- [line 79] suggestion: \"sampling operators\" --> \"row sampling operators\"\n- [line 131] \"Kronecker regression is distinct\" --> \"is a distinct\"\n- [line 141] suggestion: \"autoregressive fashion\" is an indirect way to explain\n  the procedure. Consider dropping this phrase and merging the two sentences.\n- [line 143] $I_n$ have been scalars so far. Therefore, it is better to use\n  $(i_1, \\dots, i_N) \\in [I_1] \\times \\dots \\times [I_N]$ to denote the set of\n  indices.\n- [line 144] suggestion: consider restating the dimension of $G_k$ and $G$ to\n  help the reader remember what some of these operators mean, e.g.,\n  $G := (\\text{expression}) \\in \\mathbb{R}^{R \\times R}$.\n- [line 153] Please include the citation to Malik 2022 too, so that the reader\n  can easily click through to the references.\n- [line 162] Typo: \"theorem 1.1\" --> \"Theorem 1.1\"\n- [line 165] Suggestion: Rewrite the sentence as: Let $h \\in \\mathbb{R}^{R}$ be\n  a vector and let $Y \\in \\mathbb{R}^{R \\times R}$ be a ... [delete\n  \"respectively\"].\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies algorithms to efficiently sample a row from a matrix $A = U_1 \\odot U_2 \\odot ... \\odot U_N$ with probability proportional to the leverage scores of $A$.\nHere, $B \\odot C$ denotes the _Khatri-Rao_ product of $B$ and $C$, or the column-wise kronecker-product of the columns of $B$ and $C$.\nIn particular, this means that $A$ has a small number of columns $R$, but has an exponential number of rows $\\prod_{j=1}^N |I_j|$ (where $U_j \\in \\mathbb{R}^{|I_j| \\times R}$).\n\nThis exponential blowup in the number of rows of $A$ in the core computational focus on this paper, as well as several related papers in the area.\nThis paper's technical contribution is in proposing a data structure that can preprocess $U_1, ..., U_N$ such that a row of $A$ can be sampled with probability proportional to its leverage score.\nThis data structure's space complexity matches the sizes of the input, and at query time it takes only $O(NR^2 \\log \\max \\\\{I_j, R\\\\})$ amortized time per row sampled.\nThis new data structure is based on \"binary-tree inversion sampling\", a binary search algorithm that can help sample from distributions efficiently, and is used to speed up a computational bottleneck in a known approach to sample leverage scores from Khatri-Rao products.\n\nTheoretical results demonstrate the correctness of this algorithm, and experiments demonstrate the effectiveness of this algorithm.\n\nSTRENGTHS:\nThe paper is well written, nicely motivated, pretty easy to understand, has a clear presentation of the subproblem it speeds up, and provides good context relative to existing work in the area (afaik; I'm no expert in tensor algorithms).\n\n**Originality and Quality:** A fine line needs to be drawn here, and this is the subtlest point of my review. Per the authors' account, existing work in leverage score sampling from Khatri-Rao products exists and provides nice results. A good example is Theorem 3.1 from this paper, which is an adaptation of an observation made by prior work about a particular way to compute the leverage scores of $U_1 \\odot ... \\odot U_N$. The core of this paper is taking Theorem 3.1 and finding a faster way to compute it's right hand side. This means that the prior work gave much of the framework for around paper's proposed algorithm, and we can sorta think of this paper's contribution as finding a novel way to compute a subroutine more quickly.\n\nThat said, the approach to computing the subroutine is novel and interesting. I've not seen prior work in this area use a binary-tree in such a way. There's clearly a good deal of effort put into the math (though I only verified bits and pieces of the appendix). The paper very much stands as sufficiently original, but I want to be clear that the originality lands firmly in section 3.1 of the paper -- how to efficiently sample following this framework that existed in prior work.\n\n\n**Clarity:** The paper is well written. I have only a couple mild gripes about writing and presentation (listed later in the \"Questions\" section), and I feel that I understood all the big ideas of this paper.\n\n\n**Significance:** This paper feels well motivated, and the improvement to sample a row with runtime that depends logarithmically on the number of rows in $A$ is quiet strong and interesting. The experiments suggest that this truly is a state-of-the-art algorithm for a meaningful suite of metrics. I do have some gripes with the presentation of the experiments, which undercuts my confidence in the experimental evidence a smidge, but I think this both not essential and can be easily corrected.\n\nWEAKNESSES:\nThe weaknesses are few and far between in this paper. A couple section lack a bit of clarity, like how the big-Oh rates for some of the downstream tensor algorithms are insufficiently explained. These are issues that can easily be fixed for a camera-ready version of the paper.\n\nOut of these mild gripes, there are three that do stand out a bit:\n\n1. Several experiments lack confidence intervals. This muddies the story of the algorithm's empirical efficacy, especially in Figure 3, where it feels very confusing to see the error of the green curve in the left figure jump up significantly.\n\n2. In a few place, discussion of big-Oh runtime isn't fully explained. For instance, I don't understand how the complexity of STS-CP was derived on Table 1 (appendix A.1). More details are in the \"Questions\" section.\n\n3. The last paragraph from section 3.2 takes the core technical ideas and is supposed to summarize the algorithm. However, it instead introduces some new notation and makes me wonder if it's actually summarizing an algorithm or if it's instead trying to briefly explain some of the deeper technical edgecases the algorithm has to handle. Either way, it's pretty confusing of a paragraph for me to read, and undercuts my confidence in understanding the full algorithm.\n\nNone of these are serious gripes, and call all be fixed with mild updates to the presentation. So I'm still very positive on this paper!\n\nQUESTIONS:\nI don't really have deep concerns to ask about here. I'll just enumerate a long list of minor typos and confusions.\n\n1. Why use (afaik) non-standard notation, like having $d$ be the dimension (number of cols) of A, or $p_i$ be the probability of sampling row $i$, or $m$ or $c$ or $s$ being the number of sampled rows?\n1. [Line 30] As someone familiar with leverage scores, I was surprised not to see a log appear in the sample complexity of leverage score sampling. I'd make it into a $\\tilde O$.\n1. [Lines 49 and 52] Line 49 says that the applications mainly deal with dense inputs. Line 52 says that the most practical benefit is on sparse matrices. These aren't necessarily statements that are at-odds, but it does feel like a bit of whiplash to read these statements back-to-back. The argument on lines 279-280 would be good to add here.\n1. [Line 55] Do you have any understanding of why your algorithm which boasts a new and much smaller big-Oh term has 2% slower runtime?\n1. [Table 1] I dunno how you came to the complexity-per-iteration here. As a reader, I felt like I should be able to pattern-match between Theorem 1.1 and Table 1 to understand the complexity of STS-CP, but I really couldn't make them line up.\n1. [Line 131] \"is **a** distinct\"\n1. [Line 132] \"There\" not \"Here\"\n1. [Line 194] Get rid of the square on $[0,1]^2$\n1. [Line 132] Explicitly argue why $F=1$ is affordable space-wise now, but was too expensive in the setting at the start of section 3.2\n1. [Line 241] What is $Z_j$?\n1. [Figures 2,3,5] Add confidence intervals\n1. [Figure 4] Swap the x and y axis. It's hard to read with fit being on the x-axis. Add some space between the subplots -- they're real strange shoved together like that, making it harder to read.\n1. [Line 344] Well... STS-CP does have an oscillating error pattern on Amazon too. Maybe mention it? Eh, I'm sorta torn here.\n1. [Algo 6, line 5] Missing a period in $1,..,N$\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

comments: SUMMARY:
This work proposes a new algorithm to efficiently perform exact leverage score sampling on Khatri-Rao products. This work is built on top of the TNS-CP algorithm (Malik et al., 2022), and a more efficient data structure is used to achieve better sampling computational cost. This algorithm can be used to accelerate sketching based alternating least squares (ALS) for CP decomposition, and it has been shown that state-of-the-art complexity per ALS sweep has been achieved. Experimental results show that this sampler is efficient and yields better accuracy than previous leverage score sampling based CP-ALS algorithms for large real sparse tensors.

STRENGTHS:
1. Accelerating leverage score sampling and large scale CP decomposition is an important topic, and this work achieves state-of-the-art results in both the theoretical analysis and the experimental results.

2. The data structure used to accelerate leverage score sampling of Khatri-Rao product is novel.

WEAKNESSES:
Overall I believe this is a good contribution. I only have one minor comment: it takes me a while to figure out the logics in Sections 3.1 and 3.2, and I think adding a figure to summarize the two sampling process would be good.

QUESTIONS:
n/a

LIMITATIONS:
n/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper designs a method for efficiently minimizing least squares regression, when the design matrix is equal to the Khatri-Rao product of multiple matrices. Because this Khatri-Rao product can have a huge dimension, the paper is motivated to study sampling techniques to reduce the dimension of the product matrix.\n\nThis least squares problem is motivated by applications to tensor decomposition, in particular when alternating least squares is used for finding CP tensor dimension. Each step of the alternating least squares corresponds to a least squares regression, where the design matrix is the Khatri-Rao product of the other tensor factors in the CP dimension.\n\nTo solve this sketching problem, the paper adapts the leverage score sampling technique to this Khatri-Rao product design matrix. The algorithm scales in a linear order proportionally to the sum of the dimensions of each matrix in the product.\n\nExperiments are provided, comparing the approach to methods from a prior work by Larsen and Kolda (SIAM J. Matrix Analysis and Applications (2022)). The results show moderate improvement in terms of sparse tensor decomposition.\n\nSTRENGTHS:\nS1) The main is mainly of a technical nature: the algorithm, which is based on detailed calculations of the Khatri-Rao product, appears sound, and this is validated in the experiments.\n\nS2) The sampling from Khatri-Rao product relies on a binary-tree inversion sampling technique, and this tree construction has time cost $O(I R^2)$ and storage cost $O(R^2 I / F)$.\n\nS3) The experiments compared against the method of Larsen and Kolda appear significant, especially at higher target rank values, for several sparse tensors with as many as $10^9$ nonzeros.\n\nWEAKNESSES:\nW1) From my reading, the paper does not do a good job of motivating the problem. For instance, the main motivating example is tensor decomposition, but the connection is spelled explicitly only until Section 3.3. Are there any other plausible examples for motivating the proposed problem?\n\nW2) I find the derivation steps to be quite laborious but also quite tedious. I find the illustration of Figure 1 to be very useful, and I think that a better presentation that conveys the key steps would be very helpful, for instance, in a special case of the Khatri-Rao product.\n\nW3) The significance of the results is quite limited; in my opinion, I think the paper could benefit from stronger statements about how significant the contributions are. Solely judged based on the experiments, the improvement is quite weak since only one baseline is considered in the experiments. From a theoretical standpoint, the comparison stated in the related work could be made more clear as well.\n\nQUESTIONS:\n- It would be nice to give some description of the tensor used in the experiments (without having to refer to the appendix).\n- I understand this question may be out of scope, but it would be nice to compare fit against other tensor decompositions, e.g., the Tucker decomposition. Would the newly developed techniques apply to Tucker decomposition?\n- It is mentioned in the related work that this paper is very closely related to a method by Woodruff and Zandieh (2022). I wonder how the methods would compare against each other in the experiments.\n\nLIMITATIONS:\nYes, the paper discussed limitations in Section 5. Due to the technical nature, I do not see any potential societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}