Title: Efficient Leverage Score Sampling for Tensor Train Decomposition
Venue: NeurIPS 2024
Accepted: yes
Keywords: leverage score sampling

tensor train decomposition

alternating least square
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.

Efficient Leverage Score Sampling for Tensor Train Decomposition

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 tensor size. 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.

1 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 rankr 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 rankr 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 N -dimensional 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.

4

Since TT-SVD requires performing SVDs of unfoldings of X , its cost is exponential in N . Alternating Least Square (ALS) is another popular approach [Holtz et al., 2012] to find the TT approximation. Starting with a crude guess, each iteration of ALS involves solving a sequence of least squares problems. While ALS is the workhorse algorithm in many tensor decomposition problems, the computational cost is still exponential in the order of a tensor ( N ), since each iteration requires solving least squares problems involving unfoldings of X . These issues have led to the search for alternatives based on randomization and sampling techniques. A cheaper alternative to the TT-SVD with strong accuracy guarantees can be implemented by replacing the exact singular value decomposition (SVD) with a well-studied randomized counterpart [Halko et al., 2011, Huber et al., 2017]. Randomized variants of the TT-ALS approach have received little attention. In Chen et al. [2023], the authors propose a randomized ALS algorithm that uses TensorSketch [Pham and Pagh, 2013] in each iteration. In this work, we also propose a novel randomized variant of the TT-ALS algorithm that relies on exact leverage score sampling. Notably, the sketch size in TensorSketch TT-ALS Chen et al. [2023] has an exponential dependence on the tensor dimension I whereas our algorithm avoids any dependence of the sketch size on I .

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:

