Efficient Leverage Score Sampling for Tensor Train Decomposition

Authors: vivek  bharadwaj
beheshteh t rakhshan
osman asif malik
guillaume  rabusseau
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 06:14:13

Abstract: Tensor Train~(TT) decomposition is widely used in the machine learning and quantum physics communities as a popular tool to efficiently compress high-dimensional tensor data. In this paper, we propose an efficient algorithm to accelerate computing the TT decomposition with the Alternating Least Squares (ALS) algorithm relying on exact leverage scores sampling. For this purpose, we propose a data structure that allows us to efficiently sample from the tensor with time complexity logarithmic in the product of the tensor dimensions. Our contribution specifically leverages the canonical form of the TT decomposition. By maintaining the canonical form through each iteration of ALS, we can efficiently compute (and sample from) the leverage scores, thus achieving significant speed-up in solving each sketched least-square problem. Experiments on synthetic and real data on dense and sparse tensors demonstrate that our method outperforms SVD-based and ALS-based algorithms.

Main paper: 
 introduction: 
 Tensor decomposition methods have recently found numerous applications in machine learning. Their ability to perform operations efficiently on very high-dimensional tensors makes them suitable for data science and machine learning problems. For example, they have been used for neuro-imaging, and signal processing [Zhou et al., 2013 , Sidiropoulos et al., 2017 , Cichocki and Phan, 2009] , supervised learning [Novikov et al., 2016, Stoudenmire and Schwab, 2016] , feature extraction [Bengua et al., 2015] and scaling up Gaussian processes [Izmailov et al., 2018] . The most popular decompositions are the CANDECOMP/PARAFAC (CP), Tucker [Hitchcock, 1927 , Tucker, 1966] , Tensor Train (TT) Oseledets [2011] and Tensor Ring (TR) Zhao et al. [2016] decompositions. However, finding a rank-r CP decomposition is NP-hard [Kolda and Bader, 2009, Hillar and Lim, 2013] and in the Tucker decomposition, the number of parameters grows exponentially with the order of the tensor. While the number of parameters is linear in the tensor's order for both TT and TR decompositions, TR is known to have numerical stability issues. The TT decomposition, in contrast, can be used to represent a tensor in a compressed format where the number of parameters scales linearly with the order of a tensor. Additionally, finding a good approximation to the best rank-r TT decomposition is feasible which lets the user perform various linear algebra operations efficiently and compute the TT decomposition with stable algorithms.Due to the high-dimensional nature of tensors, designing efficient algorithms for computing the TT decomposition is crucial. A popular method for computing the TT decomposition of an Ndimensional tensor X is the TT-SVD algorithm [Oseledets, 2011] which uses a sequence of singular values decompositions on the tensor unfoldings to produce the TT representation in a single pass.Our Contributions. In this paper, we propose a new sampling-based ALS approach to compute the TT decomposition: rTT-ALS. By using exact leverage score sampling, we are able to significantly reduce the size of each ALS least squares problem while providing strong guarantees on the approximation error. At the core of rTT-ALS, we leverage the TT canonical form to efficiently compute the exact leverage scores and speed up the solutions of least square problems in each iteration of ALS. To the best of our knowledge, rTT-ALS is the first efficient TT decomposition by the ALS algorithm which relies on leverage scores sampling. We provide experiments on synthetic and real massive sparse and dense tensors showing that rTT-ALS can achieve up to 26× speed-up compared to its non-randomized counterpart with little to no loss in accuracy.Our core contribution is the following theorem, which shows that we can efficiently compute a subspace embedding of a left-orthogonal chain of TT tensor cores by efficiently sampling according to their squared row norms (for technical definitions and details see subsection 3.1). Theorem 1.1 (Row-norm-squared sampling for 3D core chains). Let A 1 , ..., A j be a sequence of 3D tensors, A k ∈ R R k-1 ×I k ×R k (with R 0 = 1). Assume that the left-matricization of each core is orthogonal. Let A ≤j be the j k=1 I k × R j matrix obtained by unfolding the contraction of the tensor chain A 1 , ..., A j . Then there exists a data structure to randomly sample rows from A ≤j according to the distribution of its squared row norms with the following properties:1. The data structure has construction time Oj n=1 I n R n-1 R 2 n . When R = R 1 = ... = R j and I = I 1 = ... = I j , the runtime is O(jIR 3 ). The space overhead of the data structure is linear in the sizes of the input cores.2. The data structure produces a single row sample from A ≤j according to the distribution of its squared row norms in timeO j k=1 log (I k R k-1 /R k ) R 2 k . When all ranks R k and physical dimensions I k are equal, this complexity is O(jR 2 log I).We highlight that the runtime required to construct the data structure is asymptotically identical to the runtime required to compute the canonical form of the tensor train subchain, i.e., A ≤j , by successive QR decompositions. This implies that the data structure construction and subsequent updates of the data structure do not increase the asymptotic complexity of our method. 
