Efficient Leverage Score Sampling for Tensor Train Decomposition

We thank the reviewer for their comments and feedback, please find responses below:

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.

Correct. We added dense synthetic and real data experiments to show that our proposed approach has a better time complexity than TT-ALS and TT-SVD and works slightly better than rTT-SVD in terms of fit. The main point is to demonstrate that our new approach works as well as other types of TT decompositions for dense data. As SVD-based methods cannot handle high-dimensional sparse tensors. For the more complicated tensors, we added a sparse tensor data experiment as you can see in the second part, the setting is different and more complicated as SVD-based TT decomposition cannot be used and we just compared our approach with the classical TT-ALS method.

Questions:

Would it be possible to check how the orthogonality conditions are maintained in intermediate steps under the inexact arithmetics caused by round off errors? Aka. are the conditions for the characterizations of leverage scores preserved exactly?

We did not observe any numerical instabilities due to a round of errors during the orthogonalization step (which is performed using the stable QR decomposition routine provided by numpy.linalg) and we believe the characterization of leverage scores is preserved exactly (up to machine precision).

We thank the reviewer for their comments and feedback, please find responses below:

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.

In [Malik et al, 2021], the leverage scores were approximated using a product of simpler distributions. They sample according to the leverage scores of each core. While in our paper, we proposed a novel data structure that can compute the exact leverage scores. Moreover, the time complexity to solve one least-square of ALS with the method of Malik et al. Malik et al approach is $NIR^4$ +#iter. $NIR^{(2N+2)}/(\epsilon\delta)$, which is still exponential in the order of the tensor $N$, while our result does not have any exponential dependency to the order of the tensor. We will add and emphasize in the camera-ready version that our approach differs from theirs in two main aspects:

• We did not use approximation for finding the leverage scores, instead a novel data structure is proposed for finding exact leverage scores.
• In Malik et al, the runtime for the least-square solve still has an exponential dependency to the order of a tensor while our approach is free of this exponential dependency.

Questions

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

Indeed, TT is a special case of the TR decomposition. However, in [Malik and Becker, 2021], leverage score approximation is used for the TR-ALS with time complexities #iter · $NIR^{2N+2}/(\epsilon\delta)$ (further improved to and #iter · $N^3R^8 (R + I/\epsilon)/\delta$ in [Malik, 2022]), respectively. In contrast, our method has a lower time complexity and provides a better theoretical guarantee (please see corollary 4.4 in our paper). Therefore, our method goes beyond the works of Malik et al. As mentioned in their works, computing leverage scores requires computing the left singular values which has the same cost as solving the original least squares problem. For this reason, they approximate the leverage scores. Instead, we propose a novel data structure to sample from the exact leverage scores at a lower cost than their proposed approaches.

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.

It is true that TT is a special case of TR, but note that the main problem we address in our paper is how to scale TT-ALS to very large tensors using randomized techniques. The purpose of our experiments is to demonstrate how we efficiently achieve this goal with exact leverage score sampling. Extending our approach to TR-ALS could be interesting, however it would be very challenging due to the lack of canonical forms for TR decomposition.

We hope we have addressed all your concerns and answered your questions (in which case we kindly ask you to consider increasing your score). We are happy to clarify any additional point during the discussion period.

We thank the reviewer for their comments and respond to the feedback below:

Weaknesses:

[W1][Comparison against (Chen et al. 2023)]

The TensorSketch approach by Chen et. al. requires an exponential sample count in the tensor dimension q to achieve the $(\epsilon, \delta)$-guarantee on the accuracy of each least-squares solve. See Theorem 4.4 of their work and note the 3^q term in the sample complexity, where (in their notation) q is the tensor dimension. By contrast, our algorithm requires a sample count that has no dependence on the tensor dimension q; the sample count depends only on the column count of the design matrix and the parameters $\epsilon$nd $\delta$. The time to draw these samples in our method is linearly proportional to q, the tensoring dimension, a modest cost. The efficiency of our method stems from the fact that we sample from the exact distribution of squared row norms of the design matrix. By contrast, TensorSketch applied without tree structure modifications, such as those proposed by Ahle et al., suffers from an exponential output sketch size. We will clarify this point at the end of Section 2 and in the Introduction.

[W2][Runtime improvement]

The runtime improvements are most significant for large sparse tensors. Figure 4 shows the accuracy (y-axis, higher is better) vs. ALS iteration time for our method vs. non-randomized ALS. The speedup per iteration can be as high as 26x for lower ranks. 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 will note this at the beginning of Section 5 (Experiments) and will highlight these facts further in Section 5.2.

[W3][End-to-end analysis]

Correct. Our analysis provides guarantees on the accuracy of individual least-squares problems, in line by prior works by Cheng et. al, Larsen and Kolda, and Malik. Under additional assumptions on the tensor, convergence guarantees can be derived for sketched ALS: see, for example, https://proceedings.mlr.press/v119/gittens20a/gittens20a.pdf. The global guarantees derived in this paper by Aggour, Gittens, and Yener depend on sketching guarantees established for each linear least-squares problem, which we provide in our work.

[W4][Implementation details]

The revised version of our draft will include (in section 4.1) detailed descriptions of the BuildSampler and RowSample functions, the data structure described in the appendix to draw samples from each orthonormal core flattening. Our Git repository, which is public at https://anonymous.4open.science/r/fast_tensor_leverage-EB01, includes a simple, slow reference implementation written in Python for this data structure, providing sufficient detail to understand and fully replicate our method.

Minor comments

Thanks for catching these typos, we will correct and address them all. For Figure 3, the y-axis is the fit.

Questions:

