Adaptive Randomized Tensor Train Rounding using Khatri-Rao Products

TL;DR

The paper tackles efficient TT-rounding by introducing adaptive randomized algorithms that use Khatri-Rao product sketches to control the approximation error. It presents fixed-rank and adaptive-rank variants, provides a rigorous norm-estimation analysis with concentration results, and demonstrates up to 50x speedups over deterministic TT-rounding across synthetic, kernel-parameter, and PDE-driven problems. A key advantage is adaptivity to a prescribed tolerance, enabling efficient compression even for sums of TT tensors without forming the full sum. The authors also release code and provide extensive appendices detailing algorithm variants and problem setups, underscoring practical impact for large-scale tensor computations.

Abstract

Approximating a tensor in the tensor train (TT) format has many important applications in scientific computing. Rounding a TT tensor involves further compressing a tensor that is already in the TT format. This paper proposes new randomized algorithms for TT-rounding that uses sketches based on Khatri-Rao products (KRP). When the TT-ranks are known in advance, the proposed methods are comparable in cost to the sketches that used a sketching matrix in the TT-format~\cite{al2023randomized}. However, the use of KRP sketches enables adaptive algorithms to round the tensor in the TT-format within a fixed user-specified tolerance. An important component of the adaptivity is the estimation of error using KRP sketching, for which we develop theoretical guarantees. We report numerical experiments on synthetic tensors, parametric low-rank kernel approximations, and the solution of parametric partial differential equations. The numerical experiments show that we obtain speed-ups of up to $50\times$ compared to deterministic TT-rounding. Both the computational cost analysis and numerical experiments verify that the adaptive algorithms are competitive with the fixed rank algorithms, suggesting the adaptivity introduces only a low overhead.

Figures (8)

• Figure 1: Tensor network diagram for a TT-tensor $\boldsymbol{{\mathscr{X}}} = [ \boldsymbol{{\mathscr{X}}} _1, \ldots, \boldsymbol{{\mathscr{X}}} _d ]$.
• Figure 1: An illustration of partial contraction of tensor $\boldsymbol{{\mathscr{X}}} = [ \boldsymbol{{\mathscr{X}}} _1, \boldsymbol{{\mathscr{X}}} _2, \boldsymbol{{\mathscr{X}}} _3, \boldsymbol{{\mathscr{X}}} _4 ]$ with matrix $\bm{{\mathbf{\Omega}}} {^{\intercal}}{ {^{\newline}}} = \bm{{\mathbf{\Omega}}} {^{\intercal}}{ {^{\newline}}} _4 \odot \bm{{\mathbf{\Omega}}} {^{\intercal}}{ {^{\newline}}} _3 \odot \bm{{\mathbf{\Omega}}} {^{\intercal}}{ {^{\newline}}} _2$.
• Figure 1: Left: Relative error over 10 runs on synthetic data. Right: Average speedup with respect to TT-Rounding algorithms over 10 runs.
• Figure 2: Average estimated norm over 1000 runs for synthetic tensors with different numbers of modes ($d$). The solid lines represent the average, and the shaded region indicates the minimum and maximum values observed across the runs.
• Figure 2: Left column: Comparison of baseline methods on the Matérn kernel data. Right column: Comparison after applying a rounding pass to the left-orthonormal TT tensor obtained from adaptive randomized TT rounding algorithms. Relative errors obtained over 10 independent runs. Average speedup is measured relative to the deterministic TT-Rounding algorithm. Compression is defined as the ratio of the number of parameters in the original tensor to that in the compressed tensor and is averaged over 10 runs. Here, Adap refers to adaptive-rank algorithms, Fix to their fixed-rank counterparts, and Adap-R to the adaptive-rank algorithm followed by a rounding pass on the left-orthonormal TT tensor.

Theorems & Definitions (7)

• Remark 3.1: Input tensor norm estimation
• Remark 3.2: Additional Compression
• Remark 3.3: Sum of TT-tensors
• Theorem 3.4
• Proof 1: Proof of Theorem \ref{['thm:frobenius-preservation']}
• Corollary 3.5
• Proof 2