related work: 
 Randomized algorithms and leverage score sampling-based methods [Mahoney et al., 2011 , Woodruff et al., 2014 , Drineas et al., 2006a] have been used widely in a large body of research particularly in tensor decomposition problems over the past two decades [Malik and Becker, 2021 , Larsen and Kolda, 2022 , Fahrbach et al., 2022] just to name a few. [Cheng et al., 2016] propose SPALS, the first ALS-based algorithm relying on leverage score sampling for the CP decomposition. Their proposed method reduces the size of the least squares problem in each iteration of ALS with a sub-linear cost per iteration in the number of entries of the input tensor. Larsen and Kolda [2022] extends this method by combining repeated sampled rows in a deterministic and random sampling fashion. However, both of these methods use leverage score approximations and therefore require a number of samples which is exponential in the number of tensor modes in order to achieve relative-error performance guarantees. Malik [2022] proposes a method which avoids this exponential dependency on the number of tensor modes by using higherquality leverage score estimates for the CP decomposition. The method is further improved by [Malik et al., 2022] to use exact rather than approximate leverage scores which is applicable for arbitrary tensor decompositions. Recently, [Bharadwaj et al., 2023] provided a novel data structure to efficiently sample from the exact distribution of the factor matrices' leverage scores in the Khatri-Rao product with time complexity logarithmic in the tensor size, leading to further improvements on the work in [Malik et al., 2022] . Moreover, [Malik and Becker, 2021] proposed an ALS-based algorithm for finding the TR decomposition using leverage scores approximation. However, the runtime of their method has an exponential dependency on the order of a tensor. The sampler we propose in this paper is built on the work by Bharadwaj et al. [2023] , extending it to the TT decomposition and leveraging the canonical form for further speed-up.There are also a variety of non-ALS-based randomized algorithms for computing the TT decomposition. [Huber et al., 2017] leverages randomized SVD for the TT decomposition which accelerates the classical TT-SVD algorithm proposed by [Oseledets, 2011] . To handle situations where the exact TT rank is unknown, [Che and Wei, 2019] propose an adaptive randomized algorithm that can achieve near optimal TT approximation. [Yu et al., 2023] present a method leveraging randomized block Krylov subspace iteration for computing TT approximations. Most of the algorithms for TT decomposition are based on the randomized SVD for matrices introduced by [Halko et al., 2011] . In the quantum physics community, the ALS algorithm is widely used for finding TT decomposition and often yields more accurate results than TT-SVD. The randomized TT-SVD method struggles to scale for high-order tensors as it requires generating a random Gaussian matrix at each step that can only handle small-order tensors. More closely related to our work are those using sketching and sampling in each iteration of ALS to approximate the TT decomposition. Recently, [Chen et al., 2023] introduced an algorithm that employs TensorSketch [Pham and Pagh, 2013] in each iteration of a regularized ALS approach for TT decomposition. However, the sketch size has an exponential dependency to a tensor dimension. By contrast, our proposed algorithm requires the sketch size with no dependence on the tensor dimension and it depends only on the column size of the design matrix and ε and δ parameters. 