[Q1] [Our approach vs TensorSketch]

We believe we addressed this point in the comment above on the exponential complexity of TensorSketch without appropriate modifications. Our method offers sub-exponential sample complexity and worst-case runtime complexity to achieve guarantees on the solution to each least-squares problem. The approach by Chen. et. al. requires a worst-case sketch size exponential in the tensor dimension.

[Q2][Complexity of the algorithm]

The significance here is that data structure construction (as well as subsequent updates to the data structure) does not increase the asymptotic complexity of ALS tensor train decomposition.

[Q3][proceeding only with the canonical form]

Even after computing the canonical form (which makes the design matrix orthonormal for each linear least squares problem orthonormal), we must multiply the matricized tensor against the chain of TT cores placed in canonical form. This is the major computational bottleneck - without sketching, the runtime cost for this matrix multiplication scales as $O(nnz(T) N R^2)$, where $nnz(T)$ is the number of nonzeros in the tensor, $N$ is the order of the tensor, and $R$ is the tensor train rank. The tensor may have hundreds of millions of nonzero values. 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)$ * (number of selected nonzeros extracted by the sketch). We can update our draft to clarify this point in Section 1.1.

[Q4][Emprical comparison with Chen et al. 2023]

We did not compare empirically with Chen et. al results. To the best of our knowledge, there is no public code available to replicate their results and compare with them. Also, as explained previously, from the theoretical point of view, the method proposed by Chen et. al cannot scale to the size of the tensors that we consider in the large sparse tensors experiments.

[Q5][Updating canonical form in each iteration]

We update the canonical form without recomputing it entirely at each substep of ALS. After solving for the optimal value for each core, we compute the QR decomposition of its left or right flattening and replace the newly-computed core with the appropriately-reshaped Q factor. In this way, only one core needs to be updated after each least-squares solve to maintain the canonical form.

[Q6][Computing the leverage scores]

We do not compute the leverage scores of all rows, since, as you pointed out, there are exponentially-many rows in the tensor dimension. To design a computationally-efficient sampling algorithm, we build an efficient data structure to sample from this distribution without materializing it.

We hope we have addressed all your concerns and answered your questions (in which case we kindly ask you to consider increasing your score). We are happy to clarify any additional point during the discussion period.

We thank the reviewer for their comments and feedback, please find responses below:

Weaknesses:

The randomized SVD approach seems to be more efficient in terms cost compared the the prosed method.

Randomized SVD cannot scale to large tensors. Indeed, to decompose a tensor using randomized SVD at each step we need to produce a Gaussian tensor of size $s_j\times I^{(N-1)}$, where N is the order of the tensor, $I$ is the physical dimension of each mode and $s_j$ is the size of the sketch. This approach works well when $N$ is small. But randomized svd cannot handle very high-order tensors. In the experiment part, we only use tt-svd and randomized tt-svd for the dense tensors with small $N$, but these two algorithms cannot scale to the case of large sparse tensors in the second part of the experiments.

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

yes : TT-ALS is a very popular algorithm, notably in the quantum physics community, where it is closely related to the density renormalization group (DMRG) algorithm.

How do these methods differ in the quality of the decompositions computed?

Regarding the quality of the decomposition, TT-ALS usually finds better solutions than TT-SVD (TT-SVD can actually be used as an initialization for TT-ALS). Note however that the focus of our paper is on improving the TT-ALS algorithm using randomization. Indeed, TT-ALS is a popular algorithm, but, as mentioned previously, the cost of solving the least-squares problems in ALS is exponential in the order of the tensor which is not efficient for decomposing very large tensors. The main goal of our paper is to make the ALS efficient with randomization and doable for high-order tensors. Therefore, the improvement of a TT-ALS algorithm is the main focus of our paper.

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.

That’s correct, by increasing the sketch size we can have a better quality of sketch for rTT-SVD. However, as mentioned earlier, the SVD-based approach is not suitable for decomposing high-order tensors. Even if by increasing the sketch size rTT-SVD still suffers from the curse of dimensionality when N is very large. Therefore, the quality of the sketch of rTT-SVD cannot be compared to rTT-ALS in the case of sparse high-order tensors (section 5.2 second part of the experiments.)

Details about the sketch sizes used are missing.

In Section 5.1, the sketch sizes are set to $J=5000$ and $J=2000$ for the synthetic and real data, respectively (we mention it in the text but we will add this information in the caption of Figure 3 and Table 1 as well).

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

TT decomposition enjoys stable and numerical stabilities while finding rank-r CP decomposition is NP-hard and the number of parameters of a Tucker decomposition grows exponentially with the order of the tensor. In contrast, the number of parameters is linear in the order of the tensor for both TT and TR decompositions, but TR decomposition is known to suffer numerical stability issues. For these reasons, TT is more popular in tensors and quantum physics communities (where it is known as Matrix Product States / MPS). To address your concern about motivating the TT format, we will include a more comprehensive motivation in the introduction.

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.

If we understand correctly, your concern is mostly about the main theorem in the introduction which may cause difficulties for other audiences. We will clarify this by clearly stating before the theorem that all relevant definitions are given in Section 3.1.

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

Solving Eq (1) exactly (without randomization) has a computational cost of $O(I^N)$. It is mentioned in the introduction but we will emphasize it in the revision in section 3.2.

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

Correct. We will add j in the final version. Thank you for catching this typo.

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

This denotes the second mode unfolding of the tensor which is defined in Definition 3.1 in the paper.

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.

The paper “Fast exact leverage score sampling from Khatri-rao products with applications to tensor decomposition.” by Bharadwaj et el. has been published in Neurips 2023, we will correct the citation in the camera-ready version.