Efficient Leverage Score Sampling for Tensor Train Decomposition

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. * Equal contribution 38th Conference on Neural Information Processing Systems (NeurIPS 2024). 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.