preliminaries: 
 We use capital letters A to denote matrices and script characters A to denote multidimensional arrays. We use Matlab notation for slices of matrices and tensors. We use the tuple notation to indicate the position of entries of arrays. For example, A(i 1 , i 2 , i 3 ) indicates the (i 1 , i 2 , i 3 )-th element of A. A [i, :] and A [:, i] refer to the i-th row and column of A, respectively; for a three-dimensional tensor A ∈ R R1×I1×R2 , the matrix A [:, i, :] ∈ R R1×R2 is the i-th lateral slice of A. For a positive integer n, we use [n] to denote the set of integers from 1 to n.For i 1 ∈ [I 1 ], . . . , i N ∈ [I N ], the notation i 1 . . . i n def = 1 + N n=1 (i n -1)N -1 j=1 I j will be helpful for tensor unfoldings. We use ⊗ and ⊙ to denote the Kronecker and Khatri-Rao products, respectively (see definitions in Appendix A). We use I d to denote the d × d identity matrix, A T for the transpose of A , A + for the pseudo-inverse of A, ∥•∥ F for the Frobenius norm and ∥•∥ 2 for the Euclidean norm of a vector. We use Õ to indicate the presence of multiplicative terms polylogarithmic in R and 1/δ.Let X ∈ R I1×•••×I N be an N -dimensional array. A rank (R 1 , . . . , R N -1 ) tensor train (TT) de- composition of a tensor X ∈ R I1×•••×I N factorizes it into the product of N third-order tensors A n ∈ R Rn-1×In×Rn for n ∈ [N ] (with R 0 = R N = 1): X (i 1 , • • • , i N ) = r0,••• ,r N N n=1 A n (r n-1 , i n , r n ), for all i 1 ∈ [I 1 ], • • • , i N ∈ [I N ],where each r n ranges from 1 to R n . A tensor network representation of a TT decomposition is shown in Figure 1 . We call A 1 , A 2 , • • • , A N core tensors and we use TT((A n ) N n=1 ) to denote a TT tensor with factorsA 1 , • • • , A n . X i 1 i 5 i 2 i 4 i 3 = A 1 A 2 A 3 A 4 A 5 i 2 i 3 i 4 i 1 i 5 R 1 R 2 R 3 R 4Figure 1 : Tensor Train decomposition of a 5-dimensional tensor in tensor network notation.Definition 3.1. The mode-n unfolding of a tensorX ∈ R I1×•••×I N is the matrix X (n) ∈ R In× j̸ =n Ij defined element-wise by X (n) i n , i 1 • • • i n-1 i n+1 • • • i N def = X (i 1 , • • • , i N ).As a special case, we denote the left (resp. right) matricization of a 3-dimensional tensorA ∈ R I1×I2×I3 by A L = (A) ⊤ (3) ∈ R I1I2×I3 and A R = A (1) ∈ R I1×I2I3 .Given a TT decomposition TT((A n ) N n=1 ) and an index j, we will often use the left-chainA <j ∈ R j-1 k=1 I k ×Rj-1 and right-chain A >j ∈ R Rj × Nk=j+1 I k unfoldings obtained by matricizing the contraction of all cores on the left and on the right side of the j-th core. Formally,A <j (i <j , r j-1 ) = r0,...,rj-1 j-1 k=1 A k (r k-1 , i k , r k ) and A >j (r j , i >j ) = rj+1,...,r N N k=j+1 A k (r k-1 , i k , r k ) where i <j = i 1 . . . i j-1 and i >j = i j+1 . . . i N . We also use A ̸ =j def = A <j ⊗ A ⊤ >j ∈ R k̸ =j I k ×Rj-1Rjto denote the unfolding of the contraction of all cores except the j-th one. We conclude by introducing the canonical form of the TT decomposition [Holtz et al., 2012 , Evenbly, 2018 , 2022] which will be central to the design of our algorithm. Definition 3.2. A TT decomposition TT(( Note that any TT decomposition can efficiently be converted to canonical form w.r.t. any index j ∈ [N ] by performing a series of QR decompositions on the core tensors [Holtz et al., 2012 , Evenbly, 2018] .A n ) N n=1 ) ∈ R I1×•••×I N is in a canonical format with respect to a fixed index j ∈ [N ] if A L n ⊤ A L n = I Rn for all n < j, and A R n A R n ⊤ = I Rn-1 for all n > j (see Figure 2). A 1 i 1 A 2 i 2 A 3 i 3 A 4 i 4 A 5 i 5 R 1 R 2 R 3 R 4The TT decomposition problem consists in finding a low-rank approximation TT((A n ) N n=1 ) of a given tensor X : argmin A1,...,A N ∥X -TT(A 1 , . . . , A N )∥ F where X is the target tensor with dimensionsI 1 × • • • × I N .Since this is a non-convex optimization problem, the popular alternating least-squares (ALS) approach can be used to find an approximate solution [Kolda and Bader, 2009] . Fixing all cores except the j-th one, the low rank approximation problem can be reformulated as a linear least squares problem:EQUATIONThe ALS approach finds an approximate solution by keeping all cores fixed and solving for the j-th one. Then repeat this procedure multiple times for each j ∈ [N ] until some convergence criteria is met. While ALS is the workhorse algorithm in many tensor decomposition problems, the computational cost is still O(I N ) for I 1 = • • • = I N = I. In this work, to reduce this cost, we will combine ALS with core orthogonalization to efficiently compute the exact leverage scores. This will also lead to a stable algorithm for computing TT. To compute the orthogonalized TT approximation, we start with a crude TT decomposition in canonical form (see Definition 3.2) where all cores except the first one are right-orthonormal. After optimizing the first core, a QR decomposition is performed and the non-orthonormal part is merged into the next core. This procedure repeats until reaching the right side of the decomposition. The same procedure is then repeated from the right until reaching the left side (see the tensor network illustration in Appendix A.1). Even after computing the canonical form, which makes the design matrix orthonormal for each linear least squares problem, we still need to multiply the matricized tensor by the chain of TT cores in canonical form. Given that the tensor may have hundreds of millions of nonzeros, this multiplication is a significant computational bottleneck: without sketching, the runtime for this matrix multiplication scales as O(nnz(X )N R 2 ), where nnz(X ) is the number of nonzeros in the tensor. Sketching allows us to select only a subset of rows from the design matrix and the corresponding subset of rows from the matricized tensor, reducing the cost to O(N R 2 ). This approach also leads to providing computational benefits for computing the leverage scores and to an efficient sampling scheme which will be discussed in Section 4.There exists a vast literature on randomized algorithms [Mahoney et al., 2011 , Woodruff et al., 2014] to solve the over-determined least squares problem min x ∥Ax -b∥ F where A ∈ R I×R , I ≫ R.Regardless of the structure of both A and b, solving this least-squares problem costs O(IR 2 ). To reduce this cost, we can randomly select rows of A and b by proposing a sketching operator S with J ≪ I. Therefore, instead of solving the original least squares problem, we consider solving the downsampled version of the form min x ∥SAx -Sb∥ F , where S ∈ R J×I and reduce the cost to O(JR 2 ). The goal is to find a "good" sketch S to approximate the solution of the least squares problem at each step of the ALS algorithm. When each entry of S is selected according to the rows of A leverage scores, strong guarantees can be obtained for the solution of the downsampled problem. Definition 3.3. (Leverage scores) Suppose A ∈ R I×R with I ≫ R. The i-th leverage score of the matrix A is defined asEQUATIONDefinition 3.4. (Leverage score sampling) Let A ∈ R I×R and p ∈ [0, 1] I be a probability distribution vector with entries p i = li(A) rank(A) ; where rank(A) = i l i (A). Assume ŝ1 , ..., ŝJ are drawn i.i.d according to the probabilities p 1 , • • • , p I . The random matrix S ∈ R J×I defined element-wise by S(j, i) = 1 √ Jpi if ŝj = i and 0 otherwise is called a leverage score sampling matrix for A.The following result is well-known and appeared in several works; see, e.g., [Drineas et al., 2006b] , [Drineas et al., 2008] , [Drineas et al., 2011] , [Larsen and Kolda, 2022] . We borrow the form presented in [Malik, 2022] . Theorem 3.5. (Guarantees for Leverage Score Sampling) Suppose A ∈ R I×R . Let S ∈ R J×I be the leverage score sampling matrix defined in 3.4. For any ε, δ∈ (0, 1), if J = Õ(R 2 /εδ), then x * = min x ∥SAx -Sb∥ 2 satisfies ∥Ax * -b∥ 2 ≤ (1 + ε) min x ∥Ax -b∥ 2 , with probability 1 -δ.Computing leverage scores in Definition 3.3 requires computing the pseudo-inverse of A, which costs O(IR 2 ) and is as costly as directly solving the original least squares problem. In the following section, we will show that the leverage scores can be computed much more efficiently when A is the matrix appearing in the TT-ALS algorithm in canonical form.. 
