Fast Exact Leverage Score Sampling from Khatri-Rao Products with Applications to Tensor Decomposition
Authors: Vivek Bharadwaj
Osman Asif Malik
Riley Murray
Laura Grigori
Aydin Buluc
James Demmel
Venue: NeurIPS 2023 poster (2023)
Keywords: Tensor Decomposition
Leverage Scores
Randomized Linear Algebra
Sketching
Khatri-Rao Product
Sparse Tensors
TL;DR: We built a fast data structure to sample rows from a Khatri-Rao product according to its exact leverage scores, which leads to asymptotically faster and practically more accurate randomized CP 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 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
Vivek Bharadwaj 1.2 Osman Asif Malik 2, Riley Murray 32,1, Laura Grigori 4, Aydın Buluç 21, James Demmel
Electrical Engineering and Computer Science Department, UC Berkeley CComputational Research Division, Lawrence Berkeley National Lab IInternational Computer Science Institute
Institute of Mathematics, EPFL & Lab for Simulation and Modelling, Paul Scherrer Institute
Abstract
We present data structure 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 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 result, tractably draws samples even when the matrices forming the Khatri-Rao product have tens of millions of rows each. When used sketch the linear least squares problems arising in CANDECOMP / PARAFAC tensor decomposition, our method achieves lower asymptotic complexity per solve than recent state-of -f-the-rrt methods. Experiments on billion-scal scale sparse tensors validate our claims, with our algorithm achieving higher accuracy than competing methods as the decomposition rank grows.
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 minx |A B BlI where the design matrix U1 C.... O UN is the Khatri-Rao product of matrices U, E R x 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 [33. In this work, we focus on the case where A has moderate column count (several hundred most). Despite this, the problem remains formidable because the height of A is II, Iy. For row counts I, in the millions, it is intractable to even materialize A explicitly.
Several recently-proposee randomized sketching algorithms can approximately solve least squares problems with Khatri- product design matrices 4 12 15
18 29]. These methods apply sketching operator S the design and data matrices to solve the reduced least squares problem
minx SAX - SB where S has far fewer
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 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 Ô (R/(Ed)) samples are required (subject to the assumptions at the end of section 2.I). The challenge, then, efficiently sample according to the leverage scores when A has Rao structure.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