The data structure has construction time O ( ∑ j 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.

The data structure produces a single row sample from A ≤ j according to the distribution of its squared row norms in time O ( ∑ 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.

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

3 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 R 1 × I 1 × R 2 , the matrix A [: , i, :] ∈ R R 1 × R 2 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 ˜ O to indicate the presence of multiplicative terms polylogarithmic in R and 1 /δ .

3.1 Tensor Train Decomposition

Let X ∈ R I 1 ×···× I N be an N -dimensional array. A rank ( R 1 , . . . , R N -1 ) tensor train (TT) decomposition of a tensor X ∈ R I 1 ×···× I N factorizes it into the product of N third-order tensors A n ∈ R R n -1 × I n × R n for n ∈ [ N ] (with R 0 = R N = 1 ):

for all i 1 ∈ [ I 1 ] , · · · , i N ∈ [ I N ] , where each r n ranges from 1 to R n . Atensor 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 factors A 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

4

̸

Definition 3.1. The moden unfolding of a tensor X ∈ R I 1 ×···× I N is the matrix X ( n ) ∈ R I n × ∏ j = n I j 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 tensor A ∈ R I 1 × I 2 × I 3 by A L = ( A ) ⊤ (3) ∈ R I 1 I 2 × I 3 and A R = A (1) ∈ R I 1 × I 2 I 3 .

Given a TT decomposition TT(( A n ) N n =1 ) and an index j , we will often use the left-chain A <j ∈ R ∏ j -1 k =1 I k × R j -1 and right-chain A >j ∈ R R j × ∏ N k = 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,

̸

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 × R j -1 R j to 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(( A n ) N n =1 ) ∈ R I 1 ×···× I N is in a canonical format with respect to a fixed index j ∈ [ N ] if A L n ⊤ A L n = I R n for all n < j , and A R n A R n ⊤ = I R n -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

4

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

3.2 Alternating Least Squares with Tensor Train Structure.

The TT decomposition problem consists in finding a low-rank approximation TT(( A n ) N n =1 ) of a given tensor X : argmin A 1 ,..., A N ∥X TT( A 1 , . . . , A N ) ∥ F where X is the target tensor with dimensions I 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:

The 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 ) NR 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 ( NR 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.

3.3 Sketching and Leverage Score Sampling

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 as

Definition 3.4. (Leverage score sampling) Let A ∈ R I × R and p ∈ [0 , 1] I be a probability distribution vector with entries p i = l i ( A ) rank ( A ) ; where rank ( A ) = ∑ i l i ( A ) . Assume ˆ s 1 , ..., ˆ s 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 √ Jp i if ˆ s 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 = ˜ O ( 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..

4 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).

̸

̸

̸

̸

For each row i = j = i 1 . . . i j -1 i j +1 . . . i N of A = j , Equation (2) gives

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

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 ( NIR 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.

̸

4.1 Efficient Core Chain Leverage Score Sampling

̸

̸

̸

̸

̸

̸

As discussed above, when the TT tensor is in canonical form, the leverage score of row i = j is given by l i = 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 ˆ s 1 ∈ [ I 1 ] , ..., ˆ s j ∈ [ I j ] be random variables such that the multi-index ˆ s ≤ j = ˆ s 1 . . . ˆ s 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 ˆ s 1 . . . ˆ s j is selected with probability proportional to the squared norm of the corresponding row of A ≤ j :

Our 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 ,...,s j ) p ( s k +1 ,...,s j ) . The following lemma shows that this can be done efficiently by leveraging the underlying TT structure.

Lemma 4.1 (Conditional distribution for ˆ s k ) . Consider the events ˆ s j = s j , . . . , ˆ s k +1 = s k +1 , which we abbreviate as ˆ s >k = s >k . Then

where H >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 time O ( 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 I i be the probability vector for the leverage score distribution of A ≤ j . Then Equation (4) can be rewritten in vector form as q := 1 R j ( 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 ˆ t 1 , ..., ˆ t j follow the conditional distributions defined by

where 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 ˆ t k satisfies p ( ˆ t k = t k | ˆ t >k = t >k ) = p (ˆ s k = s k | ˆ s >k = s >k ) .

As a consequence, the joint random variable ( ˆ t 1 , ..., ˆ t j ) 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 that

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 ˆ t k 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 ˆ t k . 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:

Lemma 4.3 (Bharadwaj et al. [2023], Adapted) . Given a matrix A ∈ R I × R , there exists a data structure with construction time O ( IR 2 ) and space usage O ( IR ) such that, given any vector h ∈ R R , a single sample from the un-normalized distribution of weights ( A · h ) 2 can be drawn in time O ( R 2 log( I/R )) .

The 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 matrix A 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 requires 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

Algorithm 1 ConstructChainSampler( A 1 , ..., A N )

1: for k = 1 ..N do

2: Z k := BuildSampler ( A L k )

Algorithm 2 ChainSampleLeft( J, j )

1: for d = 1 ..J do

2: ˆ r := Uniform-sample ([1 ...R j ])

3: h := e ˆ r

4: for k = j... 1 do

5: ˆ t k := RowSample ( Z k , h ) //R k - 1

6:

h

=

h

· A

k

[

:

,

ˆ

t

k

,

:

]

7: t d = ( ˆ t k ) k ≤ j

8: return t 1 , ..., t J

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 ˆ t j -1 , ..., ˆ t 1 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 = ˜ O ( R 2 /εδ ) samples per least-square problem, we have

̸

̸

with probability (1 -δ ) , where ˜ A 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 ˜ O ( # 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 to log ⌈ 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].

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

5.1 Decomposition of Synthetic and Real Dense Datasets

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 set R 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.

100

200

300

400

500

Dimension Size (I)

0.05

0.06

0.07

0.08

0.09

Fit

tt-svd

rtt-svd

tt-als

rtt-als (proposal)

100

200

300

400

500

Dimension Size (I)

0

10

20

30

40

50

60

Time

tt-svd

rtt-svd

tt-als

rtt-als (proposal)

Real Data Experiments. For the real data experiment, we consider four real images and video datasets (more details about datasets are given in Appendix C): (i) Pavia University is a hyper-spectral image dataset of size (610 × 340 × 103) , (ii) DC Mall is also a dataset of hyper-spectral images of size (1280 × 307 × 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).

5.2 Approximate Sparse Tensor Train Decomposition

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

0

1

2

3

4

5

6

Time (s)

0.00

0.02

0.04

0.06

0.08

0.10

0.12

0.14

0.16

0.18

Fit

Uber

Exact

Random LSS

100

200

300

400

Time (s)

0.00

0.01

0.02

0.03

0.04

0.05

0.06

0.07

Fit

Enron

Exact

Random LSS

25

50

75

100

125

150

Time (s)

0.000

0.005

0.010

0.015

0.020

0.025

0.030

Fit

NELL-2

Exact

Random LSS

3.3e4

6.6e4

9.8e4

1.3e5

1.6e5

Samples (J)

0.1500

0.1502

0.1504

0.1506

0.1508

0.1510

Final Fit

Uber

3.3e4

6.6e4

9.8e4

1.3e5

1.6e5

Samples (J)

0.057

0.058

0.059

0.060

0.061

Enron

3.3e4

6.6e4

9.8e4

1.3e5

1.6e5

Samples (J)

0.0258

0.0260

0.0262

0.0264

0.0266

0.0268

Nell-2

Random

Exact

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

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

Acknowledgments and Disclaimers

V. Bharadwaj was supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC0022158. This research used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility, under Contract No. DE-AC0205CH11231 using NERSC award ASCR-ERCAP0024170. BT. Rakhshan was supported by an IVADO Fellowship. O. A. Malik contributed to this work while he was affiliated with Lawrence Berkeley National Laboratory where he was supported by the Office of Science of the DOE under Award Number DE-AC02-05CH11231. G. Rabusseau acknowledges the support of the Canadian Institute for Advanced Research (CIFAR AI chair program) and the Natural Sciences and Engineering Research Council of Canada (Discovery program, RGPIN-2019-05949).

This report was prepared as an account of work sponsored by an agency of the United States Government. Neither the United States Government nor any agency thereof, nor any of their employees, makes any warranty, express or implied, or assumes any legal liability or responsibility for the accuracy, completeness, or usefulness of any information, apparatus, product, or process disclosed, or represents that its use would not infringe privately owned rights. Reference herein to any specific commercial product, process, or service by trade name, trademark, manufacturer, or otherwise does not necessarily constitute or imply its endorsement, recommendation, or favoring by the United States Government or any agency thereof. The views and opinions of authors expressed herein do not necessarily state or reflect those of the United States Government or any agency thereof.

References

Johann A Bengua, Ho N Phien, and Hoang D Tuan. Optimal feature extraction and classification of tensors via matrix product state decomposition. In 2015 IEEE International Congress on Big Data , pages 669-672. IEEE, 2015.

Vivek Bharadwaj, Osman Asif Malik, Riley Murray, Laura Grigori, Aydin Buluc, and James Demmel. Fast exact leverage score sampling from khatri-rao products with applications to tensor decomposition. Advances in Neural Information Processing Systems , 36, 2023.

Maolin Che and Yimin Wei. Randomized algorithms for the approximations of tucker and the tensor train decompositions. Advances in Computational Mathematics , 45(1):395-428, 2019.

Zhongming Chen, Huilin Jiang, Gaohang Yu, and Liqun Qi. Low-rank tensor train decomposition using tensorsketch. arXiv preprint arXiv:2309.08093 , 2023.

Dehua Cheng, Richard Peng, Yan Liu, and Ioakeim Perros. Spals: Fast alternating least squares via implicit leverage scores sampling. Advances in neural information processing systems , 29, 2016.

Andrzej Cichocki and Anh-Huy Phan. Fast local algorithms for large scale nonnegative matrix and tensor factorizations. IEICE transactions on fundamentals of electronics, communications and computer sciences , 92(3):708-721, 2009.

Petros Drineas, Ravi Kannan, and Michael W Mahoney. Fast monte carlo algorithms for matrices i: Approximating matrix multiplication. SIAM Journal on Computing , 36(1):132-157, 2006a.

Petros Drineas, Michael W. Mahoney, and Shan Muthukrishnan. Sampling algorithms for ℓ 2 regression and applications. In Proceedings of the Seventeenth Annual ACM-SIAM Symposium on Discrete Algorithm , pages 1127-1136, 2006b.

Petros Drineas, Michael W. Mahoney, and S. Muthukrishnan. Relative-error CUR matrix decompositions. SIAM Journal on Matrix Analysis and Applications , 30(2):844-881, 2008.

Petros Drineas, Michael W Mahoney, Shan Muthukrishnan, and Tamás Sarlós. Faster least squares approximation. Numerische mathematik , 117(2):219-249, 2011.

Glen Evenbly. Gauge fixing, canonical forms, and optimal truncations in tensor networks with closed loops. Physical Review B , 98(8):085155, 2018.

Glen Evenbly. A practical guide to the numerical implementation of tensor networks i: Contractions, decompositions and gauge freedom. arXiv preprint arXiv:2202.02138 , 2022.

Matthew Fahrbach, Gang Fu, and Mehrdad Ghadiri. Subquadratic kronecker regression with applications to tensor decomposition. Advances in Neural Information Processing Systems , 35: 28776-28789, 2022.

Nathan Halko, Per-Gunnar Martinsson, and Joel A Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review , 53(2): 217-288, 2011.

Christopher J Hillar and Lek-Heng Lim. Most tensor problems are np-hard. Journal of the ACM (JACM) , 60(6):1-39, 2013.

Frank L Hitchcock. The expression of a tensor or a polyadic as a sum of products. Journal of Mathematics and Physics , 6(1-4):164-189, 1927.

Sebastian Holtz, Thorsten Rohwedder, and Reinhold Schneider. The alternating linear scheme for tensor optimization in the tensor train format. SIAM Journal on Scientific Computing , 34(2): A683-A713, 2012.

Benjamin Huber, Reinhold Schneider, and Sebastian Wolf. A randomized tensor train singular value decomposition. In Compressed Sensing and its Applications: Second International MATHEON Conference 2015 , pages 261-290. Springer, 2017.

Pavel Izmailov, Alexander Novikov, and Dmitry Kropotov. Scalable gaussian processes with billions of inducing inputs via tensor train decomposition. In International Conference on Artificial Intelligence and Statistics , pages 726-735. PMLR, 2018.

Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review , 51(3): 455-500, 2009.

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

Brett W Larsen and Tamara G Kolda. Practical leverage-based sampling for low-rank tensor decomposition. SIAM Journal on Matrix Analysis and Applications , 43(3):1488-1517, 2022.

Michael W Mahoney et al. Randomized algorithms for matrices and data. Foundations and Trends® in Machine Learning , 3(2):123-224, 2011.

Osman Asif Malik. More efficient sampling for tensor decomposition with worst-case guarantees. In International Conference on Machine Learning , pages 14887-14917. PMLR, 2022.

Osman Asif Malik and Stephen Becker. A sampling-based method for tensor ring decomposition. In International Conference on Machine Learning , pages 7400-7411. PMLR, 2021.

Osman Asif Malik, Vivek Bharadwaj, and Riley Murray. Sampling-based decomposition algorithms for arbitrary tensor networks, 2022.

Alexander Novikov, Mikhail Trofimov, and Ivan Oseledets. Exponential machines. arXiv preprint arXiv:1605.03795 , 2016.

Ivan V Oseledets. Tensor-train decomposition. SIAM Journal on Scientific Computing , 33(5): 2295-2317, 2011.

Ninh Pham and Rasmus Pagh. Fast and scalable polynomial kernels via explicit feature maps. In Proceedings of the 19th ACM SIGKDD international conference on Knowledge discovery and data mining , pages 239-247, 2013.

Nicholas D Sidiropoulos, Lieven De Lathauwer, Xiao Fu, Kejun Huang, Evangelos E Papalexakis, and Christos Faloutsos. Tensor decomposition for signal processing and machine learning. IEEE Transactions on signal processing , 65(13):3551-3582, 2017.

Shaden Smith, Jee W. Choi, Jiajia Li, Richard Vuduc, Jongsoo Park, Xing Liu, and George Karypis. FROSTT: The Formidable Repository of Open Sparse Tensors and Tools, 2017. URL http: //frostt.io/ .

Edwin Stoudenmire and David J Schwab. Supervised learning with tensor networks. Advances in neural information processing systems , 29, 2016.

Ledyard R Tucker. Some mathematical notes on three-mode factor analysis. Psychometrika , 31(3): 279-311, 1966.

David P Woodruff et al. Sketching as a tool for numerical linear algebra. Foundations and Trends in Theoretical Computer Science , 10(1-2):1-157, 2014.

Gaohang Yu, Jinhong Feng, Zhongming Chen, Xiaohao Cai, and Liqun Qi. A randomized block krylov method for tensor train approximation. arXiv preprint arXiv:2308.01480 , 2023.

Qibin Zhao, Guoxu Zhou, Shengli Xie, Liqing Zhang, and Andrzej Cichocki. Tensor ring decomposition. arXiv preprint arXiv:1606.05535 , 2016.

Zhou, L. Li, and H. Zhu. Tensor regression with applications in neuroimaging data analysis. Journal of the American Statistical Association , 108(502):540-552, 2013.

7 NeurIPS paper checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims are supported by theoretical results and experiments.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer:[NA]

Justification: We did not find any particular limitation of our approach that we thought was relevant to discuss in the paper.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The main theorems and lemmas are in section 4 with proofs in appendix.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In the experiment section we addressed all details about the datasets and the metrics to evaluate our method.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

The GitHub repository link is provided in the experiments section.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All the details all in the experiment section.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We reported all the results for several runs in plots and tables in the experiment section. However, we did not report the error bars as they are too small to report.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We reported the time each algorithms takes in the experiment section. Details are given in appendix.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper focuses on theoretical aspects. We did not feel relevant to discuss the potential positive and negative societal impacts of democratizing access to tensor decomposition methods.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer:[NA] .

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

A Additional Notation

Definition A.1. Let A ∈ R m × n and B ∈ R p × q then the Kronecker product, A ⊗ B ∈ R mp × nq is defined by

Definition A.2. Let A ∈ R m × R and B ∈ R n × R then the Khatri-Rao product, A ⊙ B ∈ R mn × R is defined by

where a 1 , . . . , a R ∈ R m are the columns of A , b 1 , . . . , b R ∈ R n are the columns of B and the columns of A ⊙ B is the subset of the Kronecker product. In the corresponding tensor network diagram, the copy tensor captures the fact that the second indices are the same.

A.1 Details about Orthogonalization of the TT Decomposition

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

B Proofs

B.1 Proof of Lemma 4.1

In the expressions above, the summation over each variable s t , 1 ≤ t ≤ k , is taken over the range [ I t ] . The first step follows by marginalizing over random variables ˆ s 1 , ..., ˆ s k -1 . The second step follows from Equation (4). The third step rewrites an inner product of two vectors as the trace of their outer product. The fourth step follows from the definition of A ≤ j . The fifth step follows from the linearity of the trace by moving the summation over s 1 into the product expression. The last step follows from the definition of the left-orthonormality property on A 1 ; that is, ∑ s 1 A 1 [: , s 1 , :] ⊤ · A [: , s 1 , :] =

A L ⊤ 1 A L 1 = I . By successively moving summation operators into the product expression to repeat the last step (exploiting the left-orthonormality of each core in the process), we find

where the last line follows from the definition of H >k .

B.2 Proof of Lemma 4.2

We write

The first step follows from Lemma 4.1. The second step follows from the definition of the trace. The third step follows from the definitions of h >k and H >k . The fourth step follows from the definition of the variables ˆ t 1 , ..., ˆ t j . Now observe that p (ˆ r = r ) = 1 /R j for 1 ≤ r ≤ R j , so we can write

which completes the proof.

B.3 Efficient Sampling Data Structure

Lemma 4.3 first appeared as Lemma 3.2 in the original work by Bharadwaj et al. [2023]. We state a condensed form of the original claim below:

Lemma B.1 (Bharadwaj et al. [2023], Original) . Given U ∈ R M × R , Y ∈ R R × R with Y p.s.d., there exists a data structure parameterized by positive integer F that requires O ( MR 2 ) time to construct and additional space space O ( R 2 ⌈ M/F ⌉ ) . After construction, the data structure can draw a sample from the distribution defined elementwise by

in time O ( R 2 log ⌈ M/F ⌉ + FR 2 ) . When Y is a rank-1 matrix, the runtime drops to O ( R 2 log ⌈ M/F ⌉ + FR ) .

In the statement above, C is an appropriate normalization constant. To prove our adapted lemma, take Y = [1] , a matrix of all ones that is rank-1, and set F = R . Then

This is the target probability distribution of Lemma 4.3, and the runtime to draw each sample is O ( R 2 log( M/R )+ R 2 ) = O ( R 2 log( M/R )) . The choice F = R also induces space usage O ( MR ) , linear in the size of the input. Our modified claim follows.

B.4 Proof of Theorem 1.1

We provide a short end-to-end proof that shows that Algorithms 1 and 2 correctly draw samples from A ≤ j (the matricization of the left-orthogonal core chain) according to the distribution of its squared row norms while meeting the runtime and space guarantees of Theorem 1.1.

Construction Complexity : The cost of Algorithm 1 follows from 4.3 with M = IR k -1 , the row count of A L k for 1 ≤ k ≤ j . Using this lemma, construction of each sampling data structure Z k requires time O ( I k R k -1 R 2 k ) . The space required by sampler Z k is O ( I k R k -1 R k ) ; summing over all indices k gives the construction claim in Theorem 1.1.

Sampling Complexity: The complexity to draw samples in Algorithm 2 is dominated by calls to the RowSample procedure, which as discussed in Section 4 is O ( R 2 k log( I k R k -1 /R k )) Summing the complexity over indices 1 ≤ k < j yields the cost claimed by Theorem 1.1 to draw a single sample. The complexity of calling the RowSample procedure repeatedly dominates the complexity to update the history vector h over all loop iterations, which is O ( ∑ j k =1 R k -1 R k ) for each sample.

Correctness: Our task is to show that Algorithm 2 each sample t d , 1 ≤ d ≤ J , is a multi-index that follows the squared row norm distribution on the rows of A ≤ j . To do this, we rely on lemmas proven earlier. For each sample, the variable ˆ r is a uniform random draw from [ R j ] , and h is initialized to the corresponding basis vector. By Equation (6) and Lemma 4.3, Line 5 from Algorithm 2 draws each index ˆ t k correctly according to the probability distribution specified by Equation (5). The history vector is updated by Line 6 of the algorithm so that subsequent draws past iteration k of the loop are also drawn correctly according to Equation (5). Lemma 4.2 (relying on Lemma 4.1) shows that the multi-index ˆ t 1 . . . ˆ t j drawn according to Equation (5) follows the same distribution as ˆ s 1 . . . ˆ s j , which was defined to follow the squared norm distribution on the rows of A ≤ j . This completes the proof.

B.5 Proof of Corollary 4.4

̸

Since A = j ∈ R ∏ N k = j I k × R j -1 R j and X ( j ) ∈ R ∏ N k = j I k × I j , we draw ˜ O ( R 2 /εδ ) samples to achieve the error bound (1 + ε ) with probability (1 -δ ) for each least squares solve in the down-sampled problem (3.5). By Theorem 1.1, the complexity of drawing J samples with our data structure is

̸

̸

̸

̸

where we suppose that R 1 = R 2 = · · · = R N -1 and I 1 = · · · = I N . The cost of sampling a corresponding subset of X ( j ) is O ( JI j ) = ˜ O ( R 2 / ( εδ ) I j ) . Solving the downsampled least squares problem also costs O ( JR 2 I j ) = ˜ O ( I j R 4 / ( εδ ) ) . Summing them all together for 1 ≤ j ≤ N gives

̸

where we wrote the last equation considering the fact that N dominates ( N -1) .

C Details about Datasets & Experiments

C.1 Datasets

For the real dense datasets experiment, we truncated and reshaped the original data tensors in to the fourth order tensors as follows.

Pavia University dataset : The original has dimensions (610 , 340 , 103) . We truncate it to (600 , 320 , 100) , permute the modes to dimensions (100 , 320 , 600) tensor and reshape it into a tensor of dimensions (100 , 320 , 24 , 25) . It is available at

http://lesun.weebly.com/hyperspectral-data-set.html

Tabby Cat dataset is permuted to (286 , 720 , 1280) and reshaped to a tensor of size (286 , 720 , 40 , 32) . The video is in color and converted to grayscale by averaging the three color channels. It is available at

https://www.pexels.com/video/video-of-a-tabby-cat-854982/ .

The MNIST dataset was reshaped into a tensor of size (280 , 600 , 28 , 10) and is available at https://www.kaggle.com/datasets/hojjatk/mnist-dataset

The Washington DC Mall dataset was truncated to dimensions (1280 , 306 , 190) before reshaping into a tensor of size (1280 , 306 , 10 , 19) . It is available at https://engineering.purdue.edu/Ëœbiehl/MultiSpec/hyperspectral.html.

The sparse tensors Uber , Enron , and NELL-2 were downloaded from the FROSTT collection [Smith et al., 2017]. The dimensions of these tensors were unchanged from the versions available online. Consistent with established practice [Larsen and Kolda, 2022], we computed the logarithm of the tensor values in the Enron and NELL-2 datasets before performing our experiments.

C.2 Computing Resources

The dense data experiments were conducted on MILA cluster nodes with 4 CPUs and 16GB of RAM each. Sparse tensor decomposition experiments were conducted on NERSC Perlmutter nodes with 2 CPUs and 512 GB of RAM each.

Peer reviews:
comments: 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.

STRENGTHS:
The approach taken is natural, and the bounds obtained are quite powerful.

The experiments considered both synthetic and real data sets, and demonstrated clear gains in the parameter regimes considered.

WEAKNESSES:
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:
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:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-07

comments: 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.

STRENGTHS:
The paper is well written with good theoretical analysis and desired experimental performance. The method using leveraging score sampling is technically sound for TT decomposition.

WEAKNESSES:
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:
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:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- Empirical evaluation show a runtime improvement with respect to unsketched TT-ALS, without reducing fit.\n- The paper presents an efficient algorithm for leverage score sampling of ALS subproblems in TT-ALS.\n\nWEAKNESSES:\nMajor comments (affecting recommendation):\n- 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.\n- 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.\n- 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.\n- Many details required for implementation of the algorithm are omitted. Many points regarding the algorithm are not clear. \n\nMinor comments (do not affect the recommendation):\n- Lines 125 - 126: add citation.\n- Figure 3 left: is this fit in the y-axis or misfit?\n- Line 108-109: I think you mean *omission* of multiplicative terms.\n- Lines 58-60: shouldn't runtime be O(jIR^3)?\n- Line 54: The term \"left matricization\" is used without defining it yet.\n- Line 55: I think it should be R_j\n- Eq (2): I think on the right it should be A[i,:]^T. Similar issue with Eq (3).\n\nQUESTIONS:\n- 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?\n- Lines 64-66: What is the significance of this observation regarding the complexity of the algorithm?\n- Lines 125 - 126: One the canonical form has been computed, can we solve the ALS subproblem efficiently?\n- Did you compare empirically your method to the method of [Chen et al. 2023].\n- Is your algorithm computing the canonical form anew in each ALS iteration, or are you somehow updating it?\n- 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?\n\nLIMITATIONS:\nNothing to add.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe strengths of the paper  are:\n1. A novel randomized approach is presented for computing the TT decomposition of large tensors.\n2. New data structure and sampling approach are proposed. Theoretical results are presented.\n3. Numerical results show that the proposed method is efficient.\n\nWEAKNESSES:\nThe weaknesses of the paper are:\n1. The randomized SVD approach seems to be more efficient in terms cost compared the the prosed method.\n2. Certain aspects of the paper can be improved (see below).\n\nQUESTIONS:\nThe paper presents an interesting approach to compute the TT decomposition of large tensors. \nHowever, I have the following comments that can further improve the paper:\n\n1. r-TT-SVD seems better - \n\nBased 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?\n\nIn 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. \n\n\n2. Presentation:\nCertain aspects of the presentation can be improved, and there are few minor typos.\n\ni. 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).\n\nii. 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. \n\niii. The computational  cost of solving eq (1) can be added.\n\niv. Theorem 1.1, point 1 has a $j$ missing in the runtime cost. \n\nv. Not sure what is the subscript (2) in Corollary 4.4.\n\nMany 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.\n\nLIMITATIONS:\nSee above comments\n\n\n------\nPost Rebuttal:\n\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}