sampling-based tensor train decomposition: 
 In this section, we show how to efficiently sample rows of A ̸ =j = A <j ⊗A ⊤ >j and X (j) in Equation (1) according to the exact leverage scores distribution. In doing so, we will also present the sketch of the proof of Theorem 1.1 (which closely mirrors that of Bharadwaj et al. [2023] with key modifications required to adapt the procedure to a tensor core chain).̸ =j = i 1 . . . i j-1 i j+1 . . . i N of A ̸ =j , Equation (2) gives l i ̸ =j (A ̸ =j ) = A ̸ =j [i ̸ =j , :](A ̸ =j ⊤ A ̸ =j ) + A ̸ =j [i ̸ =j , :] ⊤ . (3) Computing Φ def = (A ̸ =j ⊤ A ̸ =j ) +is the main computational bottleneck in finding the leverage scores of A ̸ =j . Malik et al. [2022] proposed an algorithm to compute Φ in time O(N IR 2 + R 3 ). In this paper, we leverage the fact that when the TT tensor is in canonical form w.r.t. mode j, A ̸ =j is orthogonal, and thus Φ = I R 2 . Therefore, computing Φ is free of cost. By maintaining the canonical form of the TT tensor throughout the ALS algorithm, we can sketch the least square problems from the leverage score distributions with almost no computational overhead. We now explain how to efficiently sample rows of A ̸ =j from the leverage scores distribution.As discussed above, when the TT tensor is in canonical form, the leverage score of row i ̸ =j is given by li ̸ =j (A ̸ =j ) = A ̸ =j [i ̸ =j , :]A ̸ =j [:, i ̸ =j ] ⊤ . Leveraging the Kronecker structure of A ̸ =j = A <j ⊗ A ⊤ >j , one can easily show that l i ̸ =j (A ̸ =j ) = l i<j (A <j ) • l i>j (A ⊤ >j ).Sampling from the leverage scores distributions thus boils down to sampling rows of A <j and A ⊤ >j with probability proportional to their squared row norms (due to the orthogonality of A <j and A >j inherited from the canonical form). Without loss of generality, we detail the sampling procedure for A ≤j (the difference between A ≤j and A <j amounts to reindexing). The sampling procedure for A >j will be the same and straightforward.Let ŝ1 ∈ [I 1 ] , ..., ŝj ∈ [I j ] be random variables such that the multi-index ŝ≤j = ŝ1 . . . ŝj follows the leverage score distribution of A ≤j . Since TT((A n ) N n=1 ) is in canonical form w.r.t. j + 1, A ≤j is an orthonormal matrix, hence ŝ1 . . . ŝj is selected with probability proportional to the squared norm of the corresponding row of A ≤j :EQUATIONOur sampling procedure will draw a lateral slice from each core starting from A j and ending with A 1 , corresponding to a single row of A ≤j . Suppose we have drawn s k+1 , . . . , s j , for some k < j. To sample the k-th index, we need to compute the conditional probability p(s k |s k+1 , . . . , s j ) = p(s k ,...,sj ) p(s k+1 ,...,sj ) . The following lemma shows that this can be done efficiently by leveraging the underlying TT structure. Lemma 4.1 (Conditional distribution for ŝk ). Consider the events ŝj = s j , . . . , ŝk+1 = s k+1 , which we abbreviate as ŝ>k = s >k . Thenp(ŝ k = s k | ŝ>k = s >k ) ∝ Tr H ⊤ >k • A k [:, s k , :] ⊤ • A k [:, s k , :] • H >k ,whereH >k := A k+1 [:, s k+1 , :] • ... • A j [:, s j , :] .The proof is given in Appendix B. Intuitively, H >k acts as a "history matrix" conditioning on s >k , while the trace operation corresponds to marginalization over s <k . Unfortunately, updating H >k through matrix multiplication as each index is selected still requires timeO(R 3 ) (assuming R 1 = ... = R j = R).In order to further improve the runtime and reach the quadratic complexity in R claimed in Theorem 1.1, we make the following observation: let q ∈ R i≤j Ii be the probability vector for the leverage score distribution of A ≤j . Then Equation (4) can be rewritten in vectorform as q := 1 Rj A ≤j [:, 1] 2 + ... + A ≤j [:, R j ] 2 .Here, the square of each column vector is an elementwise operation. Observe that each A ≤j [:, r] 2 is a probability vector (positive entries summing to one) due to the orthonormality of A ≤j . Hence q is a mixture distribution. To sample from q, it thus suffices to select a single column r of A ≤j uniformly at random and restrict the sampling procedure to A ≤j [:, r]2 . More formally, let r be uniformly distributed over [R j ] and let t1 , ..., tj follow the conditional distributions defined byEQUATIONwhere h >k = A k+1 [:, t k+1 , :] • ... • A j [:, t j , r]. We have the following result.Lemma 4.2. For any choice of s j , ..., s k , fix s j = t j , s j-1 = t j-1 , ..., s k = t k . After marginalizing over r, the conditional distribution of tk satisfies p(tk = t k | t>k = t >k ) = p(ŝ k = s k | ŝ>k = s >k ).As a consequence, the joint random variable ( t1 , ..., tj ) follows the desired squared row-norm distribution of A ≤j after marginalizing over r. The proof appears in Appendix B.2. Notice that the "history matrix" H >k has been replaced by a vector h >k . This vector can be updated by matrix-vector multiplication, yielding a reduced sampling complexity with only a quadratic dependency on R.Our final improvement is to show that each sample from the distribution in Equation ( 5) can be drawn in time sublinear in the dimension I k (after appropriate preprocessing). Letting A L k be the left unfolding of A k , one can check thatEQUATION)The probability of selecting the slice s k is thus the sum of R k-1 consecutive entries from the probability vector (A L k • h >k ) 2 . As a result, we can sample tk by first sampling an index in the range[I k R k-1 ] given by (A L k •h >k ) 2, then performing integer division by R k-1 to obtain the corresponding slice index tk . The advantage here lies in an efficient data structure for sampling from the weight vector (A L k • h >k ) 2 , given by the following lemma: for k = j...1 do 5:tk := RowSample(Z k , h)//R k-1 6: h = h • A k :, tk , : 7: t d = ( tk ) k≤j 8: return t 1 , ..., t JThe adaptation of this lemma is given in Appendix B.3. Lemma 4.3 enables us to efficiently draw samples according to the distribution in Equation 6, and therefore gives us a procedure to sample from the entire core chain. Constructing the data structure above for each matrixA L k , 1 ≤ k ≤ j, costs O(IR k-1 R 2 k )with a linear space overhead in the input core sizes. Drawing a sample from the k-th data structure re-quires time O(R 2 k log(I k R k-1 /R k )). Summing up this runtime over 1 ≤ k ≤ j gives the stated complexity in Theorem 1.1. Algorithms 1 and 2 summarize the procedures to efficiently draw J samples from a left-orthogonal core chain. The construction procedure builds a set of data structures Z k of the form given by Lemma 4.3 on the left-matricization of each tensor core. For each of J rows to draw, the sampling algorithm selects a column t uniformly at random from the left matricization. It then initializes the history vector h and successively samples indices tj-1 , ..., t1 according to the conditional distribution, updating the history vector at each step. Appendix B.4 provides a rigorous proof of the correctness of the procedure sketched in this section.While our procedure shares similarities with the Khatri-Rao product leverage score sampler, significant adaptations are required to sample from a tensor train core chain. The factors of a Khatri-Rao product can be sampled in any order, since the Khatri-Rao product of several matrices is commutative up to a permutation of its rows. By contrast, our sampling procedure requires us to sample from core A j down to A 1 , since Lemma 4.1 exploits the left-orthogonality of the each core in its derivation. Starting the sampling procedure at A j leads to a "history matrix" to keep track of prior draws instead of the vector that would arise starting from core A 1 . Here, our second innovation of sampling a column uniformly at random is required to bring down the overall sampling complexity. We can now state the following guarantee for Randomized-TT-ALS (rTT-ALS) applying the data structure in Theorem 1.1. The proof is given in Appendix B.5.Corollary 4.4. (rTT-ALS) For any ε, δ ∈ (0, 1) the sampling procedure proposed above guarantees that with J = Õ(R 2 /εδ) samples per least-square problem, we haveA ̸ =j ( Ãj ) ⊤ (2) -X ⊤ (j) ≤ (1 + ε) min (Aj ) (2) A ̸ =j (A j ) ⊤ (2) -X ⊤ (j) ,with probability (1 -δ), where Ãj is the solution of the sketched least-squares problem, for all least-squares solve. The efficient sampling procedure of Theorem 1.1 brings the overall complexity toÕ #it εδ R 4 • N j=1 N log I j + I j ,where "#it" is the number of ALS iterations.Algorithms 1 and 2 refer to procedures "BuildSampler" and "RowSample" that were first used to sample from the Khatri-Rao product. The k-th BuildSampler data structure creates a full binary tree that truncated tolog⌈I k R k-1 /R k ⌉ levels, each caching an R k × R k matrix containing information from A L k .To draw a sample, the RowSample procedure executes a random walk from the root to the leaves requiring O(R 2 k ) work at each internal node. The sampler performs matrix-vector multiplication with the cached data at each internal node and compares the output value to a threshold, using the comparison to branch either left or right in the random walk. Assuming R 1 = ... = R j and I = I 1 = ... = I j , the storage cost of each sampler is O(IR 2 ) (computed by multiplying the matrix size stored at each node by the maximum node count). By multiplying the tree depth by the matrix-vector multiplication cost at each node, we get runtime cost O(R 2 log I) to draw one sample from A L k . For the motivation behind the procedure, details of its correctness, and pseudocode, we refer the reader to the original work [Bharadwaj et al., 2023] . 