We propose a leverage-score sampler for the Khatri- Chatri-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 U1,-.. UN with Un RR there exists data structure satisfying the following:
1. The data structure has construction time O E===I,,² and requires additional storage space (Ei= I;
If single entry in matrix Uj changes, it can be updated in time O log R)) the entire matrix U3 changes, it can be updated in time O (1₂R²).
2. The data structure produces J samples from the Khatri-Rao product U1 UN according to the exact leverage score distribution its rows in time
o N +
Σ R² log max (Ia, R) k=1
using O(T²) scratch space. The structure can also draw samples from the Khatri-Rao product of any subset of U1.... UN.
The efficient update property and ability to exclude one matrix are important in CP decomposition. When the inputs U1. .... UN are sparse, analogous data structure with O R \ nnz(U,)
construction time and O [-== mnz(U,) storage space exists with identical sampling time. Since the output factor matrices U1,..., UN are typically dense, we defer the proof to Appendix AA.8 Combined with error guarantees for leverage-score sampling, achieve an algorithm for alternating least squares CP decomposition with asymptotic complexity lower than recent state-of-the-art methods (see Table T
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, algorithm STS-CP can achieve 95% of the fit of ron-manddd ALS between 1.5x and 2.5x faster than a gigh-perforaacce implementation of the state-of- state-of-the-ar ttthher-tt CP-ARLS-LES 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 --V with even 54 times many samples.
Table I: Asymptotic Complexity to decompose an N- d-dimensional I x ... I dense tensor via CP alternating least squares. For randomized algorithms, each approximate least- squares solution has residual within (1 + 8) of the optimal value with high probability 1 5. Factors involving log Rt and (og( are hidden (Ō notation). See A.1 for details.
ALGORITHM
COMPLEXITY PER ITERATION
CP-ALS 13]
N(N + --R N(R 1)R ((zō)
CP-ARLS-LEV 15
TNS-CP [19]
2 (cō)
GAUSSIAN TNE 17
(N (1
1/2
STS-CP (OURS)
N(NR log I (RR2)
2 Preliminaries and Related Work
Notation. We use [N) to denote the set [1,..., NY for positive integer N. We use Ô notation to indicate the presence of multiplicative terms polylogarithmic R and (1/8) runtime complexities. For the complexities of our methods, these logarithmic factors are more than (((og( We

---

use Matlab notation A [:. to index rows, resp. columns, of matrices. For consistency, we use the convention that A [ row vector. We use for standard matrix multiplication, @ as the elementwise product, 69 to denote the Kronecker product, and @ for the Khatri-Rao product. See Appendix A.2 for definition of each operation. Given matrices E TR x B mR the j-th column of the Khatri-Rao product A o B E R R x is the Kronecker product A [: @ B [:
We use angle brackets (:.....) to denote a generalized inner product. For identically-sizee vectors matrices, it returns the sum of all entries in their elementwise product. For A, B, C' E R=*xx
(A, = Σ AL..B[C.
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 minx ||AX BII when A has no special structure 30 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 [1], or Countsketch- type operators [. For tensor decomposition, however, the matrix B may be sparse or implicitly specified as black-box function. When B is sparse, Countsketch-tyee algorithm to iterate over all nonzero values in B. As Larsen and Kolda [S]] note, operators similar to the FFT induce fill-in when applied to 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 .,.,, æ, be a selection of J indices for the rows of A E 2 sampled i.i.d. according to a probability distribution Q1,..., Qr- The associated sampling matrix S E Jxl specified by
operators still require the
if =
S $
:.n. 0,
otherwise
where the weight of each nonzero entry corrects bias induced by sampling. When the probabilities gj 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 matrix assign a measure of importance to each of its rows. The leverage score of row from matrix A R is given by
e, A [i. (AT A)+
for S i Leverage scores can be expressed equivalently as the squared row norms the matrix Q in any reduced QR factorization of A [ 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 the residual of the original problem.
Theorem 2.1 (Guarantees for Leverage Score Sampling). Given A R and E, E (0, 1), let S J be leverage score sampling matrix for A. Further define X = arg min ||SXX SBI F- If J Z max (log (R/8). 1/(=ō)), then with probability at least - 8 it holds that
| B < (1++ min |AX- - Bllr-
X
For the applications considered in this work, R ranges up few hundred. As e and ő tend to 0 with fixed R, 1/(5 dominates log( R/5). Hence, we assume that the minimum sample count J to achieve the guarantees of the theorem is (2RR(E=)).
3

---

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 = U1 O ... @ Un due to the height of the Khatri- product. Cheng et al. [ instead approximate each score as product of leverage scores associated with each matrix Uj- Larsen and Kolda [IS] propose CP-ARL---LEE
which uses 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-Rat product of N matrices, both require OR (e))) samples to achieve the (s, 5) 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 It and N increase. The TNS-CP algorithm by Malik et al. [9] samples from the exact leverage score distribution, thus requiring only
O O(R( (s)) samples per least squares solve. Unfortunately, it requires time O (Σi= 1,1²²) each sample.
to draw
Comparison to Woodruff and Zandich. 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 a prior oblivious method by Ahle et al. M. which sketches a Khatri-Rao product using 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 E nnz(UU)) and requires an embedding dimension quadratic in R to achieve the (s, 5) solution-ques quallty guarantee. Woodruff and Zandieh use a collection of these tree sketches, each with carefully-controlln runtime dependence on the column count R. On the other hand, the method exhibits O(N") scaling in the number of matrices involved, has O(e scaling 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 selicientl-uuddatabble sampler with runtime linear in both N and = These properties are attractive when the column count R is below several thousand and when error as low as E R 10-3 is needed in the context of an iterative solver (see Figure 5p. Moreover, the term O Dj=1 I2) in our construction complexity arises from symmetric rank-k updates, highly-optimized BLAS3 kernel on modern CPU and GPU architectures. Appendix A.3 provides a more detailed comparison between the two approaches.
approximation error, to design an algorithm with linear
Kronecker Regression. Kronecker regression is a distinct (but closely related) problem to the one we consider. There, A U1 @ ... & UN and the matrices U, have potentially distinct column counts Ri.... RN. While the product distribution of leverage scores from U1.... UN 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 moaaiically-updttale tree-sketches 21,, to produce accurate and updatable Kronecker sketching methods. None of these results apply directly in our case due to the distinct propertics of Kronecker and Khatri-Raa products.
3 An Efficient Khatri-Rao Leverage Sampler
Without loss of generality, we will prove part 2 of Theorem I.I for the case where A = U1 o UN; the case that excludes a single matrix follows by reindexing matrices Uk- We further assume that A is nonzero matrix, though it may be rank-deficient. Similar to prior sampling works [88 29). our algorithm will draw one sample from the Khatri-Rao product by sampling a row from each of U1, U..... in sequence and computing their Hadamard product, with the draw from U, conditioned on prior draws from ......4 UZ-1-
Let us index each row of A by a tuple (i1,..., iN) E [hi x... x [Ix]. Equation T gives
lu.... = A ((i1,..,, iN), :] (A T A)* [( [I]
For 1 S k < NN define G3 := U Uv E R Rx and G = + k= Gk
E R Rx it is a well-known
fact that G AT. A [33. For single row sample from A, let $1, ..., SN be random variables for the

---

draws from multi-index set [11] x x [Ix] according to the leverage score distribution. Assume, for some k, that we have already sampled an index from each of ..., [1), and that the first k - 1 random variables take values & = $1:..., 8 sk-1- We abbreviate the latter condition as Sck = Sck- To sample from [s. we seek the distribution of & conditioned on $kk Define
heck as the transposed elementwise product of rows already sampled:
hack =
Us [st,
(3)
==1
Also define G>k as
G>k:= G
Gi.
(4)
i=k+1
Then the following theorem provides the conditional distribution of 5k- Theorem 3.1 (Malik 2022, [8]. Adapted). For any E [Tx).
p(sk 8k Beck SCk) ( Un [| Un [sk: G2*)
(5)
:= Ue. [sk]
where C = (Meckl cki UK Uk. G>k) is nonzero.
We include the derivation of Theorem 3.1 from Equation in Appendix A.4 Computing all entries of the probability vector U.. would cost O R²) per sample, too costly when Uj 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 Uj. since the
conditioning random variable has II | Is potential values. Our key innovation is data structure to sample from a discrete distribution of the form Uh... without materializing all of its entries or
incurring superlinear cost in either or = 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 UYYU,Y
We introduce slight change of notation in this section to simplify the problem and generalize our sampling lemma. Let h E IR be a vector and let Y E Ritx be positive semidefinite (p.s.d.) matrix, respectively. Our task is sample rows from matrix U gxxR according to the distribution
Qu,U, [s] := Chh U [s (s.),,)
(6)
provided the normalizing constant C 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.
UTU,YY Y), nonzero. We impose that all rows are
Lemma 3.2 (Efficient Row Sampler). Given matrices U U R R with Y p.s.d., there exists a data structure parameterized by positive integer F that satisfies the following:
1. The structure has construction time O (IR?) and storage requirement O (R²(I/F)).
1 < F, the storage requirement drops to O(1).
2. After construction, the data structure can produce a sample according to the distribution U.y time ORR [ + FR2) for any vector
3. Y is a rank-l matrix, the time per sample drops to O log [TF] + FR).
This data structure relies an adaptation of classic binary-tree inversion sampling technique 22 Consider a partition of the interval [0, 1] into bins, the i-th having width qu, y [i]. We sample Uniform (0, 1] and return the index of the containing bin. We locate the bin index through binary search terminated when at most F' bins remain in the search space, which are then scanned in linear time. Here, F is tuning parameter that we will usage.
control sampling complexity and space
For > 6, assume that @ (...) produces vector / matrix filled with ones.
5

---

We can regard the binary search as a walk down
full, complete binary tree Tr with [I/F]
(1..8]
leaves, the nodes of which store contiguous, dis-
joint segments S(e) = (()) - [T] of size at most F. The segment of each in- ternal node is the union of segments held by its children, and the root node holds {1.... I]. Suppose that the binary search reaches node with left child L(u) and maintains the inter- val [low, high] - [0, 1] as the remaining search space to explore. Then the search branches left in the tree < low + Eies(L(v) Uh,U,Y [i]
[1..4]
[5..8]]
1.2]
[3,4]
[6,
[7.8]
2:
9x
Figure I: A segment tree T8,2 and probability dis- tribution {91...., Q8} on (1,...88]
This branching condition can be evaluated effi- ciently if appropriate information is stored at each node of the segment tree. Excluding the offset "low", the branching threshold takes the form
Σ (ES(v)
Qu,U,y [i] (hh Σ U... = (hh G).Y). VES(v)
(7)
Here, we call each matrix G E IR I partial Gram matrix. In time O and space O [1/P1), can compute and cache G for each node of the construct our data structure. Each subsequent binary search costs ORR2) time to evaluate Equation (7) at each of log [g[1 internal nodes and O(FR²) to evaluate fb,0, 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 we directly apply Lemma 3.2 to sample from the conditional distribution in Theorem 3.1 Because G> k is not rank-1 general, we must use point of the lemma, where no selection of the parameter allows us to simultaneously satisfy the space and runtime constraints of Theorem Selecting
results in cost O(R)) per sample (violating the runtime requirement in point 2), whereas F = results in superlinear storage requirement O(IR²) (violating the space requirement point and becoming prohibitively expensive for I > 10%). To avoid these extremes, we break the sampling procedure into two stages. The first stage selects a I-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 achieve better time and space complexity.
U..s. and h = heck When sampling from It. observe that G> is the same for all samples. We compute symmetric cigendecomposition G> = VAVT, where each column of V is eigenvector of G>s and A diag (Aw))()) contains the eigenvalues along the diagonal. This allows us to rewrite entries of as
Below, we abbreviate
B q [s] = C-1 Σ Au Un [s Un
(8)
Define matrix W e R4 x elementwise by
W [1 (hhT, Un [6:] Un [4:1) V u [[:uu]
and observe that all its entries are nonnegative. Since A 2 for all u (G>* is p.s.d.), we can write mixture of probability distributions given by the normalized columns of W:
W F. u
where the vector of nonnegative weights is given by w [M] (C-1X,, [[[;,)]]) in this form gives us the two stage sampling procedure: first sample component a of the mixture according to the weight vector then sample an index in {1..k according to the probability vector
w]ll)) Rewriting

---

defined by W [t. ul [W [: uI][ Let ty be a random variable distributed according to the probability
mass vector w. We have, for C' taken from Theorem 3.1
p(ū = =1)) - Xun
W [1, wR]
(9)
=C' Xux (hh
V lt us] V usl" G.)
= =VVVAV .G2 [
Hence, we can use point 2 of Lemma 3.2 to sample a value for úg efficiently. Because VVVT has only Rt rows with Rt 10a, we can choose tuning parameter F achieve lower time per sample while incurring modest (R²) space overhead. Now, introduce random variable fg with distribution conditioned on ü₃ = 44 given by
pf4- = tx ûU = ux) == W [tk, ux. || [t. uxll11
(10)
This distribution is well-defined, since we suppose that ui uy occurs with nonzero probability e [ug]] which implies that ||W [t, ux]ll # O. Our remaining task is efficiently sample from the distribution above. Below, we abbreviate h = V [., usl @ and derive
p(tx = tx ug. = =u)) (hh T, UN [t :] UR [tx, :I, V [t. upl V [t upl uxl")
I| [., ux]]l
(hhT. Up [ta, :] UR [tg, [1]]
(11)
I [s, uxlll
[ [tx]
Based on the last line of Equation sample row of UR. Since Y [1] a rank-1 matrix, can use point of the lemma and select
we apply Lemma 3.2 again to build an efficient data structure
a larger parameter value F = R to reduce space usage. The sampling time for this stage becomes O(R² log R/)
To summarize, Algorithms [ and 2 give the construction and sampling procedures for our data and "RowSample" procedures from Algorithms and
structure. They rely the "BuildSamphe in Appendix A.55 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 (1I.. Construction costs
O (E)=ı ()? and if any matrix U, changes, we can rebuild Zj in isolation. Because F = R.
the space required for Z is O (T, In the sampling phase, the procedure in Algorithm accepts an optional index of a matrix to exclude from the Khatri-Rao product. The procedure begins by computing the symmetric eigendecomposition of each matrix G>+- The eigendecomposition computed only once per binary tree structure, and its computation samples. It then creates data structures Eg for each of the distributions specified by Equation (9). These data structures (along with those from the construction phase) are used to draw ûg and fę in succession. The random variables ig follow the distribution Theorem 3. conditioned on prior draws, so the A.6 proves the complexity claims in the theorem and provides further details about the algorithms.
is amortized over all J
(fx kez follows the leverage score distribution on A, as desired. Appendix
3.3 Application to Tensor Decomposition
A tensor is multidimensional array, and the CP decomposition represents tensor sum of outer products [13]. See Appendix A.9 for an overview. To approximately decompose tensor the popular alternating least squares (ALS) algorithm begins with randomly initialized factor matrices Uj. U, E x for 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 single factor matrix while holding the others constant. The j-th least squares problem occurs in the update
T E R/xx...xxx 4
U, = arg min U4z X
mat(T,))
where Uz₁ = UN &... @ U @ U, o... & U, the Khatri-Rao product of all matrices excluding Ug and mat(-) denotes the mode-j matricization of tensor T. Here, we reverse the order of matrices
7

---

in the Khatri- -aao 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 12 15 and applying the data structure in Theorem I.1 gives us the STS-CP algorithm.
Corollary 3.3 (STS-CP). Suppose T is dense, and suppose we solve each least squares problem ALS with a randomized sketching algorithm. A leverage score sampling approach as defined in section 2 guarantees that with (R/ (Eō) samples per solve, the residual of each sketched least squares problem is within (1 + E) of the optimal residual with probability (1 - 8). The efficient
sampler from Theorem 1.1 brings the complexity of ALS
Σ
(Nn I, + 1,
E6
where "#ii"" the number of ALS iterations, and with any term log 15 replaced by log R 1f < R.
The proof appears in Appendix A.9 and com- bines Theorem [[.]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 de- pends heavily on the sparsity structure and is difficult to predict. Nevertheless, we expect significant speedup based on prior works that use sketching to accelerate CP decomposition
Algorithm 1 ConsttcckRPPSamplerCC.... UN)
1: for 1..N do 2:
Z, = BuulldSamplerU [1)
3: G, U, U
Algorithm 2 >KRPSampplej.. J)
1: G = Øzzj G2
2: for k j
3: G>a:= G+ kkjtt1 Gk
N
4 Experiments
4: Decompose G> = ViAA VT
5: Eg = BuildSampler A V/ F
Experiments were conducted on CPU nodes of NERSC Perlmutter, HPE Cray EX supercomputer, and our code is available at httcps://github... com www/wbaaradwaf... vbharadwajj Bast_tensor__eeerrgge git On tensor de- composition experiments, we compare our algo- rithms against the random and hybrid versions of CP-ARLS-LEV proposed by Larsen and Kolda
6: for -1.JJ do
7:
h [1, ".... 1]
8:
for k + j do
9:
ûg RowSample Ex, h)
10:
lR RowSample Zx, h @ (Vr [:, ix)))
11:
h UR :
[15]. These algorithms outperform uniform sam- pling and www-norm-squrred sampling, achiev- ing excellent accuracy and runtime relative to exact ALS. contrast to TNS-CP and the Gaus- sian tensor network embedding proposed by Ma and Solomonik (see Table 1), CP-ARLS-LEV one of the few algorithms that can practically decompose sparse tensors with mode sizes in the millions. In the worst case, CP-ARLS-LEV requires Ö(R~-1/ (s=)) samples per solve for an N-dimensional tensor to achieve solution guarantees like those in Theorem 2.1 compared to Ö(R((sd)) samples required by STS-CP. Appendices 10 A.11 and A. 13 provide configuration details and additional results.
12:
8d = (Ix)sz)
13:
return 81,
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 [.I is reflected in our performance measurements, with the exception of the plot that varies R. Here, construction becomes disproportionnatlly cheaper than sampling due to cache-efficient BLASS calls during construction. Even when the full Khatri-Rao product has as 3.78 x 1022 rows (for I 225 N = 3, R 32), we require only 0.31 seconds on average for sampling (top plot, rightmost points).
8

---

10
32.2
10-1
10-9
304
100
10a
109
157
[m2: New
0.6
Becibal
100
Proaxt ApSIVEEEIIO Sampler
128
N
13-4
+ Cuastr ariin
Sunpling
Figure 2: Average time (5 trials) construct our proposed sampler and draw samples from U1 @...UUN. with U, E R/x Vj. Error bars indicate 3 standard deviations.
Figure 3: Distortion and residual error (50 tri- als) for varying and N on least squares, / 2 J = 5000. "X" marks indicate outliers 1.5 times the interquartile range beyond the median, stars indicate means.
50. 000
4.2 Least Squares Accuracy Comparison
We now test our sampler on least squares problems of the form minz ||Ar - bll, where A = U1 @ @ UN with Uy E R for all j. We initialize all matrices U, entrywise i.d. from standard normal distribution and randomly multiply 1% of all entries by 10. We choose b as a Kronecker product Cj X ... @ CN, with each vector E R 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 only permitted on-demand, black-box access to its entries.
For each problem instance, define the distortion of our sampling matrix S with respect to the column space of A as
D(S, = x(SQ)
(12)
x(SQ) +
where Q is an orthonormal basis for the column space of A and k(SQ) is the condition number SQ. A higher-quality sketch S exhibits lower distortion, which quantifies the preservation of distances from the column space of the column space of SA [20]. For details about computing x(5Q)
efficiently when A is a Khatri-Rao product, see Appendix A.12 Next, define &
residualapees 1,
residualies
where residualapprox 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 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 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 decompose several large sparse tensors from the FROSTT collection 25 (see Appendix A.I1 for more details on the experimental configuration). Our accuracy metric the tensor fit. Letting T be our low-rank CP approximation, the fit with respect to ground-truth tensor T
fit( (i,(T,
7
F [|]]
Table 4 in Appendix A. 133 compares the runtime per round of our algorithm against the Tensorly Python package 144 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 the rank increases for

---

Earon"
0.20
5.24
0.07
0.22
0.05
0.00
75 Thrgett
Thrget
Trrpe
Amazzo (ell
Reddic (=d.7e6 nx)
3.10
(yyherd)
10
touns)
5.30
0.08
Figure 4: Average fits (8 trials) achieved by randomized (J = 21 and exact ALS for sparse tensor CP decomposition. Error bars indicate standard deviations. See Appendix A. A.
for details.
Uber
Amazon
0.100
10
0.955
0.9
CP CRVAEEUUT J- 60.02
M
GIMALSSEE 1-180.000
5.05
20
30
60
10
20
30
LST50 Problems Number
Prabblm
0.000
Figure Average (5 runs) for randomized least squares solves in 10 ALS rounds, R 50.
Figure Fit vs. time, Reddit tensor, 100. Thick lines are averages trial interpolations.
fixed sample count. By contrast, STS-CP improves the fit consistently, with significant improvement at rank 125 over CP-ARLS-LEV. Timings for both algorithms are available in Appendix 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, compute the exact solution to each least squares problem before updating the factor matrix with a randomized algorithm's 's solution. Figure[]plote e as ALS progresses for hybrid CP-ARLS-LEEY 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 = 26 against ---VRL--LEV with progressively larger sample count. Even with 2 samples per randomized least squares solve, CP-ARLS LEV cannot achieve the maximum fit STS-CP, Furthermore, STS-CP makes progress more quickly than CP-ARLS-LENV See Appendix 3.4 for similar plots for other datasets.
Discussion and Future Work
Our method for exact Khatri-Rai leverage score sampling enjoys strong theoretical guarantees and practical performance benefits. Especially for massive tensors such as Amazon and Reddit, our randomized algorithm's 'i guarantees translate faster progress 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 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.
10

---

Acknowledgements, Funding, and Disclaimers
We thank the referees for valuable feedback which helped improve the paper.
Bharadwaj was supported by the U.S. Department of Energy, Office Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC0022158, O. A. Malik and A. Buluç were supported by the Office of Science of the DOE under Award Number by the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation program through grant agreement 810367. R. Murray was supported by Laboratory Directed Research and Development (LDRD) funding from Berkeley Lab, provided by the Director, Office of Science, of the U.S. DOE under Contract No. DE-ACO2-05CH 1231. R. Murray was also funded by an NSF Collaborative Research Framework under NSF Grant Nos. 2004235 and 2004763. This research used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility, under Contract No. DE-ACO2-05CHI 1231 using NERSC award ASCR-ERCAP0024170.
1231. L. Grigori was supported
This report was prepared as an account of work sponsored by an agency of the United States Govermment Neither the United States Government any agency thereof, nor any of their employees, makes any warranty, express or implied, or assumes any legal liability or responsibility for the accuracy, completeness, usefulness of any information, apparatus, product, process disclosed, or represents that its use would not infringe privately owned rights. Reference herein to any specific commercial product, process, service by trade name, trademark, manufacturer, or otherwise does not necessarily constitute or imply its endorsement, recommendation, or favoring by the United States Government any agency thereof. The views and opinions of authors expressed herein do not necessarily state or reflect those of the United States Government or any agency thereof.
References
[1] Thomas 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.
[2] Nir Ailon and Bernard Chazelle. The fast Johnson-LLindenstruus transform and approximate nearest neighbors. SIAM Journal on computing. 39(1)300--322, 2009.
[3] Brett W. Bader and Tamara G. Kolda. Efficient matlab computations with sparse and factored tensors. SIAM Journal on Scientific Computing, 30(1):00--231, 2008. doi: 10. 1,.117706666488)
[4] Casey Battaglino, Grey Ballard, and Tamara G. Kolda. A practical randomized CP tensor decomposition. SIAM Journal on Matrix Analysis and Applications, 39(2):77-9911, 2018. doi: 10.1137/17M11 11 12303.
[5] Ke Chen, Qin Li, Kit Newton, and Stephen Wright. Structured random sketching for PDE inverse problems. SIAM Journal on Matrix Analysis and Applications, 41(4):1742-1770, 2020. doi: 10.1137/20M13104497
[6] Dehua Cheng, Richard Peng, Yan Liu, and Ioakeim Perros. SPALS: Fast alternating least squares via implicit leverage scores sampling. In Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
[7] Huaian Diao, Rajesh Jayaram, Zhao Song, Wen Sun, and David Woodruff. Optimal sketching for Kronecker product regression and low rank approximation. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[8] Petros Drineas, Malik Magdon-Ismail, Michael W. Mahoney, and David P. Woodruff. Fast approximation of matrix coherence and statistical leverage. J. Mach. Learn. Res., 13(1): dec 2012. ISSN 1532-4435.
[9] Matthew Fahrbach, Gang Fu, and Mehrdad Ghadiri. Subquadratic Kronecker regression with applications to tensor decomposition. In Advances Neural Information Processing Systems, volume 35, pages 18776-28788 Curran Associates, Inc., 2022.
11

---

[10] Azzam Haidar, Tingxing Dong, Stanimire Tomov, Piotr Luszezek, and Jack Dongarra. Frame- work for batched and GPU-resident factorization algorithms block Householder transforma- tions. In ISC High Performance, Frankfurt, Germany, 07-2015 2015. Springer, Springer.
[11] N. Halko, P. G. Martinsson, and J. A. Tropp. Finding structure with randommess: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM Review, 53(2):217-288, 2011. doi: 10. 10.1137/000771806
[12] 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. 10.1099//maia//aaa0223
[13] Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Review, 51(3):455-500 August 2009, ISSN 9036-1445 doi: 10. 1137/07070111X for Industrial and Applied Mathematics.
Publisher: Society
[14] Jean Kossaifi, Yannis Panagakis, Anima Anandkumar, and Maja Pantic. Tensorly: Tensor learning in python. Journal of Machine Learning Research (JMLR), 20(26), 2019.
[15] 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-1511 August 2022. doi: 10.113721114411544
[16] 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.
[17] Linjian Ma and Edgar Solomonik. Cost-efficient gaussian tensor network embeddings for tensor-structured inputs. In Advances Neural Information Processing Systems, volume 35, pages 38980-38993. Curran Associates, Inc., 2022.
[18] 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.
[19] Osman Asif Malik, Vivek Bharadwaj. and Riley Murray. Sampling- based decomposition algorithms for arbitrary tensor networks, October 2022. arXiv:2210.03828 [cs, math].
[20] Riley Murray, James Demmel, Michael W. Mahoney, N. Benjamin Erichson. Maksim Mel- nichenko, Osman Asif Malik, Laura Grigori, Piotr Luszezek, Michal Dereziáski, Miles Lopes, Tianyu Liang, Hengrui Luo, and Jack Dongarra. Randomized numerical linear algebra: perspective on the field with an eye to software, 2023.
[21] Aravind Reddy, Zhao Song, and Lichen Zhang. Dynamic tensor product regression. In Advances in Neural Information Processing Systems, volume 35, pages 4791-4801 Curran Associates, Inc., 2022.
[22] 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.
[23] N.D. Sidiropoulos and R.S. Budampati. Khatri-Rao space-time codes. IEEE Transactions Signal Processing, 50 50(10))2366-2407 2002. doi: 10.11099SS22002.83333
[24] Shaden Smith and George Karypis. SPLATT: The Surprisingly ParalleL spArse Tensor Toolkit.
https://gi tht.://gttttuu.... hwt..///aaaaaeeeennt//ssss////// splatt 2016.
[25] Shaden Smith, Jee W. Choi, Jiajia Li, Richard Vudue, Jongsoo Park, Xing Liu, and George Karypis. FROSTT: The formidable repository of open sparse tensors and tools, 2017. URI.
http://frostt.io/
12

---

[26] Linghao Song, Yuze Chi, Atefeh Sohrabizadeh, Young- kyu Choi, Jason Lau, and Jason Cong. Sextans: A streaming accelerator for general-purpose sparse-matris catuse-matrix vense-matric multiplication. In Proceedings of the 2022 ACM/SIGDA International Symposium on Field-Proge Proild-Pogrumaalle Gate Arrays, FPGA '22, page 63-77, New York, NY, USA, 2022. Association for Computing Machinery. ISBN 7781450391499 doi: 10. 10.1145/334042223552555
[27] Yining Wang, Hsiao- Yu Tung. Alexander J Smola, and Anima Anandkumar. Fast and guaranteed tensor dccomposition via sketching. Advances in neural information processing systems, 28, 2015.
[28] Sasindu Wijeratne, Ta- Yang Wang, Rajgopal Kannan, and Viktor Prasanna. Accelerating sparse MTTKRP for tensor decomposition on FPGA. In Proceedings of the 2023 ACM/SIGDA International Symposium Field Programmable Gate Arrays, FPGA '23, page 259-269, New York, NY. USA, 2023, Association for Computing Machinery. ISBN 9781450394178. doi: 10.1145/553322..5777779
[29] David Woodruff and Amir Zandieh. Leverage score sampling for tensor product matrices in input sparsity time. In Proceedings of the 39th International Conference on Machine Learning, volume 162 Proceedings of Machine Learning Research, pages 23933-23964. PMIR, 17-23 Jul 2022.
[30] David P Woodruff et al. Sketching tool for numerical linear algebra. Foundations and Trends@ in Theoretical Computer Science, 10(1-2):1-157, 2014.
[31] Yao Yu, Athina P. Petropulu, and H. Vincent Poor. MIMO radar using compressive sampling. IEEE Journal of Selected Topics in Signal Processing, 4(1): 14))144616,,, February 2010. doi: 10.1109/JSTSP22009 00000900SSSSP200022837....
13

---

A Appendix
A.1 Details about T
CP-ALS [3 is the standard, non-randomized alternating least squares method given by Algorithm 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 [1]] that samples rows from the Khatri-Rao product according to 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 [8] proposed the CP-ALS-ES algorithm (not listed in the table), which is superseded by the TNS- -PP algorithm [99]. We report the complexity from Table I 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 for Algorithm in their work.
Table T 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²) initialization time. CP-ARLS-LEV uses the initialization phase to compute the initial leverage scores of all factor matrices. TNS-CP uses initialization step to compute and cache Gram matrices of factors Uj. STS-CP must build the efficient sampling data structure described in Theorem I.] The algorithm from Ma and Solomonik requires an initialization cost of O m), where m is sketch size parameter that O (N R/s²) to achieve the (z, 8) accuracy guarantee for each least squares solve.
A.2 Definitions of Matrix Products
Table defines the standard matrix product :, Hadamard product @, Kronecker product @, and Khatri Rao product O, as well as the dimensions of their operands.
Table Matrix product definitions.
OPERATION
SIZE OF A
SIZE OF B
SIZE OF C'
DEFINITION
C A. B
(m, k)
(k,n) m)
(m, n) (m,n)
Cii,j = Σ
A a
B [a,j]
C &
(m,
(m, m)
BijBu B(i
C &
(ma, m)) (m,nn)
(ma n2)
(mam), n2n1)
C (ia..)
Ai
C A B
(12 nt)
(Mam) n)
C (11,, AA B(iz.]
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 + Eu-subspace embedding for A U1 @ @ UN is given in Appendix B of their work (proof of Theorem 2.7). Adapted to our notation, their runtime is
o (1og* Roo N Σ
N
(U,) + 87 log Rlog N
where = ER A a regularization parameter. For comparison, our runtime for constant failure probability 6 is
X \1 AR are the eigenvalues of the Gram matrix G of matrix A, and À 2 0
R3
N log R
R T=1
nnz (U,)
log II
4,
Woodruff and Zandieh's 's method provides significant advantage for large column count R or high regularization parameter X. As result, well-suited to the problem regularized low-rank approximation when the column count R is given by the number of data points dataset. On the
14

---

other hand, the algorithm has poor dependence on the matrix count N and error parameter For typically no larger than a few hundred, while high accuracy (e 10-3) is required for certain tensors to achieve fit competitive with non-randomized methods (see section Figures
tensor decomposition,
5 When À small, we have 82 R. Here, Woodruff and Zandieh's s runtime has an O dependence similar to ours. When R < log" A R log N, our sampler has faster construction time as well.
4 and
Finally, note that our sampling data structure can be constructed using highly cache-efficient, parallelizable symmetric rank-At updates (BLAS3 operation dSYRK). As result, the quadratic dependence 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 appeared in a modified form Lemma 10 in the work by Malik [8]. This original version used the definition G> = + @ + x=++ G3 in place of G>x defined in Equation ( where was a sketched approximation of G+. Woodruff and Zandieh [29] exhibit version of the theorem with similar modifications. We prove the version stated in our work below.
Proof of Theorem 3.7 We rely the assumption that the Khatri-Rao product A is nonzero matrix (but it may be rank-deficient). We begin by simplifying the expression for the leverage score of a row of A corresponding multi-index (x.). Beginning with Equation (2). we derive
bi = A [(i,, iN). G+ [[(1,...iNN).
(A ((11...,, iN), ((i,,...,, in), G+)
*
Ua Nas:l]
*
Ua [a G+)
Ua [ar :] Ua (ar
a=1
k-
*
Ua :
Va Na. :) Un :] Un lix. : c
*
Ua :a
Ua i
a=l
a=k+1
k- 1
N
*
Ua [aa :] Ua [a. :) Un :] Un [ix.
G
G
c
U :a Ua [aa
a=1
a=k+1
(13)
We proceed to the main proof of the theorem. To compute p(3, over random variables $k+1..... Recalling the definition of heck from Equation 33. we have
8cck sck), marginalize
( ($ck
N
p(Sx = Skk 8cck = Sek) x Σ
((Ect= Sek) (Sk -k) A
A u=#+1
(Bu = in)
(14)
Σ
The first line above follows by marginalizing over kk+l,..., SN. The second line follows because the joint random variable ($1,..,, &x) follows the distribution of statistical leverage scores on the rows of
15

---

to get A. We now plug in Equation 13
Σ Nkt1onn
N
Σ
Us [sa: i] Ua [sar Us [sk,:]
Us [skri], G
*
Us Bas :
Ua Bayi])
a=1
1
Σ iRtinndy
(Mecherk Un [s, Un [swr)) G @
*
Un [a
Un [i,
a=k+1
-hethh Uk [sks Us [,t,
Ua Ra. :1 Un [R,
6=
-hacherk Up [s Us [sk, ;
G
G.)
a a=k+1
-heche Us [] Us [skr Gsa))
(15)
We now compute the normalization constant for the distribution by summing the last line of Equation (15) over all possible values for $y:
C
Σ Ceec Un [s UR (s
4x
(16)
(heahekk
Σ
Un [skyi]
Un ( Gzk)
8k-1
-hechh Gky Gzk)-
For k 1, we have h2ck = [1,.... 1)_ so C = (Gk, G2R). Then C is the sum of all leverage scores, which is known be the rank of A [30] Since A assumed nonzero, C 0. For k > assume that the conditioning event &C* = 82* occurs with nonzero probability. This a reasonable assumption, since our sampling algorithm will never select prior values $1,..., Sk-1 that have probability of occurrence. Let Ć' be the normalization constant for the conditional distribution on $4-1- Then we have
p( 8k-1 8cc-1 sck-1)
E- (hckk-hh Un-1 [s Un-1 (1 G>k-1)
C- T
<hi G
(17)
C- G + Gzk)
C-I(leclize Gk. Gzk) = =
Since δ > 0, 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 Ty,r denotes the collection of nodes in full, complete binary tree with [1/F] leaves. Each leaf u E T7.F holds a segment S(u) = (S) S (1..7), with [S(v) < and S()) Ø for distinct leaves u, For each internal node u. S(v) S(L(e))US(R()) where L(v) and R(v) denote the left and right children node v. The root node r satisfies S(r) (1..))
Construction:
Algorithm 3 gives the procedure to build the data structure. We initialize segment tree TI,F and compute G for all leaf nodes e E Tr.F as sum of outer products of rows from U (lines 1-3). Starting at the level above the leaves, then compute G for each internal node as
16

---

the sum of GT.() and GR(O), the partial Gram matrices of its two children. Runtime O(IR²) is required to compute I outer products across all iterations of the loop on line 3. Our segment tree has [T/F] - internal nodes, and the addition in line 6 contributes runtime (R2) for each internal node. This adds complexity ((R2((T// - 1)) S O(TR"), for total construction time O(TR²).
To analyze the space complexity, observe that we store matrix G" E R at all 2[1/ - nodes of the segment tree, for asymptotic space usage O(T1/F D(TI/FTRR²) We can cut the space usage half by only storing G" when is either the root or a left child in our tree, since the sampling procedure in Algorithm 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 Finally, in the special case that I < F, the segment tree has depth and the initial binary search can be eliminated entirely. As result, the data structure has O(1) space overhead, since we can avoid storing any partial Gram matrices G". This proves the complexity claims in point of Lemma 3.2
Algorithm 3 BuildSampler(U E IR F, F.Y)
1: Build tree Tı with depth d = [log [/
2: for E leaves(T),r) do
3: G := (v) U [ U [
4: for = d d 2..0 do
5: for v E (e,e, do
6: G = (v))) + (())
Sampling: Algorithm gives the procedure draw a sample from our proposed data structure. It is easy to verify that the normalization constant C for Uu, is (hh Groo(TT.F)) Y), since (r) = U. Lines 8 and 9 initialize pair of templated procedures ni 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 Uu, [So(v) S1 S1())] for the segment (S.(v)..Si(t)) maintained by a leaf node. To see this last fact, observe for E [1] that
q(v) ji- - -So(v)]
-C- [ (hh + &.)
(18)
= C- (hh U [ V
9h.u.y [1]
The loop on line 12 performs the binary search using the two templated procedures. Line 18 uses the procedure ğ to scan through at most F bin endpoints after the binary search finishes early.
The depth of segment tree F is log] [ F1. As a result, the runtime of the sampling procedure is dominated by log [Tg[ evaluations m and a single evaluation of ÿ during the binary search. Each execution of procedure v. requires time O O(R2), relying on the partial Gram matrices G computed during the construction phase. When Y is general p.s.d. matrix, the runtime of q O(F This complexity is dominated by the matrix multiplication W. (hh" T @ Y) on line 5. In this case, the runtime of the "RowSampler" procedure to draw one sample is OR² log [/F] + F FR2), satisfying the complexity claims in point 2 of the lemma.
Now supppose Y is a rank-1 matrix with Y = =w for some vector u. We have hh (h s (h @ u)T. This gives
+ + =
(h, C, v) = diag(W (hh" T @ uu W) = (W. (h @ x))²
where the square is elementwise. The runtime of the procedure q is now dominated by a matrix-vecto multiplication that costs time O(FR). In this case, we have per-sample complexity Ó(R² log [/F
FR), matching the complexity claim in point of the lemma.
Recall that the inversion sampling procedure partitions the interval [0. 1] into I bins, the i-th bin having width yu,u. [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 U.Y, At the end of the loop, c is a leaf node that
Correctness:
17

---

Algorithm 4 Row Sampling Procedure
Require: Matrices U, Y saved from construction, partial Gram matrices {G" Tr,F).
1: procedure my(h, C, v)
2: return C- G", Y)
3: procedure 4y (h, v)
4: W := U [
5: return C- diag( (W- (hh" -Y-WT)
6: procedure RowSamplc(h)
C== (hh", (T.)(Tr.)) Y)
7:
8:
mn-)) == mp(h, C,.)
9:
q(-) == go (h, C,-)
c : roo
low 0.0, high 1.0 Sample d ~ (0.0,
11:
while e & leaves(Ti.F) do cutoff == low + m(L()) if cutoff 2 d then c: T((c), high == cutoff
13:
15:
else
17:
c: [[(c),((ooo low cutoff
return So(v) + arg min:20 (low - Σj= 7(c) <
18:
maintains collection S(c) of bins, one of which contains the random draw d. Since the procedure ạ correctly returns probabilities fu, [i] for E S(c) for leaf node e, (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 T and in the context of the sampling procedure outlined We verify the complexity claims first and then prove correctness.
in section 3.2
For each matrix Uy. Algorithm [ builds an efficient row sampling data structure Zj as specified by Lemma 3.2 We let the p.s.d. matrix Y that parameterizes each sampler be matrix of ones, and we set F R. From Lemma 3.2 the time to construct sampler Z, R() The space used by sampler Zj is O(TL R?a) O(1 since F' R. In case I₁ R. we use the special case described in Appendix A.5 to get a space overhead O(1), avoiding a term O(R²) in the space complexity.
Construction and Update:
Summing the time and space complexities over all j proves part 1 of the theorem. To update the data structure if matrix U, changes, we only need to rebuild sampler Z, for cost of R(I, The construction phase also computes and stores the Gram matrix G, for each matrix Uj. We defer the update procedure in case single entry of matrix U, changes Appendix A.7
Sampling: For all indices (except possibly 1), lines 1-5 from Algorithm compute G and its eigendecomposition. Only single pass over the Gram matrices G& is needed, so these steps cost ( for each index k. Line builds an efficient row sampler Eş for the matrix of scaled eigenvectors VAx- Vy. For sampler k. we Y Ge with cutoff parameter F = From Lemma 3.2 the construction cost is O(R²) for each index k. and the space required by each sampler is (
Summing these quantities over all # j gives asymptotic runtime O R&9) for lines 2-5.
The loop spanning lines 6-12 draws row indices from the Khatri-Rao product U+ For each sample, we maintain "history vector" h to write the variables h from Equation (3) For each index A * j. we draw random variable åg using the row sampler Es. This random draw indexes scaled eigenvector of G We then use the history vector h multiplied by the eigenvector sample row index &g using data structure Z₁. The history vector h is updated, and proceed to draw the next index fx.
18

---

As written, lines 2-5 also inuur scratch space usage O(NRR). The scratch space can be reduced to O(R)) by exchanging the order of loops on line and line 8 and allocating J separate history vectors hi once for each draw. Under this reordering, we perform all J draws for each variable ûn and &z before moving to 0x+1 and 1 In this case, only a single data structure Es is required 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 implemented the memory-saving approach in our code.
From Lemma 3.2 lines 9 and 10 cost O log R) and O (R² log [+/R]), respectively. Line 11 costs O(R) R.) and contributes lower-order term. Summing over all & + j. the runtime to draw single sample
Σ (R² log[I</R] kØj
R² R2 log R) = =0
Σ R log max (Ik, R)
Adding the runtime for all J samples to the runtime of the loop spanning lines 2-6 gives runtime O (N E R log max (1x+ Rt) and the complexity claims have been proven.
We show correctness for the case where j and we sample from the Khatri-Rao product of all matrices Us. since the proof for any other value of j requires a simple reindexing of matrices. To show that our sampler is correct, it enough to prove the condition that for 1 S k S N,
Correctness:
p(4+ = 1x heck) = Uhen.Us...s. [(x]
(19)
since, by Theorem 3.1 p(52 = 8% B = )+)x = Qhex.Us. [sw]. This would imply that the joint
random variable (4,, (x) has the same probability distribution as (81, AN), which by definition
follows the leverage score distribution on Un @ .... UN. To prove the condition in Equation (19)
we apply Equations (9) and (T]) derived earlier:
plle = la heR)
Σ
pffx th 0b uk. h<k)p(u up her)
(Bayes' Rule)
n
[Kx us.
Σ w [ux]
(Equations 9 and 11 in reverse)
[ [ ux
04.
Uhet,
[tx].
(20)
A.7 Efficient Single-Element Updates
Applications such as CP decomposition typically change all entries of single matrix Uj between iterations, incurring an update cost O(I)R²) for our data structure from Theorem 1.1 In case only single element of U, changes, our data structure can be updated in time O (R log 15).
Proof. Algorithm 5 gives the procedure when the update U, [r, c] = â is performed. The matrices Go refer to the partial Gram matrices maintained by each node v of the segment trees in our data structure, and the matrix U, refers the matrix U, before the update operation.
Let T be the segment tree corresponding to matrix U, in the data structure, and let u Ti,,R be the leaf whose segment contains Lines 3-5 of the algorithm update the row and column indexed by e in the partial Gram matrix held by the leaf node.
The only other nodes requiring an update are ancestors of each holding partial Gram matrix that is the sum 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 change between the old value of G and its updated version. Thus, the runtime of this procedure O R log / [/) from multiplying the cost to update single node by the depth of the tree.
D
19

---

Algorithm 5 UpdaateSampler(,rsccc
1: Let u = U, [r,
2: Locate such that r E S(v)
3: Update G" [c, += (ú u U) |,
4: Update G [.c]] += (Ú u U) [,,
5: Update G [c, += (û u)
6: while o roou(T1,..) do
7: Uproy == v, ;= A(v)
8: Update G = (riblingg(perr) + G
A.8 Extension to Sparse Input Matrices
Our data structure is designed to sample from Khatri-Rart products U1 .... @ UN where the input matrices U1, .... Un 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.I.
Corollary A.1 (Sparse Input Modification). When input matrices U1.... UN are sparse, point 1 of Theorem 1.1 can be modified so that the proposed data structure has O (
construction time and O (Σ nnz(U,)) storage space. The sampling time and scratch space usage in point 2 of Theorem does not change. The single-element update time in point I 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, [I] in relation to the data structure in Lemma 3.2 These choices for F and are used the construction phase given by Algorithm | The proof in Appendix A.5 constrains each leaf of segment tree T,, hold contiguous segment
S(v) S [Z] of cardmallty at most F. Instead, choose each segment S(v) = {Su(v)...S(o)) so that U (S)(()) (t),.: has at most R² nonzeros, and the leaf count of the tree is at most [nn (U) +1 for input matrix U E I Assuming the nonzeros of U are sorted in row-major order, we can construct such a partition of [T] into segments in time (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 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
constructs the partial Gram matrix for each leaf node in the segment tree. Each nonzero in the segment U (v) S₁ [o).:] contributes time (RR) during line 3 of Algorithm to update a single row and column of G 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") R2 internal nodes and the addition on line 6 costs O(R²) per node, the remaining steps of the construction procedure cost (Unz(U)), lower-order term. The construction time is therefore ((nzz(U) ( R)..
Since we store a single partial Gram matrix of size R² at each of O ([nnz(U)//21) internal nodes, the space complexity of our modified data structure is (mnz(('))
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 (U) + 1 is upper-bounder by [I/ R] + 1, the runtime of the update procedure in Algorithm stays the same.
Updated Sampling Complexity of Lemma 3.2 R, = [1]: The procedure "RowSample" in Algorithm now conducts traversal of tree of depth O ((mz(U))R))) As result,
20

---

can still upper-bound the number of calls to procedure m as [I/F]. The runtime of procedure
in unchanged. The runtime of procedure for leaf node is dominated by the matrix-vector multiplication U [So(e) (c), h. This runtime is O (nnz [So(c) (c),,:l) < O (R2)
Putting these facts together, the sampling complexity of the data structure Lemma 3.2 does not
change under our proposed modifications for F R. YR,Y= [1]-
Updated Construction Complexity of Theorem 1.1
Algorithm now requires
O OR R (RE)" nnz(U,)) construction time and O (Σ (n)) storage space, summing
the costs for the updated structure from Lemma 3.2 over all matrices U1.... UN. 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 R4 x as weighted sum of generalized outer products. Formally, let U1, UN with U, E R4ix& be factor matrices with each column having unit norm, and let a E R be nonnegative coefficient vector. We call R the rank of the decomposition. The tensor T that the decomposition represents is given elementwise by
7 Tir iN] = (oT,Us (e.. UN :i.
Σ a u ..N...N..)
which generalized inner product between o T and rows U, [ for 1 J S N. Given an input tensor T and target rank R, the goal of approximate CP decomposition is to find rank-R representation / that minimizes the Frobenius norm T T-
Definition of Matricization. The matricization mat(T, .) flattens tensor T E R into a matrix and isolates mode j along the row axis of the output. The output of matricization has dimensions x II,y In. We take the formal definition below from survey by Kolda and Bader The tensor entry T [11,-..,8 in] is equal to the matricization entry mat(T,)) [iv, u). where
13
II Im.
Details about Alternating Least Squares. Let U1,..., UN be factor matrices of a low-rank CP decomposition, U* IR We use Uęi denote Ok= Ug. Note the inversion order here to match indexing in the definition of matricization above. Algorithm| gives the non-randomized alternating least squares algorithm CP-ALS that produces a decomposition of target rank R given input tensor T E | xg.xxx.xxyy general format. The random initialization on line I of the algorithm can be implemented by drawing each entry of the factor matrices U, according to a standard normal distribution, or via a randomized range finder [ The vector a stores the generalized singular values of the decomposition. At iteration j within a round, ALS holds all factor matrices except Uj constant and solves a linear-least squares problem on line for new value for Uy. In between least squares solves, the algorithm renormalizes the columns of each matrix U, to unit norm and stores their original norms in the vector a. Appendix A. 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 Uzi for optimization problem j within round of ALS
has dimensions 1I 4 x R. The observation matrix mat(T.j) has dimensions II 4x x 4z. To achieve error threshold + e with probability å on each solve, we draw J = Ô from both the design and observation matrices and solve the downsampled problem (Theorem 2.1
rows
21

---

Algorithm 6 CP-ALS(T.
1: Initialize U3 RR,xx randomly for S N.
2: Renormalize U [: / [. S S N, SisR.
3: Initialize d E RR to [1].
4: while not converged do
5: for j = 1...N do
6: U = arg min x [0y X mat 11)
7:
= J [:
<ii
8:
Renormalize U 1 - [ isish R.
9:
return [o; U1.. UN.
These rows are sampled according to the leverage score distribution on the rows of Uzj. for which we use the data structure in Theorem 1.1 After a one-time initialization cost 0(Σ, - ) before the ALS iteration begins, the complexity to draw J samples (assuming 2 R) is
o ( N + 1
R² log = NR
(
R Σ
R² log J
The cost to assemble the corresponding subset of the observation matrix is (JII) O(RI The cost solve the downsampled least squares problem is O(JR²) ((I,R\(EE)), which dominates the cost of forming the subset of the observation matrix. Finally, we require additional time OiT, R update the sampling data structure (Theorem 1.1 part 1). Adding these terms together and summing over v S N gives
N
+
Σ kデj
Rd log Lx
Eδ
j=1
(21)
N
=
[I (N -1 log 1,]
Rounding N - to and multiplying by the number of iterations gives the desired complexity. When I, R for any j. the complexity changes in Theorem I.] propagate to the equation above. The col- umn renormalization line 8 of the CP-ALS algorithm contributes additional time O (EA- per round, a lower-order term.
A.10 Experimental Platform and Sampler Parallelism
We provide two implementations of our sampler. The first 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 efficient C++ implementation using 128 OpenMP threads. We link our code against Intel Thread Building blocks to call 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 Khatri-Rao product, thereby enjoying better sample efficiency than alternative approaches such as CP-ARLS-LEV 15). The cost draw each sample, however, is O log H), where is the number of rows the Khatri-Rao product. Methods such row-norm-squarde sampling or CP-ARLS-LEV can draw each
22

---

sample in time (log H) after appropriate preprocessing. Therefore, efficient parallelization of our sampling procedure is required for competitive performance, and we present two strategies below.
1. 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 team, each of which makes calls to KRPSampleDraw asynchronously. This strategy works well on a CPU. but less attractive on a SIMT processor like a GPU where instruction streams cannot diverge without significant performance penalties.
2. 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 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 is followed by single call to the function g. Observe that procedure m in Algorithm 4 can be computed with 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 distinct set of samples. The MAGMA linear algebra library provides batched version of genv [TO], while batched dot product can be implemented with an ad hoc kernel. MAGMA also offers a batched version of the symmetric rank-k update routine syrk, which is helpful to parallelize row sampler construction (Al (Algorithm B). 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 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: Sparse Tensors from FROSTT collection.
TENSOR
DIMENSIONS
NONZEROS
PREP.
INIT.
UBER PICKUPS
183 x 24 x 1,140 1,717
3,309,490
NONE
IID
ENRON EMAILS
6,066 x 5.699 x 244.268 x 1,176
54,202.099
LOG
RRF
NELL-2
12,092 x 9,184 x 28,818 4,821,207 X 1,774,269 x 1,805,187
76,879.419
LOG
IID
AMAZON REVIEWS
1,741,809,018
NONE IID
REDDIT-2015
8,211,298 x 176,962 x 8,116,559
4,687,474,081
LOG
IID
Table 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 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..d. entries from the standard normal distribution. As suggested by Larsen and Kolda 115, the Enron tensor's factors were initialized with a randomized range finder [I]] The range finder algorithm initializes each factor matrix U, as ma)(T,j))S, a sketch applied the mode-j matricization of T with S T Ix 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.) were nonzero. Thus, we computed at most mnz (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 I epoch), and we used an early
23

---

stopping condition to terminate runs before the maximum round count. The algorithm was terminated epoch if the maximum fit in the last 3 epochs did not exceed the maximum fit from epoch through epoch T - 3 by tolerance 10-1
Hybrid CP-ARLS-LEVV deterministically includes from the Khatri-Rao product whose probabil- ities exceed 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 7 = 1/J [1].
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
non-randomized ALS implementation. As a result, Figure 3 verifies that randomized algorithm STS-CP produces tensor decompositions with accuracy comparable to those by highly-optimizedd state-of-the-ar CP decomposition software. We leave a implementation of our randomized algorithms to future work.
A.12 Efficient Computation of Sketch Distortion
The definition of a in this section is different from its definition in the rest of this work. The condition number k of a matrix M is defined as
0 max (M)
K (M)
Fmin (M)
where Omin (M) and max (M) denote the minimum and maximum nonzero singular values of M. Let A be Khatri-Rao product of N matrices U1,..., UN with II, I rows, R columns, and rank r S R. Let A = QEVT be its reduced singular value decomposition with Q E RII, 1şXr E E R" and V E RT*R, R, Finally, let S R R 4, be a leverage score sampling matrix for A. Our goal is to compute x(SQ) without fully materializing either A or its QR decomposition. We derive
k(SQ) =(SQEV x(SAVE-))
(22)
The matrix SA E JJJRR is efficiently computable using leverage sampling data structure. We require time OJJR²) to multiply it by VE- and compute the singular value decomposition of the product to get the condition number. Next observe that A VV²V,, so we can recover V and E-1 by eigendecomposition of A " gdx ft in time O(RS). Finally, recall the formula
A A
A
UJU, U
U
used at the beginning of Section 3 that enables computation of A A in time O 1, without materializing the full Khatri-Rao product. Excluding the time to form SA (which is given by Theorem 1.1 x(SQ) is computable in time
N o JR + R
Σ R
3=1
Plugging R(SQ) into Equation 112) gives an efficient method compute the distortion.
A.13
Supplementary Results
A.13.1
Comparison Against Standard CP Decomposition Packages
Table f 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 version of CR-SS-LLE in Tensor Toolbox.
24

---

Table 4: Average time (seconds) per ALS round for our method vs. standard CP decomposition packages. OOM indicates an out-of-memory error. All experiments were conducted on a single LBNL Perlmutter CPU node. Randomized algorithms were benchmarked with 2 samples per least-squares solve.
METHOD
UBER
ENRON
NELL-2
AMAZON
REDDIT
TENSORLY, SPARSE BACKEND
64.2
0OM
759.6
0OM
0OM
MATLAB TTOOLBOX STANDARD
11.6
294.4
177.4
>3600
OOM
MATLAB TTOOLBOX CP-ARLS-LEV
0.5
1.4
34.2
OOM
STS-CP (ours)
0.2
0.5
0.6
3.4
26.0
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-Rat 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 fair comparison with CP- ARLS-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 problem that our sampler works as expected. For the Khatri-Rao product of three matrices A = U, @ U2 @ U3. it plots the true distribution of leverage scores against a normalized histogram of 50,000 draws from the data structure in Theorem [.]] We choose U1. U2 U3 E R 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 distribution.
0.025
Tran Lomsaye Soom Butritation Thetograre of frowe thiss Our Sunpuer
0.615
0.000
0.000
0,000
Towe Tndror frres ESD
Figure 7: Comparison of true leverage score distribution with histogram of 50,000 samples drawn from U1 @ U2 Uz.
A.13.3 Fits Achieved for J = 2
Table 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 Function of Time
Figures Sa and 8b shows the fit as a function of time for the Amazon Reviews and NELL2 tensors. The hybrid version of CR-ARL--EVV was used for comparison in both experiments. As in section 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-ARIS-LEV do not achieve the same final fit STS-CP. CP-ARI---LEV finally achieves comparable fit to STS-CP when the former uses 196, 608 samples, compared to 65, 536 for our method.
25

---

Table Fits Achieved by Randomized Algorithms for Sparse Tensor Decomposition, = 216, and randrandomieed ALS. The best result among randomized algorithms is bolded. "CP-ARLS-LEV-H" refers to the hybrid version CP-ARLS-LES and "Exact" refers to non-randomized ALS.
TENSOR
R
CP-ARLS-LEV
CP-ARLS-LEV-H
STS-CP (OURS)
EXACT
25
.187 2.30F-03
.189 1.52E-03
.190 + 1,41E-03
50
.211 + 1.72E-03
1.27-3
1.18-.1
.218 1.61E-03
UBER
75
.218 + 1.76E-03
2.05
2.24942
.232 9.29E-04
100
.217 1 3.15--03
1.7
237
2211 1.00E-03
125
.213 + 1.96E-03
.213 2.47
2 1.78E-03
.247 1.52E-0
25
.8811 + 1.02E-02
.0882 + 9.01E-03
.0955 + 1.19E-02
.0978 + 8,50E-03
50
.8833 1.72B-02
.0920 6.32E-33
.15 1.03E-02
112 1.51E-22
ENRON
75
.0899 6.100-03
.0855 6.39E-03
1.25=..5
1.77
100
.0809 1.26F-02
.0787 1.00F-02
5.90=-0
176
125
.0625 + 1.526-02
.0652 + 1.00F-02
.182 + 1.04E-02
25
.045
9.52E-04
.0467 4.61E-04
.0470 + 4.69E-04
.0478 + 7.20E-04
50
.0590 + 5.33E-04
4.34E
.0608 + 5.44E-04
.0618 + 4.211-04
NELL-2
75
00588 + 6.84n-04
0660
- 2.96E-04
.0708 + 3.11E-04
100
.0700 4.91
4.48
.0760 + 6.52E-04
.0779 5.09E-04
125
.0729 8.56E-04
7.22E--
.0814 + 5.03E-04
.0839 8.476-04
25
.338 1 6.63E-04
3.9.669968.
.30 6.61E-04
1.78
50
.359 1 1.09E-03
3.04
1.22E.0
.366 166:111 1.01E-03
AMAZON
75
.36 + 1.82E-03
1.74-.77
382
.382 5.90
100
.366 3.05E-03
271
:.32 6.776667
125
.358 = 6.51E-03
36424.228-03
.400 3.67-.67
.401 3.58E-04
25
.0581 1.02B-03
.0583 2.78
.0592 + 3.07E-04
.0596 + 4.27
.0766 1.03E-03
.0738 4.85E-03
.0774 7.88E-04
.0783 2.60E-04
REDDIT
75
.0845
1.64E-03
.8499 = 8.96E-04
.0909 + 5,49E-04 04
.0922 = 3.69E-04
100
.0904 + 1.35E-03
.0911 = 1.59E-03
.101 = 6.255-04
.103 + 7.14E-04
125
.0946 + 2.13R-00
.0945 + 3.17F-03
.109 + 7.71R-04
.111 + 7.98F-04
0.3600
0.3775
0.355
0.3725
0.0700
0.355
CPARSSS
0. 0.7555
CY-BLSS-EE 4.%% 153,000
LP-ARLS-LWN 1-131 272
CY-ALSS-EE 1313,0072
CPARLS-LEN 1-211-3
CY-ABLS LE% J-98.364 20.
SPARLSILP L.N.S-LV////ee -65.66
C2 ARLS LLX J-65,300
0.355
0 300
Time:
(a) Amazon
(b) NELL-2
Figure 8: Fit as a function time, 100.
A.13.5 Speedup 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 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 range of sample counts. We tested STS-CP on values of from the set (2 1 < 4] for all tensors. CP. ARLS-LEV required sample count that varied significantly between datasets to hit the required thresholds, and we report the sample counts that we tested in Table Because CP-ARIS-LEV has poorer sample complexity than STS-CP, we tested wider range of sample counts for the former algorithm.
26

---

Table 6: Tested Sample Counts for hybrid CP-ARLS-LEV
TENSOR
VALUES J TESTED
UBER
f2' x E {1..13}}
ENRON
12'x x (1..7) U (10, 12, 14, 16, 18, 20, 22, 26, 30, 46,50, 50,
NELL-2
(2
x E (1..7)}
AMAZON
(2 r
x E (1..7))
REDDIT
{2 {1.112
For each configuration of tensor, target rank R. sampling algorithm, and sample count J, we ran trials using the configuration and stopping criteria in Appendix [A.I]] The result of each trial 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 RT.R.A.J R+ + [0, 1]. The value fr. "T.... (t) is the average fit at time achieved by algorithm to decompose tensor T with target rank R using samples per least squares solve. Finally, let
miny argmin [FT.RRCP. C...L.E.-LL.EE (t) > PI
Speedup
R,M TR
miny argmin [fr R.STS--CP. s Pl
be the speedup of STS-CP with target rank R. We let the threshold P for each tensor T be fixed fraction of the fit achieved by non-randomized ALS (see Table 5
CP-ARLS-LEV (hybrid) to achieve threshold fit P on tensor T
Figure reports speedup STS-CP over hybrid CP-ARLS-LEV for 0.95 on tensors except Enron. For large tensors with over one billion nonzeros, report 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-CL TPP performs comparably to CP-ARLS-LEV at most ranks, with significant slowdown only rank 25 the smallest tensor Uber.
Attacoo
Ubee
NELL2
2.5-
1.50
2.0-
1.29
1.00
2.0
0.73
0.53
63
0.2
60
75
3.00
1.2
125
000
75
3500
125
(a) Large Tensors
(b) Small Tensors
Figure 9: Speedup of STS-CD over CP-ARLS-LEV hybrid (LK) to reach 95% of the fit achieved by non-randomized ALS. Large tensors have more than billion nonzero entries.
On the Enron tensor, hybrid CP-ARLL-EEV could not achieve the 95% accuracy threshold for any rank above 25 for the sample counts tested in Table STS-CP achieved the threshold accuracy for all ranks tested. Instead, Figure reports the speedup achieve 85% of the fit non non-rand randomized ALS on the Enron. Beyond rank 25, our method consistently exhibits more than 2x speedup reach the threshold.
Guide to Sampler Selection. Based on the performance comparisons in this section, we offer the following guide 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 offset by reductions in the least squares solve time. For smaller tensors, our sampler may still offer significant performance benefits (Enron). other cases (Uber, NELL-2), CP-ARL--LEEV exhibits better performance, but by small margins for rank beyond 50.
27

---

Squndor over
05-
Figure 10: Speedup of STS-CP over CP-ARLS-LEV hybrid (LK) reach 85% of the fit achieved by non-randomized ALS on the Enron Tensor.
STS-CP reduces the cost of each least squares solve through sample selection process that relies dense linear algebra primitives (see Algorithms [ [)) Because these operations can be expressed standard BLAS calls and can be carried out parallel (see Appendix AA.1 hypothesize that STS-CP is favorable when GPUs or other dense linear algebra accelerators are available.
Because target tensor is sparse, the least squares solve during each ALS iteration requires sparse matricized-tenn -enzor times Khatri-Rat product (spMTTKRP) operation. After sampling, this primitive reduced 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-LEE 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.
28


Reviews:
[{"review_id": "JOERmzJsmf", "reviewer": "Reviewer_u8zc", "summary": "This 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.", "questions": "- 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.", "limitations": "Yes, the paper discussed limitations in Section 5. Due to the technical nature, I do not see any potential societal impact of their work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1) 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.", "weaknesses": "W1) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and useful feedback, which have strengthened the quality of our draft.\n\nRegarding the third weakness [W3] about the novelty of our contribution: our approach offers a significant advance in the best theoretical complexity for sketched ALS CP Decomposition. By improving the per-leverage-score sample complexity from linear in $\\left| I_j \\right|$ (Malik 2022) to $\\log \\left| I_j \\right|$, we enable computational cost savings of **several thousand times** for approximations of **identical** quality. This is the critical improvement that allows exact leverage-score sketching for large sparse tensor decomposition, which we illustrate by the FLOP complexity calculation below.\n   \n   Suppose the algorithm in [Malik 2022] was applied to the Amazon Reviews tensor in our experiments with dimensions $4,821,207 \\times 1,774,269 \\times 1,805,187$ to produce a rank $R=25$ decomposition. After extra non-asymptotic improvements made in their paper, the exact floating point operation (FLOP) count to draw a sample from the KRP excluding the third mode is lower-bounded by $\\left| I_2 \\right| R^2$, or 1.12 gigaFLOPs *per row*. Including multiplicative constants but excluding the small $(25 \\times 25)$-sized eigendecompositions, **our** approach requires only 53.7 kiloFLOPs per row sample, a more than **20,000x reduction** that is due to the asymptotic improvement from $\\left| I_j \\right|$ to $\\log \\left| I_j \\right|$.\n\nWe plan to add the statements above to the paper to highlight our contributions over existing work.\n\nOther Responses:\n\n**Questions**\n1. [Tensor Descriptions] Agreed. We omitted these descriptions due to space constraints in the main body, but have added them in the extra space the revision affords.\n\n2. [Tucker Decomposition] The techniques developed here are less applicable to ALS Tucker decomposition due to the different properties of the Kronecker product on which it is based. The leverage scores on the rows of $A \\otimes B$ are products of the leverage scores of $A$ and $B$, a property that the Khatri-Rao product does not share. Drawing a sample from the Kronecker product, therefore, is simpler than sampling from a Khatri-Rao product, although the number of samples required to achieve least-squares solution guarantees is exponential in the tensor dimension for even state-of-the-art sketched ALS Tucker methods (see, e.g. https://arxiv.org/abs/2209.04876). As a result, the runtime to compute the Tucker decomposition may be significantly higher than CP decomposition for even 4-dimensional tensors.\n\n3. [Woodruff Zandieh Comparison] The algorithm developed by Woodruff and Zandieh is intricate (requiring three distinct sketching operators composed in non-trivial ways) with no publicly available code, to the best of our knowledge. Furthermore, their near input-sparsity time algorithm confers no benefit for tensor decomposition. The least-squares problems $\\min_X \\left| \\left| AX - B \\right| \\right|$ in CP decomposition have observation matrices $B$ with $\\left| I_j \\right|$ columns for each $1 \\leq j \\leq N$. Regardless of the efficiency of the sketching mechanism, the column count of matrix $B$ introduces an unavoidable cost $O(\\left| I_j \\right| R^2)$, destroying the benefit of the $\\tilde O(\\left| I_j \\right| R)$ input-sparsity time cost to form the sketching data structure for each factor matrix $U_j$. Furthermore, Woodruff and Zandieh do not provide a readily evident method to update their sketch when entries of each factor matrix change, although their approach could likely be adapted. Efficient factor matrix updates are essential for CP decomposition, and our algorithm provides methods to update each factor when even a single entry changes.\n\n\n**Weaknesses**\n\n1. [W1] Our introduction includes several motivating examples for linear systems with Khatri-Rao design matrices, including compressed sensing and signal processing besides tensor decomposition. Systems of this exact form also occur in PDE-inverse problems (see, for example, equations (2) and (4) from https://arxiv.org/pdf/1909.11290.pdf); like us, they also consider the case with low column count and hundreds of thousands of rows per matrix forming the Khatri-Rao product. We have added this reference to the paper.\n\n2. [W2] In the global response at the top of these reviews, the PDF attached contains figures  that illustrate Theorem 3.1 and our two-stage leverage-score sampling procedure. We have added these to our draft. We hope these aid the reader and are happy to engage in further discussion.\n\n3. [W3] The main baseline that we selected for comparison, CP-ARLS-LEV, is a state-of-the-art algorithm for sparse tensor decomposition with performance exceeding well-known software packages on the market. The table below compares the runtime per iteration of our algorithm with three other pieces of software for sparse tensor decomposition. Figure 4 (original document) and Table 4 demonstrate that our accuracy is comparable to non-sketched tensor decomposition, recovering 99.7% of the fit consistently for the Amazon tensor.\n\n| Method    | Uber  | Enron | NELL-2 | Amazon | Reddit|\n| - | - | - | - | - | - |\n| Tensorly, Sparse Backend | 64.2 | OOM  | 759.6 | OOM   | OOM    |\n| Matlab TToolbox Standard | 11.6  | 249.4  | 177.4   | >3600 | OOM |\n| Matlab TToolbox CP-ARLS-LEV | 0.5 | 1.4 | 1.9  | 34.2  | OOM|\n| **STS-CP (Ours)** | 0.2 | 0.5 | 0.6  | 3.4  | 26.0 | \n\nAll times reported are average seconds \nper single ALS iteration; randomized algorithms were benchmarked \nwith $2^{16}$ samples per LSTSQ solve. We used Python \nTensorly 0.81 and Matlab Tensor Toolbox version \n3.5 on an identical system\nconfiguration to other experiments. OOM means out-of-memory. **The timings demonstrate that our algorithm STS-CP can quickly decompose tensors requiring hundreds of gigabytes of disk space in a fraction of the time required by standard packages**. We have added these clarifications to a revised copy of our draft."}, {"review_id": "sVR2CyDCwE", "reviewer": "Reviewer_Njzt", "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.\n\n", "questions": "n/a", "limitations": "n/a", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\n2. 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.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your review of our paper. Per the instructions, the PDF attached to our global response at the top of these reviews has a three-part figure that illustrates the sampling process, including the matrices involved, the significant operations, and the two-stage sampling procedure. These diagrams have been added to our draft."}, {"review_id": "yVpFCAKBgn", "reviewer": "Reviewer_qocR", "summary": "The 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.", "questions": "I 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$", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The 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.", "weaknesses": "The 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!", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your comments, which have strengthened the quality of our draft. The global rebuttal response to all reviewers includes an extra PDF that contains updated figures (2, 3, 4, 5) with error bars / axes swaps mentioned in this review. It also includes a frequency spectrum analysis (figure 6 in the extra PDF) regarding the error oscillation in CP-ARLS-LEV (see below).\n\nResponses:\n\n**Questions:**\n1. [Notation] We chose the notation for column count [Larsen & Kolda 2022] and [Malik 2022], which use $R$ or $r$ as the column count (more standard in the tensor decomposition literature). We see your point with the other notation; happy to change $q_i$ to $p_i$. We were a bit worried that $s$ as the sample count could cause confusion with $\\hat s_1, ..., \\hat s_N$, denoting random variables for the multi-indices of a single sample, so we took $J$ as the sample count from [Malik 2022]. Open to changing this if necessary.\n\n2. [O tilde] Agreed, and pointed out by another reviewer. We have updated our draft.\n\n3. [Dense vs. Sparse Inputs] Agreed, thanks. We meant in lines 48-49 that the matrices $U_j$ in the Khatri-Rao product are typically dense, and will clarify this.\n\n4. [2% Slower Runtime] This is because the sample counts for our algorithm vs CP-ARLS-LEV were equal for this first experiment (line 53), and CP-ARLS-LEV has faster runtime per sample but a poorer accuracy for the same sample budget, both in theory and practice.\n\n   In hindsight, this was not the best way to highlight our contributions, since the runtime, accuracy, and sample couns must all be accounted for. Our\n   revised draft makes reference to the last section of the appendix and reads as follows: \"On the billion-scale Amazon and Reddit tensors,\n   our algorithm STS-CP can achieve 95% of the fit of exact ALS between 1.5x and 2.5x faster than the high-performance, state-of- \n   the-art competing CP-ARLS-LEV method. Our algorithm is significantly more sample-efficient; on the Enron tensor, only $2^{16}$ samples were required to achieve the 95% \n   accuracy threshold above a rank of 50, which could not be achieved by CP-ARLS-LEV with even 54 times as many samples.\n\n5. [Complexity Derivation] per the comment above, we will work in $\\tilde O$ notation. A version of the derivation below appears in appendix A.9, and we have added a reference to the main body of the text referencing it.\n\n   a. Each ALS iteration contains $N$ least-squares problems. The $j$'th problem, for $1 \\leq j \\leq N$, is $\\min_X \\left| \\left| U_{\\neq j} X - B \\right| \\right| $, where $U_{\\neq j}$ is the Khatri-Rao product of all matrices but $U_j$. The matrix $B$ has $\\left| I_j \\right|$ columns.\n\n   b. For each problem, we need $\\tilde O(R / (\\epsilon \\delta))$ samples from $A$ at a cost of $O(R^2 \\sum_{k \\neq j} \\log \\left| I_k \\right|)$ per sample, plus a one-time cost (not-per-sample) of $O(NR^3)$ for the eigendecompositions. To avoid the dependence on $j$, we will round up the cost per sample to $O(R^2 \\sum_{k=1}^N \\log \\left| I_k \\right|)$\n\n   The cost to compute a $QR$ decomposition on the downsampled design matrix is $\\tilde O(R^3 / (\\epsilon \\delta))$, a lower-order term. Because the observation matrix $B$ has $\\left| I_j \\right|$ columns, the cost of multiplying by $Q$ and back-substituting\n   with respect to $R$ is $\\tilde O(\\left| I_j \\right| R^2 / (\\epsilon \\delta))$. The total complexity per solve is:\n\n   $\\tilde O\\left(NR^3 + (R / (\\epsilon \\delta)) \\left(R^2 \\sum_{k=1}^N \\log \\left| I_k \\right| \\right) + \\left| I_j \\right| R^2 / (\\epsilon \\delta)\\right).$\n\n   For $\\epsilon, \\delta < 1$ and $\\left|I_k\\right| \\geq 2$ for all $k$, observe that the first term $NR^3$ is at most the second term in the expression above, and we can eliminate it; simplifying slightly gives\n\n   $\\tilde O\\left((1 / (\\epsilon \\delta)) \\left( \\sum_{k=1}^N R^3 \\log \\left| I_k \\right| \\right) + \\left| I_j \\right| R^2 / (\\epsilon \\delta)\\right)$\n\n   c. Finally, we sum the expression above over $1 \\leq j \\leq N$ to get:\n\n   $\\tilde O\\left((N / (\\epsilon \\delta)) \\left( \\sum_{k=1}^N R^3 \\log \\left| I_k \\right| \\right) + \\sum_{j=1}^N \\left| I_j \\right| R^2 / (\\epsilon \\delta)\\right)$\n   \n   Combining the index variables $k$ and $j$ over independent sums into a single summation over $j$ gives:\n\n   $\\tilde O\\left((1 / (\\epsilon \\delta)) \\sum_{j=1}^N \\left( N R^3 \\log \\left| I_j \\right| + \\left| I_j \\right| R^2 \\right) \\right)$\n   \n   and setting all $\\left| I_j \\right|$ equal to $I$ gives the Table 1, last row.\n\n6-9. Thanks, fixed.\n\n10. The variables $Z_j$ are the segment trees (with cached matrices) for the second-stage eigensampling from each factor matrix, and the variables $E_k$ are the segment trees for the first stage involving the small $R \\times R$ gram matrices. We clarified this and also revised the last paragraph of Section 3.2. See the global rebuttal for the revision text.\n\n11, 12. Agreed, see PDF. We observe, in the updated Figure 3, that the median error of our method (orange bars) across several randomly generated input matrices actually decreases as the tensor dimension $N$ increases. On the other hand, there are more outliers (defined in the caption) that drive up the mean at increased dimension. We carefully checked that this behavior does not arise from experimental error. Both algorithms run on identical matrices regenerated once per new trial, with identical solve procedures after the sample indices are identified.\n\n13. [Oscillating error pattern on Amazon] Our algorithm does exhibit some oscillation on Amazon, but we find no clearly-defined period. The period for CP-ARLS-LEV **exactly matches the dimension of the tensor**. This is an artifact that the worst-case theoretical guarantees for CP-ARLS-LEV do not capture. Figure 6 in the PDF attached is a frequency spectrum of the errors. Observe that CP-ARLS-LEV exhibits a clear peak at 4 for the Uber tensor and 3 for Amazon, while our algorithm does not suffer from such an artifact.\n\n14. Thanks, corrected."}, {"review_id": "xIe5InSekp", "reviewer": "Reviewer_LxVj", "summary": "This 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).", "questions": "**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\"].", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer, and we have already updated our draft to address the typos and suggestions. The table below gives the requested comparisons to out-of\nthe-box methods. All times reported are average seconds \nper single ALS iteration; randomized algorithms were benchmarked \nwith $2^{16}$ samples per LSTSQ solve. We used Python \nTensorly 0.81 and Matlab Tensor Toolbox version \n3.5 on an identical system\nconfiguration to other experiments. OOM \ndenotes out-of-memory, for Tensorly due to explicit materialization \nof the Khatri-Rao product . Our method requires significantly\nless runtime across all tensors. Due to the slow \nperformance of CP-ARLS-LEV in Matlab, we developed a \nhigh-performance multithreaded version in C++, which serves \nas the baseline in the rest of the paper. See figures 9 and 10 in the appendix (supplement to original submission)\nfor more runtime comparisons.\n\n| Method    | Uber  | Enron | NELL-2 | Amazon | Reddit|\n| --------- | ----  | ----- | ------ | ------ | ----- |\n| Tensorly, Sparse Backend    | 64.2 | OOM  | 759.6 | OOM   | OOM    |\n| Matlab TToolbox Standard     | 11.6  | 249.4  | 177.4   | >3600 | OOM  |\n| Matlab TToolbox CP-ARLS-LEV  | 0.5 | 1.4 | 1.9  | 34.2  | OOM |\n| **STS-CP (Ours)** | 0.2 | 0.5 | 0.6  | 3.4  | 26.0  | \n\n\nOther Responses:\n\n**Questions:**\n\n1. [Single Element Factor Updates] Yes! There is a faster update time, which is mentioned in Theorem 1.1 (line 41 of the submitted draft): “If a single entry in a matrix $U_j$ changes, it [the data structure] can be updated in time $O(R \\log(\\left| I_j \\right| / R)$”. The procedure is detailed in Algorithm 5 in Appendix A.7, “Efficient Single-Element Updates”. The relatively simple method involves selectively updating matrices stored by binary tree nodes on the path from a leaf to the root, with the elements updated in each cached matrix depending on the index of the entry changed within $U_j$. Due to space constraints, we did not highlight this contribution more in the main body of the draft; we are happy to do so in the revision.\n\n2. [Table 1 I vs I_j] We intended (but did not make clear in the caption) that the input tensor has dimensions $I \\times I \\times … \\times I$ (i.e. all mode sizes $I_j = I$) to simplify the complexities. Thanks, we will clarify this.\n\n3. [S[j, i] vs S[i, j] ] Corrected, thanks.\n\n4. [Tilde on lower bound for J] The tilde denotes a hidden constant multiplying the right-hand side $R \\max (\\log (R / \\delta), 1 / (\\varepsilon \\delta)))$. Woodruff [1, Theorem 2.11] reports the constant value as 144 multiplying the term $R \\log (R / \\delta)$, while Malik [2, S1 supplement] reports a value of $8 / 3$ as sufficient. Line 99: agreed, we have changed this to $\\Omega$.\n\n5. [Kronecker Sketching Connection] Thanks for pointing out this connection - we plan to include this in the section on Kronecker sketching.\n\n6. [Increasing Sample Counts] We agree that the sample count should increase with the target rank in theory. However in figure 4, we would have to increase the sample count by different rates for different algorithms, since our algorithm STS-CP requires $O(R)$ samples, while CP-ARLS-LEV requires $O(R^{N-1})$ samples (see also Figure 6). Furthermore, our algorithm STS-CP performs better than worst-case analysis suggests when the sample budget is fixed, achieving 99.7% of the fit of exact ALS even at rank 125 on the Amazon tensor.\n\n   To avoid these confounding effects that vary between different algorithms and tensors, we used a fixed sample count throughout figure 4 to **directly compare the sample efficiency of our methods for a fixed budget**, as well as quantify the accuracy degradation as the rank increases. See figures 6, 8, 9, and 10 for experiments that vary the sample counts.\n\n**Weaknesses**\n1. [Thereom 3.1 from Malik 2022]: You are correct. The novelty in our approach lies in the strategy to sample from the distribution given in Theorem 3.1, a critical improvement that enables scaling to massive sparse tensors with up to **several thousand times** less compute for decompositions of *identical* quality.\n\n   Suppose the algorithm in [Malik 2022] was applied to the Amazon Reviews tensor in our experiments with dimensions $4,821,207 \\times 1,774,269 \\times 1,805,187$ to produce a rank $R=25$ decomposition. After extra non-asymptotic improvements made in their paper, the exact floating point operation (FLOP) count to draw a sample from the KRP excluding the third mode is lower-bounded by $\\left| I_2 \\right| R^2$, or 1.12 gigaFLOPs *per row*. Including multiplicative constants but excluding the small $(25 \\times 25)$-sized eigendecompositions (performed only per batch of rows, not once per row), **our** approach requires only 53.7 kiloFLOPs per row sample, a more than **20,000x reduction** that is due to the asymptotic improvement from $\\left| I_j \\right|$ to $\\log \\left| I_j \\right|$.\n\n   We have revised our draft to emphasize these points.\n\n2. [Interval Endpoints] Agreed, we revised our draft to make the cut points explicit. If we let $v_1, …, v_{\\lceil I / F \\rceil}$ be leaf nodes such that the intervals $S(v_1), …, S(v_N)$ are ordered from left to right, the explicit formula for $S_0(v_i)$ is\n$S_0(v_i) = (i-1) * F$, so that each segment has at most $F$ rows. Our draft has been updated to reflect this. The method to choose the cut points for sparse factors is slightly more involved, and the procedure is given in appendix A.8.\n\n3. [Corollary 3.3] Corrected, thank you.\n\n4. [Runtime Comparison] Agreed! Beyond Figures 6, 8a, and 8b, a thorough runtime comparison against CP-ARLS-LEV is made in Appendix A.12.4, figures 9 and 10, which measure the speedup of our algorithm over CP-ARLS-LEV. We tested a range of sample counts for both CP-ARLS-LEV and our algorithm, testing the former method on a significantly larger range of sample counts for fairness (see Table 5)."}]

Author rebuttal:
We thank all the reviewers for their constructive comments. The PDF attached to this rebuttal contains figures that were requested by some reviewers. These include diagrams that illustrate our sampling procedure and updated graphs with error bars. Below, you can find a summary of some revisions to our draft common to multiple reviewers, with specific responses in the individual rebuttals to each review.

1. We revised our draft to use $\tilde O$ notation instead of $O$ notation, where applicable, to account for the factors $\log(1 / \delta)$ that appear in leverage-score sampling.

2. We provided diagrams illustrating Theorem 3.1 and the two-stage sampling process, Figure 1 in the attached PDF.

3. We have provided baseline comparisons to state-of-the-art tensor decomposition libraries, such as Tensorly and Tensortoolbox, that demonstrate the dramatic efficiency of our algorithm while maintaining approximation quality. The responses to reviewers u8zc and LxVj contain this table with average runtime per ALS iteration for our method vs. Tensorly and Matlab TensorToolbox 3.5.

We hope that these changes improve the clarity of the draft, and are happy to engage in further discussion.

Finally, reviewer qocR requested a revision of the last paragraph of section 3.2, which we include here due to space constraints in the comment below that review. Point taken, this last paragraph is dense and can be simplified. Here is our revision of that paragraph:

"To summarize, Algorithms 1 and 2 give the construction procedure and two-stage sampling algorithm described above. The subroutines "BuildSampler" and "RowSample" relate to the procedure to build and sample, respectively, from the data structure in Lemma 3.2. The construction procedure builds the samplers $Z_j$, $1 \leq j \leq N$, for the second phase of sampling. The construction cost is $O(\left| I_j  \right| R^2)$ per matrix $U_j$. The sampling algorithm returns $J$ samples from the Khatri-Rao product of all matrices (excluding possibly index $U_j$, a useful feature for tensor decomposition applications). Lines 2-5 construct the sampling data structures for the first phase of sampling, while lines 9-11 implement the two-stage procedure by calling RowSample twice in succession and updating the running vector $h_{<k}$."

Meta-review: Dear Authors,

Thank you for your valuable contribution to Neurips and the ML community. Your submitted paper has undergone a rigorous review process, and I have carefully read feedback provided by the reviewers and considered the author rebuttal in detail.

This paper proposes a data structure that enables sampling rows from a Khatri-Rao product with respect to leverage scores. The key idea is using a "binary-tree inversion sampling" strategy. Since such matrices arise in least squares problems arising in tensor decompositions, the proposed method can be used in accelerating approximate solutions of such problems. The reviewers all agree that the contribution is solid and interesting.

Given this positive assessment, I am willing to recommend the acceptance of your paper for publication.

I would like to remind you to carefully review the reviewer feedback and the resulting discussion. While most reviews were positive, the reviewers have offered valuable suggestions that can further strengthen the quality of the paper. Please take another careful look a the 'weaknesses' section of each reviewer comment. Please also review the discussions with Reviewers LxVj and u8zc in order to resolve any potential confusion. I encourage you to use this feedback to make any necessary improvements and refinements before submitting the final version of your paper.

Once again, thank you for submitting your work to Neurips.

Best,
Area Chair
Decision: Accept (poster)