experiments: 
 All experiments were conducted on CPU nodes of the NERSC Perlmutter, an HPE Cray EX supercomputer, and the Mila Quebec AI Institute compute cluster. Our code is available at https://github.com/vbharadwaj-bk/ortho_tt_subspace_embedding. In this section, we demonstrate the effectiveness of the proposed rTT-ALS on two types of tensors: (i) synthetic and real dense datasets and (ii) real sparse datasets. We use the fit as evaluation metric (higher is better): fit( X , X ) = 1 -∥ X -X ∥ F /∥X ∥ F , where X is the TT approximation and X is the target tensor. The goal of the dense tensor experiments is to show that rTT-ALS has a better time complexity than TT-ALS and TT-SVD while matching rTT-SVD in terms of fit. The sparse tensor experiments show that SVD-based decompositions cannot handle high-order (sparse) tensors. We compare rTT-ALS with the classical TT-ALS. The runtime improvements are most significant for large sparse tensors. Figure 4 compares accuracy (y-axis, higher is better) against ALS iteration time, for rTT-ALS versus non-randomized ALS. The speedup per iteration can be as high as 26x for lower ranks. Particularly, for the NELL-2 tensor, the plot shows that accuracy within three significant figures of non-randomized ALS was achieved roughly 3-4x faster than an optimized non-randomized ALS baseline.We compare rTT-ALS to three other methods; TT-SVD [Oseledets, 2011] , Randomized TT-SVD (rTT-SVD) [Huber et al., 2017] and TT-ALS [Holtz et al., 2012] . We use TensorLy [Kossaifi et al., 2019] for SVD-based methods and our own implementation for deterministic TT-ALS. For simplicity, we setR 1 = • • • = R N -1 = R for all experiments.For all algorithms, we illustrate the quality of performance by fit and runtime.Synthetic Data Experiments. For the synthetic data experiment, we generate random tensors of size I × I × I for I ∈ {100, . . . , 500} and of TT rank R = 20 (by drawing each core's components i.i.d. from a standard normal distribution). A small Gaussian noise with mean zero and standard deviation of 10 -6 is added to each entry of the resulting tensor. We then run the four methods to find a rank R = 5 approximation of the target tensor. ALS-based methods are initialized using their SVD-based counterpart (TT-ALS with the output of TT-SVD and rTT-ALS with the output of rTT-SVD) and are run for 15 iterations. The sample count for rTT-ALS is fixed to J = 5000 for all values of I. The average fit over 5 trials for all four algorithms are reported as a function of the dimension in Figure 3 . rTT-ALS is about 2× faster than TT-ALS and 3× faster than TT-SVD for I = 500. Although rTT-SVD is the fastest method, it achieves poor performance in terms of fit. 191) . Both datasets are three-dimensional tensors where the first two dimensions are the image height and width, and the third dimension is the number of spectral bands, (iii) the MNIST dataset is of size (60000 × 28 × 28), and iv) Tabby Cat is the three-dimensional tensor of size (720 × 1280 × 286) which contains grayscale videos of a man sitting on a park bench and a cat, respectively. The first two dimensions are frame height and width, and the third dimension is the number of frames. For all datasets, the preprocessing step is done by tensorizing data tensors into higher-dimensional tensors. Table 1 illustrates the results for a single trial when R = 5. For all datasets we keep the sample count fixed to J = 2000. Similarly to the synthetic data experiments, rTT-ALS is faster than TT-ALS and TT-SVD (up to 10× faster than TT-ALS).We next apply rTT-ALS to three large sparse tensors from FROSTT [Smith et al., 2017] . Table 2 gives the fits achieved by our method to decompose these tensors. The largest of these tensors, NELL-2, has around 77 million nonzero entries with mode sizes in the tens of thousands. Fits for sparse tensor decomposition are typically low, but the factors of the resulting decomposition have successfully been mined for patterns [Larsen and Kolda, 2022] . For these experiments, we chose all decomposition ranks equal with R 1 = ... = R N = R and tested over a range of values for R.The fits produced by rTT-ALS match those produced by the non-randomized ALS method up to variation in the third significant figure for Uber and NELL-2, with slightly higher errors on the Enron tensor. We kept the sample count for our randomized algorithms fixed at J = 2 16 throughout this experiment. As a result, the gap between the fit of the randomized and exact methods grows as the target rank increases, which our theory predicts.Table 2 also reports the average speedup per ALS sweep of rTT-ALS over the exact algorithm. On the NELL-2 sparse tensor with target rank 12, the non-randomized ALS algorithm requires an average of 29.4 seconds per ALS sweep, while rTT-ALS requires only 1.87 seconds. Figure 4 shows that our method makes faster progress than its non-randomized counterpart across all three tensors. Because we could not find a well-documented, high-performance library for sparse tensor train decomposition, we wrote a fast multithreaded implementation in C++, which serves as the baseline method in these figures and tables (the code will be made publicly available).Figure 5 shows the impact of varying the sample count on the final fit. We find modest increases in accuracy for both Uber and NELL-2 as the sample count increases by a factor of 5 (starting from J = 2 15 ). Increasing J has a smaller impact for the Enron tensor, which is generally more difficult to decompose beginning with i.i.d. random factor initialization [Larsen and Kolda, 2022]. 
conclusion: 
 We proposed a sampling-based ALS method leveraging an efficient data structure to sample from the exact leverage scores. More precisely, we show that by exploiting the canonical form of the TT decomposition, leverage scores can be computed efficiently for all the least squares problems of ALS. We provide strong theoretical guarantees for the proposed data structure. Experiments on massive dense and sparse tensors confirm the theoretical results. The sampling algorithm we proposed could be extended to more general tree-based tensor network structures, leveraging canonical forms in a similar spirit to rTT-ALS. A Additional Notation Figure 6 illustrates the single-site TT-ALS method, which begins with a TT decomposition in canonical form initialized by a crude guess. Core A 1 of the decomposition is non-orthogonal; in sweeps from left-to-right and right-to-left, the algorithm holds all but one core constant and solves for the optimal value for the remaining core. After updating each core, by a QR decomposition the non-orthonormal part is merged to the left or right (depending on the direction of the sweep), a step which is called core orthogonalization. A 1 i 1 A 2 i 2 A 3 i 3 A 4 i 4 A 5 i 5 R 1 R 2 R 3 R 4 step: 1 A 1 i 1 A 2 i 2 A 3 i 3 A 4 i 4 A 5 i 5 R 1 R 2 R 3 R 4 QR A 1 i 1 A 2 i 2 A 3 i 3 A 4 i 4 A 5 i 5 R 1 R 2 R 3 R 4 step: 2 A 1 i 1 A 2 i 2 A 3 i 3 A 4 i 4 A 5 i 5 R 1 R 2 R 3 R 4 QR A 1 i 1 A 2 i 2 A 3 i 3 A 4 i 4 A 5 i 5 R 1 R 2 R 3 R 4 step: 3 
 
 Appendix:  


Decision: After the rebuttal period and inter-reviewer discussion, the final ratings for this paper are {7, 5, 6, 7} since Reviewer BTzr raised their score from 4 to 5, allowing us to reach a unanimous accept decision.

I read the paper myself and am also fine with it being accepted. It's not as crisp as it should be (theory and experiments), but a paper on leverage score-based sampling for Tensor Train (TT) decompositions is missing from "faster tensor decomposition via leverage scores" line of work (CP decomposition in NeurIPS 2016, Tucker decomposition in NeurIPS 2022).

Reviews: review: paper_summary: paper_summary: This paper considers the problem computing the tensor train TT decomposition of large tensors and proposes a novel randomized approach for efficiently solving the problem. In particular, the Alternating Least Squares (TT-ALS) algorithm is considered and exact leverage scores sampling approach is proposed to accelerate the approach. A data structure based approach is devised to efficiently compute the leverage scores of the appropriate matricization of the tensors and sample rows from them. Theoretical analysis show the sampling complexity. Numerical experiment results are presented which to illustrate the performance of the proposed algorithm.
strength_weakness: strength_weakness: The weaknesses of the paper are:
1. The randomized SVD approach seems to be more efficient in terms cost compared the the prosed method.
2. Certain aspects of the paper can be improved (see below).
questions: questions: The paper presents an interesting approach to compute the TT decomposition of large tensors. 
However, I have the following comments that can further improve the paper:

1. r-TT-SVD seems better - 

Based on numerical results presented, it appears  the randomized SVD based approach  (r-TT-SVD) takes significantly lower time, yet achieves comparable performance (fit) as other methods. Are TT-ALS and rTT-ALS popular and relevant for TT decomposition. How do these methods differ in the quality of the decompositions computed?

In some cases, rTT-SVD might have a slightly lower fit. But, since the runtime is so low, perhaps a larger sketch size  would yield similar fit at lower cost. Details about the sketch sizes used are missing. 


2. Presentation:
Certain aspects of the presentation can be improved, and there are few minor typos.

i. Are TT decompositions popular in tensor applications? Overall, they seem to have higher computational cost and are not very interpretable. Few alternate tensor decompositions achieve better compression with lower cost. A discussion on the motivation for the use of TT decomposition will benefit the paper (readers can better appreciate the results).

ii. Introduction has tensor jargons which might not be known to general AI audience, such as 3D core chains, left-matricization, contraction of tensor, etc. It is better to avoid these or define them before use. 

iii. The computational  cost of solving eq (1) can be added.

iv. Theorem 1.1, point 1 has a $j$ missing in the runtime cost. 

v. Not sure what is the subscript (2) in Corollary 4.4.

Many of the results in the paper is based on another Arxiv paper (which is not peer-reviewed). So, correctness of these results is not established.
limitations: limitations: See above comments


------
Post Rebuttal:

I have read other reviews and authors' responses to all reviews. The responses have adequately addressed my concerns. The work has many merits. I have raised my score.
paper_summary: paper_summary: This paper presents an efficient algorithm to use leverage-score sampling to solve least squares problems arising as subproblems in a larger alternating least squares (ALS) algorithm for building an approximate Tensor Train (TT) decomposition. The paper reports empirical evaluation of the proposed algorithm, showing runtime improvement compare to the baseline un-sampled problems.
strength_weakness: strength_weakness: Major comments (affecting recommendation):
- The proposed method is compared empirically only to a small number of baselines. Seems that the most likely competatitor (Chen et al. 2023) is not compared against.
- Runtime improvement in the experiments is modest. This is a sampling based method, so if sampling is very efficient (due to efficient leverage score approximation) I would expect a big improvement in runtime. This might not be the case if either more iterations are now needed for convergence or computing the sampling probabilities is very expensive. Unclear from then experiments which is the case.
- The paper does not present an analysis end-to-end. It only analyzes how to subsample the ALS subproblems, and thus accelerate them. Not analyzed is how the solution of the subsampled system in lieu of the full subproblem affect ALS convergence.
- Many details required for implementation of the algorithm are omitted. Many points regarding the algorithm are not clear. 

Minor comments (do not affect the recommendation):
- Lines 125 - 126: add citation.
- Figure 3 left: is this fit in the y-axis or misfit?
- Line 108-109: I think you mean *omission* of multiplicative terms.
- Lines 58-60: shouldn't runtime be O(jIR^3)?
- Line 54: The term "left matricization" is used without defining it yet.
- Line 55: I think it should be R_j
- Eq (2): I think on the right it should be A[i,:]^T. Similar issue with Eq (3).
questions: questions: - Lines 96-97: Sketching is a good way to circumvent the need to compute leverage scores, and TensorSketch allows us to sketch matrices with Kronecker structure. And indeed, the authors point out that [Chen et al. 2023] do this. How do their results compare to the present paper? Why is what you are proposing a better idea than using TensorSketch?
- Lines 64-66: What is the significance of this observation regarding the complexity of the algorithm?
- Lines 125 - 126: One the canonical form has been computed, can we solve the ALS subproblem efficiently?
- Did you compare empirically your method to the method of [Chen et al. 2023].
- Is your algorithm computing the canonical form anew in each ALS iteration, or are you somehow updating it?
- Does your algorithm maintain/computes the leverage scores of each row (the number of such scores is exponential), or do you present only an efficient way to sample from it without forming it?
limitations: limitations: Nothing to add.
paper_summary: paper_summary: The authors proposed a leverage score sampling-based TT-ALS method to reduce the computational complexity of the traditional TT-ALS. Experimental results verify the performance of the proposed method.
strength_weakness: strength_weakness: The contribution and novelty are not clearly stated compared with [Malik and Becker, 2021]. In [Malik and Becker, 2021], the leverage sampling was applied to TR. As TT can be treated as a special case of TR, the author does not clearly state that it is necessary to develop a new method for TT, or that there are some new theories or findings that are distinct from the previous TR structure.
questions: questions: 1. My main concern about this work is its relation to the leverage score sampling-based TR-ALS [Malik and Becker, 2021]. As TT can be seen as a special case of TR, what is the main contribution of the current work compared with leverage score sampling-based TR-ALS? Please clarify that it is necessary to develop a new method beyond [Malik and Becker, 2021] (i.e. has a better theoretical guarantee or has better performance than leverage score sampling-based TR-ALS).

2. For real data experiments, how about the performance compared with other tensor decomposition structures, such as TR? As you can set the first and last TR-rank as 1 thus it reduces to a TT-like structure.
limitations: limitations: N/A
paper_summary: paper_summary: This paper gives a better randomized alternating least squares algorithm for computing tensor factorizations. It's based on an exact characterization of leverage scores of the matrization of tensors via a suitable intermediate orthonormal representation. This is justified rigorously, and significant empirical gains were demonstrated.
strength_weakness: strength_weakness: I'm a bit concerned about the setting of the experiments, which seem to be 3-dimensional dense tensors. My understanding is that a lot of the more complicated tensor instances are sparse and in higher dimensions. However, I'm not sure whether those have low rank representations.
questions: questions: Would it be possible to check how the orthogonality conditions are maintained in intermediate steps under the inexact arithmetic caused by round off errors? Aka. are the conditions for the characterizations of leverage scores preserved exactly? (this has been answered)
limitations: limitations: yes

Comments: title: [W1][Comparison against (Chen et al. 2023)]
thank you
comment: >Thank you for the detailed answers and clarifications. It will be beneficial to include many of the clarifications in revised versions of the manuscripts. 

Thank you again for you review and questions. We will include all the points mentioned in our rebuttal in the revisions which will improve the clarity.  

> However, let me followup the discussion on Chen et al. 2023. I see your point. However, note that Thm 4.4 in Chen et al. 2023 gives exact constants (no O() expression) and this is where the 3^q comes from. In contrast, your results is J=O(R^2/(\eps * \delta)) where the constant is unclear. For all I know, the O() hides a 3^q constant! So: if you flesh out the constant, do you really do not have any dependence on the tensor order?

There is no hidden dependency on the tensor order in the big O. 

In Corollary 4.4, the number of samples required for each least square problem, J=O(R^2/(\eps * \delta)), is independent on the tensor order N (N is not hidden in the constant). Thus, for one sweep of the ALS algorithm, the complexity is linear in the tensor order, as stated in the second part of the corollary (i.e. the overall ALS complexity is O(#it εδ R^4  N \sum_{j=1^N} (log I_j + I_j), where N is the order of the tensor). We will clarify this point in the revision. 

> even if we accept that there is a theoretical advantage in the form of a dependence of the order, this does not necessarily tell the whole story for a small q (where typically q=3 is the most common use-case). In fact, if q is fixed to 3, then 3^q is now a constant, and it really matters how it compares the constant of the your algorithm. 

We agree that if q is treated as a constant, then the asymptotic complexity is the same. However, we believe that designing an algorithm with only a linear dependency on the order, instead of the exponential one given in Chen et al, constitute a non-trivial contribution and advancement compared to Chen et al. In addition, the TT decomposition is particularly suited and relevant for (very) high order tensors where the tensor order cannot be treated as constant  (as demonstrated by its uses in quantum physics for simulating many-body systems where the order corresponds to the number of particles in the system). 

> I still hold that Chen et al. 2023 is the most appropriate empirical baseline, and without it the empirical evaluation is lacking.

Unfortunately, the lack of a publicly available code for the method presented in (Chen et al, 2023) did not allow us to include it in the synthetic experiment. Moreover, to the best of our understanding, the method proposed by Chen et. al could not scale to the size of the tensors that we consider in the large sparse tensors experiments.
Thank you for the detailed answers and clarifications. It will be beneficial to include many of the clarifications in revised versions of the manuscripts. I will consider how your answers affect my final score, which is "borderline" to begin with. 

However, let me followup the discussion on Chen et al. 2023. I see your point. However, note that Thm 4.4 in Chen et al. 2023 gives exact constants (no O() expression)  and this is where the 3^q comes from. In contrast, your results is J=O(R^2/(\eps * \delta)) where the constant is unclear. For all I know, the O() hides a 3^q constant! So: if you flesh out the constant, do you really do not have any dependence on the tensor order?

Re empirical comparison: even if we accept that there is a theoretical advantage in the form of a dependence of the order, this does not necessarily tell the whole story for a small q (where typically q=3 is the most common use-case). In fact, if q is fixed to 3, then 3^q is now a constant, and it really matters how it compares the constant of the your algorithm. 

I still hold that Chen et al. 2023 is the most appropriate empirical baseline,  and without it the empirical evaluation is lacking.
Thank you for the detailed responses.

My concern w.r.t. roundoffs is that equality conditions like equation (6) might not be robust to perturbations to the original matrix. However, I now see that it's only for the initial matrix, not for the intermediate products of sampling. So I agree that it's inherently robust.

I will raise my presentation score, and leave the overall unchanged.
Dear reviewer, considering that the end of the discussion period is approaching, do you have any further questions? We think we have answered all your concerns in our rebuttal and hope you'll consider increasing your score.
Reviewers: The author-reviewer discussions end tomorrow (August 13), so please read the authors' rebuttal carefully and respond if you haven't.
Thank the reviewer for their reply. I've increased the score.