Performance of linear solvers in tensor-train format on current multicore architectures

Abstract

Tensor networks are a class of algorithms aimed at reducing the computational complexity of high-dimensional problems. They are used in an increasing number of applications, from quantum simulations to machine learning. Exploiting data parallelism in these algorithms is key to using modern hardware. However, there are several ways to map required tensor operations onto linear algebra routines (“building blocks”). Optimizing this mapping impacts the numerical behavior, so computational and numerical aspects must be considered hand-in-hand. In this paper we discuss the performance of solvers for low-rank linear systems in the tensor-train format (also known as matrix-product states). We consider three popular algorithms: TT-GMRES, MALS, and AMEn. We illustrate their computational complexity based on the example of discretizing a simple high-dimensional PDE in, for example, 50¹⁰ grid points. This shows that the projection to smaller sub-problems for MALS and AMEn reduces the number of floating-point operations by orders of magnitude. We suggest optimizations regarding orthogonalization steps, singular value decompositions, and tensor contractions. In addition, we propose a generic preconditioner based on a TT-rank-1 approximation of the linear operator. Overall, we obtain roughly a 5× speedup over the reference algorithm for the fastest method (AMEn) on a current multicore CPU.

1. Introduction

Low-rank tensor methods provide a way to approximately solve problems that would otherwise require huge amounts of memory and computing time. Many ideas in this field arise from quantum physics. For example, the global state of a quantum system with N two-state particles can be expressed as a tensor of dimension 2 ^N . In this setting, interesting states are for example given by the eigenvectors of the smallest eigenvalues of the Hamiltonian of the system–a Hermitian linear operator that describes the energy of the system. Therefore, most work focuses on solving eigenvalue problems for Hermitian/symmetric operators using the DMRG method (Schollwöck, 2005; White, 1992; Wilson, 1983). However, linear systems in low-rank tensor formats also arise in interesting applications for example for solving high-dimensional or parameterized partial differential equations, see, for example, Kressner and Tobler (2011); Dahmen et al. (2015); Dolgov and Pearson (2019). In addition, linear solvers and eigenvalue solvers are closely related and many successful methods for finding eigenvalues are based on successive linear solves. This paper addresses iterative methods for solving linear systems (symmetric and non-symmetric) in the tensor-train (TT) format for the case where the individual dimensions are not tiny, that is, for systems of dimension n ^d × n ^d with n ≫ 2. We employ the TT format (called matrix-product states in physics) as it is a simple and common low-rank tensor format. Most of the ideas are transferable to other low-rank tensor formats (at least to loop-free tensor networks). Our work considers the TT-GMRES algorithm (Ballani and Grasedyck, 2012; Dolgov, 2013), the modified alternating least-squares (MALS) algorithm (DMRG for linear systems) (Holtz et al., 2012; Oseledets and Dolgov, 2012), and the Alternating Minimal Energy (AMEn) algorithm (Dolgov and Savostyanov, 2014). We show numerical improvements and performance improvements of the underlying operations and focus on a single CPU node. These improvements are orthogonal to the parallelization for distributed memory systems presented in Daas et al. (2022), so the suggestions from both Daas et al. (2022) and this paper could be combined in the future. As the resulting complete linear solver requires a tight interplay of different algorithmic components, we discuss the behavior of the different numerical methods involved for the TT format. An alternative class of methods to solve linear systems in TT format consists of Riemannian optimization on the manifold of fixed-rank tensor-trains, see Kressner et al. (2016). This results in a nonlinear optimization problem and is therefore not in the scope of this paper. However, it partly requires similar underlying operations.

The paper is organized as follows: First, we start with required numerical background in Section 2 and also introduce relevant performance metrics for today’s multicore computers. Then in Section 3, we discuss the involved high-level numerical algorithms TT-GMRES, MALS, and AMEn. Based on an example, we illustrate the numerical behavior of these algorithms and compare their complexity in Section 4. Afterwards, in Section 5, we analyze and optimize the underlying building blocks. We conclude in Section 6 with a short summary and open points for future work.

2. Background and notation

In this section, we provide the required background concerning numerics and performance.

2.1. Numerical background

We first introduce required matrix decompositions and a notation for the considered algorithms.

2.1.1. Matrix decompositions

As matrix decompositions are heavily used as steps in tensor-train algorithms, we repeat some basic properties of QR and SVD decompositions from the literature, see for example, Golub and Van Loan (2013); Higham (2002).

A QR-decomposition is a factorization of a matrix $M\in {\mathbf{R}}^{n_1\times n_2}$ with column rank r into a matrix Q with orthonormal columns and an upper triangular part R:

$M=QR,\text{with}{Q}^{T}Q=I,Q\in {\mathbf{R}}^{n_1\times r},R\in {\mathbf{R}}^{r\times n_2}.$

(1)

For (numerically) rank-deficient M, one can employ a pivoted QR-decomposition

$MP=QR,\iff M=QR{P}^{-1},$

(2)

where P is a permutation matrix. The pivoted QR-decomposition can be computed in a numerically robust way. However, it cannot (safely) be used to approximate M with a lower rank matrix based on the size of the pivot elements (diagonal entries of R) by Q_:,1:r′R_1:r′,:, r′ < r as the worst case error grows with O (2 ^r ), see Higham (1990); Kawamura and Suda (2021).

The singular value decomposition (SVD) in contrast provides the best approximation of lower rank:

$M=US{V}^T\Rightarrow \Vert M-U_{:,1:r^{\prime }}{S}_{1:r^{\prime },1:r^{\prime }}{V}_{:,1:r^{\prime }}^T{\Vert }_F=\underset{\tilde{M},\mathrm{rank}(\tilde{M})=r^{\prime }}{\min }\Vert M-\tilde{M}{\Vert }_F,$

(3)

where $U\in {\mathbf{R}}^{n_1\times r}$ , $V\in {\mathbf{R}}^{n_2\times r}$ are the matrices of the orthonormal left/right singular vectors and S = diag (s₁, s₂, …, s _r ) is composed of the singular values s₁ ≥ s₂ ≥⋯ ≥ s _r > 0.

2.1.2. Tensor-train decomposition

In higher dimensions, there is no unique way to decompose a tensor into factors and to define its rank(s) (variants are e.g., the Tucker and the CANDECOMP/PARAFAC (CP) decompositions, see the review Kolda and Bader, 2009). We focus on the tensor-train format which decomposes a tensor $X\in {\mathbf{R}}^{n_1\times n_2\times \cdots \times n_d}$ into d three-dimensional sub-tensors X₁, …, X _d :

$X=X_1\bowtie X_2\bowtie \cdots \bowtie X_d\text{with}{X}_k\in {\mathbf{R}}^{r_{k-1}\times n_k\times r_k},r_0=r_d=1.$

(4)

Here, (⋅ ⋈ ⋅) is the contraction of the last dim. of the left operand with the first dim. of the right operand:

$X_k\bowtie X_{k+1}=\sum _i{(X_k)}_{:,:,i}{(X_{k+1})}_{i,:,:}\in {\mathbf{R}}^{r_{k-1}\times n_k\times n_{k+1}\times r_{k+1}}.$

(5)

The TT decomposition of a given tensor is not unique: it is invariant with respect to multiplying one sub-tensor by a matrix and the next with its inverse. More precisely, for $M,N^T\in {\mathbf{R}}^{r_k\times r_k^{\prime }}$ with MN = I:

${\overline{X}}_k\bowtie {\overline{X}}_{k+1}=X_k\bowtie X_{k+1}\text{for}{\overline{X}}_k≔X_k\bowtie M,{\overline{X}}_{k+1}≔N\bowtie X_{k+1}.$

(6)

The smallest possible dimensions (r₁, …r_d−1) that allow to represent X denote the TT ranks of X with the maximal rank r≔rank (X) = max (r₁, …, r_d−1). If X has rank-1 in the TT format, we can also write it as a generalized dyadic product of a set of vectors:

$X={(X_1)}_{1,:,1}\otimes {(X_2)}_{1,:,1}\otimes \cdots \otimes {(X_d)}_{1,:,1}.$

2.1.3. Tensor unfolding and orthogonalities

We define a general reshape operation to reinterpret the entries of a tensor as a tensor of different dimensions:

$\begin{array}{ll}\hfill \mathrm{reshape}\left(X,\left(\begin{array}{ccc}\hfill {\overline{n}}_1\hfill & \hfill \cdots \hfill & \hfill {\overline{n}}_{\overline{d}}\hfill \end{array}\right)\right)& ≔\overline{X}\in {\mathbf{R}}^{{\overline{n}}_1\times \cdots \times {\overline{n}}_{\overline{d}}}\text{with}\hfill \\ \hfill {(\overline{X})}_{{\overline{i}}_1,\dots ,{\overline{i}}_{\overline{d}}}={(X)}_{i_1,\dots ,i_d}& \text{for}{\overline{i}}_1+{\overline{i}}_2{\overline{n}}_1+{\overline{i}}_3{\overline{n}}_1{\overline{n}}_2+\cdots =i_1+i_2{n}_1+i_3{n}_1{n}_2+\cdots .\hfill \end{array}$

With this, we define the left-unfolding that combines two dimensions of a 3d tensor X _k to obtain a matrix:

${(X_k)}_{\text{left}}≔\mathrm{reshape}\left(X_k,\left(\begin{array}{cc}\hfill r_{k-1}{n}_k\hfill & \hfill r_k\hfill \end{array}\right)\right)\in {\mathbf{R}}^{r_{k-1}{n}_k\times r_k}.$

(7)

Similarly, we define the right-unfolding:

${(X_k)}_{\text{right}}≔\mathrm{reshape}\left(X_k,\left(\begin{array}{cc}\hfill r_{k-1}\hfill & \hfill n_k{r}_k\hfill \end{array}\right)\right)\in {\mathbf{R}}^{r_{k-1}\times n_k{r}_k}.$

(8)

We denote a 3d tensor X _k as left-orthogonal if the columns of the left-unfolding are orthonormal $({(X_k)}_{\text{left}}^T{(X_k)}_{\text{left}}=I)$ and right-orthogonal if the rows of the right-unfolding are orthonormal $({(X_k)}_{\text{right}}{(X_k)}_{\text{right}}^T=I)$ . From the TT format, we can build an SVD of an unfolding of $X\in {\mathbf{R}}^{n_1\times n_2\cdots \times n_d}$ into a set of left and right dimensions

$US{V}^T={(X)}_{{\text{unfold}}_j}≔\mathrm{reshape}\left(X,\left(\begin{array}{cc}\hfill n_1\cdots n_j\hfill & \hfill n_{j+1}\cdots n_d\hfill \end{array}\right)\right)$

(9)

by left-orthogonalizing (X₁, …, X _j ) and right-orthogonalizing (X_j+1, …, X _d ):

$X=\underset{U}{\underbrace{{\overline{X}}_1\bowtie \cdots \bowtie {\overline{X}}_j}}\bowtie S\bowtie \underset{V^T}{\underbrace{{\overline{X}}_{j+1}\bowtie \cdots \bowtie {\overline{X}}_d}}.$

(10)

2.1.4. Tensor-train vectors and operators

A tensor-train operator is a tensor in TT format where we combine pairs of dimensions (n _i × m _i ) with the form $\mathcal{A}\in {\mathbf{R}}^{(n_1\times m_1)\times (n_2\times m_2)\times \cdots \times (n_d\times m_d)}$

$\mathcal{A}=A_1\bowtie A_2\bowtie \cdots A_d\text{with}{A}_k\in {\mathbf{R}}^{r_{k-1}^A\times n_k\times m_k\times r_k^A},r_0^A=r_d^A=1$

(11)

which defines a linear mapping: $\mathcal{A}:{\mathbf{R}}^{m_1\times m_2\times \cdots \times m_d}\to {\mathbf{R}}^{n_1\times n_2\times \cdots \times n_d}$ . For the scope of this paper, we only consider quadratic regular operators (m _k = n _k ). Again, the sub-tensor dimensions $(r_1^A,\dots r_{d-1}^A)$ denote the TT ranks of the operator with a maximal rank of $r^A≔\mathrm{rank}(\mathcal{A})=\max (r_1^A,\dots r_{d-1}^A)$ . This definition allows an efficient application of a TT operator on a TT “vector” directly in the TT format:

$\begin{array}{ll}\hfill \mathcal{A}X& =Y=Y_1\bowtie Y_2\bowtie \cdots \bowtie Y_n\text{with}\hfill \\ \hfill Y_k& ≔\mathrm{reshape}\left(\sum _i{(A_k)}_{:,:,i,:}{(X_k)}_{:,i,:},\left(\begin{array}{ccc}\hfill r_{k-1}^A{r}_{k-1}\hfill & \hfill n_k\hfill & \hfill r_k^A{r}_k\hfill \end{array}\right)\right).\hfill \end{array}$

(12)

The resulting TT decomposition has ranks $(r_1^A{r}_1,\dots ,r_{d-1}^A{r}_{d-1})$ and thus rank (Y) ≤ r ^A r.

With all these definitions at hand, we can specify the main problem considered in this paper: Given a low-rank linear operator in TT format ${\mathcal{A}}_{\text{TT}}≔(A_1,\dots A_d)$ , a low-rank right-hand side (RHS) B_TT ≔ (B₁, …, B _d ), and a desired residual tolerance ϵ_tol, find an approximate low-rank solution X_TT ≔ (X₁, …, X _d ) with

${‖{\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-B_{\text{TT}}‖}_F\le {ϵ}_{\text{tol}}.$

(13)

2.2. Performance characteristics on today’s multicore CPU systems

Today’s hardware features multiple levels of parallelism and memory that we need to exploit to efficiently use the available compute capacity: A supercomputer is composed of a number of nodes connected by a network (distributed memory parallelism). Each node contains one or more multicore CPUs with access to a main memory (shared memory parallelism). However, the different CPUs access different memory domains with different speed (NUMA architecture): Usually the access to the memory banks directly connected to one CPU is faster. Each CPU “socket” consists of multiple cores (≲ 100 in 2023) with a hierarchy of caches. Inside one core, SIMD units perform identical calculations on a small vector of floating-point numbers. We focus on the performance on a single multicore CPU, but the algorithms considered can also run in parallel on a cluster (see Daas et al., 2022). Many supercomputers nowadays use GPUs which is not discussed further in this paper.

2.2.1. Roofline performance model

To obtain a simpler abstraction of the hardware, we employ the Roofline (Williams et al., 2009) performance model. The Roofline model distinguishes between computations (floating-point operations) and data transfers: The maximal performance one can achieve on a given hardware is P_peak [GFlop/s]. The bandwidth of data transfers from main memory is b _s [GByte/s]. If all data fits into some cache level, the appropriate cache bandwidth is used instead. These are the required hardware characteristics. The considered characteristics of one building block of an algorithm (e.g., one nested loop) are the number of required floating-point operations n_flops and the volume of the data transfers V_read/write. Their ratio is called computational intensity I _c ≔n_flops/V_read/write [Flop/Byte]. Assuming that the data transfers and operations overlap perfectly, this results in the following performance:

$P_{\text{roofline}}=\min \left(P_{\text{peak}},I_c{b}_s\right).$

(14)

If the compute intensity is low (I _c ≪ P_peak/b _s ), the building block is memory bound. If in contrast the compute intensity is high (I _c ≫ P_peak/b _s ), the building block is compute bound. We specifically split the algorithms considered in this paper into smaller parts (“building blocks”) because they feature not one dominating operation but are composed of multiple different blocks with different performance characteristics.

2.2.2. Memory and cache performance

In addition to the model above, some details of the memory hierarchy play a crucial role for the algorithms at hand (see Hager and Wellein (2010) for more details): First, modifying memory is often slower than reading it. In order to write to main memory, the memory region is usually first transferred to the cache, modified there and written back (write-allocate). A special CPU instruction allows to avoid this and directly stream to memory (non-temporal store). A common technique to improve the performance of data transfers is to avoid writing large (temporary) data when the algorithm can be reformulated accordingly (write-avoiding), see Carson et al. (2016). Second, the CPU caches are organized in cache lines: This means that, e.g., 8 double precision values are transferred together, always starting from a memory address divisible by the cache line size. So the data locality–e.g., which index is stored contiguously–is important.

In addition, today’s CPUs use set-associative caches that allow the mapping of one memory address to a fixed set of cache lines. Due to this, memory addresses with a specific distance (e.g., 1024 double numbers) are mapped to the same cache set and the cache effectiveness is dramatically reduced when data is accessed with specific “bad” strides (cache thrashing). This easily occurs for tensor operations if the product of some dimensions is close to a power of two. A common solution for operations on 2d arrays is padding: adding a few ignored zero rows in a matrix such that the stride is at least a few cache lines bigger than some power of two. This becomes more complicated in higher dimensions as one either obtains a complex indexing scheme or one needs to perform calculations with zeros. This is discussed in more detail in Section 5.3.

3. Numerical algorithms

In this section, we discuss three different methods to approximately solve a linear system in TT format as in equation (13). We start with a general purpose (“global”) approach based on Krylov subspace methods, TT-GMRES (Dolgov, 2013), and present some improvements for the TT format. Then, we consider the more specialized (“local”) MALS (Holtz et al., 2012; Oseledets and Dolgov, 2012), which optimizes pairs of sub-tensors similar to DMRG. Afterwards, we discuss the (“more local”) AMEn (Dolgov and Savostyanov, 2014) method, which iterates on one sub-tensor after another. Finally, we present a simple yet effective preconditioner in TT format to accelerate convergence of TT-GMRES.

3.1. Krylov methods: TT-GMRES

All methods that apply the linear operator ${\mathcal{A}}_{\text{TT}}$ on linear combinations of previously calculated directions produce solutions from the Krylov subspace ${\mathbf{K}}_k({\mathcal{A}}_{\text{TT}};V_{\text{TT}})≔\mathrm{span}\left\{V_{\text{TT}},{\mathcal{A}}_{\text{TT}}{V}_{\text{TT}},\dots ,{\mathcal{A}}_{\text{TT}}^{k-1}{V}_{\text{TT}}\right\}$ where V_TT is usually the initial residual of the problem. Different Krylov subspace methods then select the “best” solution from the subspace K _k according to different definitions of “best,” see van der Vorst (2003) and Saad (2003) for a detailed discussion. As we usually only approximate intermediate steps in TT arithmetic, we effectively employ inexact Krylov methods which are discussed thoroughly in Simoncini and Szyld (2003); van den Eshof and Sleijpen (2004). In this paper, we consider the TT-GMRES method (Dolgov, 2013) for non-symmetric problems. For symmetric operators ${\mathcal{A}}_{\text{TT}}$ , we simply omit unneeded steps to obtain a MINRES variant. However, all considerations here effectively hold for other Krylov subspace methods as well.

3.1.1. Arithmetic operations in tensor-train format

Krylov methods require the following operations which can be performed directly in the TT format (all introduced in Oseledets, 2011): Applying the operator to a vector $(Y_{\text{TT}}\leftarrow {\mathcal{A}}_{\text{TT}}{X}_{\text{TT}})$ , dot products (α ← ⟨X_TT, Y_TT⟩) and scaled additions (Y_TT ← αX_TT + Y_TT) of two tensor-trains. To reduce the computational complexity, TT truncation $({\tilde{X}}_{\text{TT}}\leftarrow {\mathrm{trunc}}_{\delta }(X_{\text{TT}}))$ approximates a tensor-train with another tensor-train with lower rank (see TT-rounding algorithm in Oseledets, 2011):

${‖X_{\text{TT}}-{\mathrm{trunc}}_{\delta }(X_{\text{TT}})‖}_F\le \delta .$

With these operations, we can perform a variant of the GMRES algorithm with additional truncation operations, see Alg. 1. This idea was first discussed in Ballani and Grasedyck (2012) for the more general H-Tucker format with a slightly different projection. We employ a variant based on Dolgov (2013). The numerical stability of TT-GMRES is analyzed in more detail in Coulaud et al. (2022). We remark that we use a more strict truncation tolerance than suggested in Dolgov (2013) based on the analysis of inexact Krylov methods in Simoncini and Szyld (2003). However, in Simoncini and Szyld (2003) only inaccurate applications of the linear operator are considered (as in line 5 of Alg. 1). We also truncate in each step of the orthogonalization (line 8) and once afterwards (line 10). We can still express the error as an error in the operator of the form (see eq. (2.2) in Simoncini and Szyld, 2003):

$\left(A+E_i\right)v_i=V_{i+1}{H}_i,i=1,\dots ,m.$

We denote the errors of all truncation operations in one Arnoldi iteration with Δw⁽⁰⁾ (line 5), Δw^(j) (line 8) and Δw⁽ⁱ⁺¹⁾. Then, we obtain for the error:

$E_i{v}_i=\mathrm{\Delta }{w}^{(i+1)}+\sum _{j=0}^i\left(I-\sum _{k=j+1}^i{v}_k{v}_k^T\right)\mathrm{\Delta }{w}^{(j)}.$

For orthogonal basis vectors v _k of the Krylov subspace, the error of the Arnoldi iteration is bounded by:

$\Vert E_i\Vert \le \Vert \mathrm{\Delta }{w}^{(i+1)}{\Vert }_F+\sum _{j=0}^i\Vert \mathrm{\Delta }{w}^{(j)}{\Vert }_F\le 0.5{\delta }_i\Vert w^{(i+1)}{\Vert }_F+\sum _{j=0}^i\frac{0.5{\delta }_i}{i+1}\Vert w^{(j)}{\Vert }_F\le {\delta }_i\Vert A\Vert .$

Of course, due to the truncations, one easily looses the orthogonality of the basis vectors v _k , see discussion in Section 3.1.2 below. Suitable tolerances δ _i require the condition number c = κ (H _m ) that we estimate using the parameter c ≈ κ (A) as suggested in Simoncini and Szyld (2003). So we obtain the following bound for the difference between the true residual vector r_∗ and the inexact residual vector ${\tilde{r}}_{\ast }$ (see equation (5.8) in Simoncini and Szyld (2003)):

$\Vert r_{\ast }-{\tilde{r}}_{\ast }{\Vert }_F\le 0.5ϵ.$

The factor 0.5 ensures that the true residual norm is smaller than the desired tolerance:

$\Vert r_{\ast }\Vert =\Vert r_{\ast }-{\tilde{r}}_{\ast }+{\tilde{r}}_{\ast }\Vert \le \Vert r_{\ast }-{\tilde{r}}_{\ast }\Vert +\Vert {\tilde{r}}_{\ast }\Vert \le ϵ.$

In our experiments, we use an optimized variant (see Section 5.1) of the standard TT truncation algorithm. An alternative randomized truncation algorithm is presented in Daas et al. (2022) for truncating sums of multiple tensor-trains (e.g., only truncating in line 10 of Alg. 1 and not in line 8).

3.1.2. Improved Gram-Schmidt orthogonalization

Above, we assumed that the resulting Krylov basis vectors are orthogonal. However, as the modified Gram-Schmidt orthogonalization is only applied approximately (truncations in line 8 and 10), this assumption is usually violated. As a result, the true residual norm might not be smaller than the prescribed tolerance ϵ even though the approximate residuals converge. To compensate for the inaccurate orthogonalization, one can prescribe a more strict truncation tolerance as discussed in Section 6 of Simoncini and Szyld (2003). Another common approach consists in re-orthogonalization: We employ the following specialized variant of a modified Gram-Schmidt iteration.

In the TT format calculating a scalar product is much faster than a truncated scaled addition (axpby) as discussed in more detail in Section 5. So we can perform additional scalar products to reorder Gram-Schmidt iterations and perform selective re-orthogonalization as shown in Alg. 2 to increase the robustness. This omits subtracting directions that are already almost orthogonal in order to avoid growing the TT-ranks. See Leon et al. (2012) and the references therein for a detailed discussion on different Gram-Schmidt orthogonalization schemes. Here, we again need to use sufficiently small truncation tolerances to fulfill the requirements of the outer inexact GMRES method. The factor 2i is an estimate for the number of inner iterations (line 9) as usually orthogonalizing “twice is enough” (Giraud et al., 2005; Parlett, 1998). And we choose k_max = 4 in all our experiments as this was sufficient for the cases we investigated.

3.1.3. Tensor-train ranks for problems with a displacement structure

Even with truncations after each tensor-train addition, the tensor-train ranks can grow exponentially in the worst case:

$\mathrm{rank}(V_{\text{TT},i+1})\le \mathrm{rank}({\mathcal{A}}_{\text{TT}})\mathrm{rank}(V_{\text{TT},i})+\sum _{j=0}^i\mathrm{rank}(V_{\text{TT},j}).$

We observe only a much smaller growth for some special linear operators ${\mathcal{A}}_{\text{TT}}$ . In particular, we consider linear operators with a displacement/Laplace structure:

${(\mathcal{A}X)}_{{\text{unfold}}_j}=A_{j,\text{left}}{X}_{{\text{unfold}}_j}+X_{{\text{unfold}}_j}{A}_{j,\text{right}},j=1,\dots ,d-1.$

For those operators, the rank of the solution is bounded if the right-hand side also has small rank as discussed for several tensor formats in Shi and Townsend (2021). We obtain the following expression for applying the operator k times:

${({\mathcal{A}}^{k}X)}_{{\text{unfold}}_j}=\sum _{i=0}^k\left(\begin{array}{c}\hfill k\hfill \\ \hfill i\hfill \end{array}\right)A_{j,\text{left}}^{k-i}{X}_{{\text{unfold}}_j}{A}_{j,\text{right}}^i.$

Similarly, for any matrix polynomial p _k of degree k, we get

${(p_k(\mathcal{A})X)}_{{\text{unfold}}_j}=\sum _{i=0}^k{\overline{p}}_{k-i}(A_{j,\text{left}})X_{{\text{unfold}}_j}{\widehat{p}}_i(A_{j,\text{right}})$

with appropriate sets of polynomials ${\overline{p}}_i$ and ${\widehat{p}}_i$ . As the jth TT rank is just the rank of the jth unfolding, this results in at most a linear growth of ranks of Krylov subspace basis vectors:

$\mathrm{rank}(V_{\text{TT},i})\le (i+1)\mathrm{rank}(V_{\text{TT},0}).$

(15)

However, we will see in Section 4.1 that this only holds in exact arithmetic. If we do not calculate the Krylov basis accurately enough, the TT ranks might again grow exponentially.

3.2. Modified alternating least-squares (MALS)

Krylov methods like TT-GMRES consider the linear operator as a black box. However, we can also exploit the tensor-train structure of the problem and project it onto the subspace of one or several sub-tensors. This is the idea of the ALS and MALS methods discussed in Holtz et al. (2012). In principle, MALS is identical to the famous DMRG method (Schollwöck, 2005; White, 1992) for eigenvalue problems from quantum physics. Here, we describe it from the point of view of numerical linear algebra. Using all but two sub-tensors of the current approximate solution, we define the operator:

$\begin{array}{ll}\hfill {\mathcal{V}}_{j,2}& :{\mathbf{R}}^{r_{j-1}\times n_j\times n_{j+1}\times r_{j+1}}\to {\mathbf{R}}^{n_1\times \cdots \times n_d},\text{with}\hfill \\ \hfill {\mathcal{V}}_{j,2}Y& =X_1\bowtie \cdots \bowtie X_{j-1}\bowtie Y\bowtie X_{j+2}\bowtie \cdots \bowtie X_d.\hfill \end{array}$

(16)

If sub-tensors (X₁, …X_j−1) are left-orthogonal and sub-tensors (X_j+2, …X _d ) are right-orthogonal, the operator ${\mathcal{V}}_{j,2}$ has orthonormal columns: ${\mathcal{V}}_{j,2}^T{\mathcal{V}}_{j,2}=I$ . Using ${\mathcal{V}}_{j,2}$ , we can project the problem onto the subspace of the sub-tensors {X _k , k < j ∨ k > j + 1}. This results in a smaller problem of the form:

${\mathcal{V}}_{j,2}^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,2}Y={\mathcal{V}}_{j,2}^T{B}_{\text{TT}}.$

(17)

For s.p.d. operators A_TT, this much smaller problem typically has a better condition number $\left(\kappa \left({\mathcal{V}}_{j,2}^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,2}\right)\le \kappa \left({\mathcal{A}}_{\text{TT}}\right)\right)$ , and its solution minimizes the error in the induced operator norm:

$Y=\arg \hspace{0.17em}\underset{Y}{\min }{‖{\mathcal{V}}_{j,2}Y-{\mathcal{X}}_{\text{TT}}^{*}‖}_{{\mathcal{A}}_{\text{TT}}}^2\text{for}Y={\left({\mathcal{V}}_{j,2}^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,2}\right)}^{-1}{\mathcal{V}}_{j,2}^T{B}_{\text{TT}},$

where ${\mathcal{X}}_{\text{TT}}^{*}$ denotes the true solution. For non-symmetric operators ${\mathcal{A}}_{\text{TT}}$ , this projection (Ritz-Galerkin) is often still successful (Dolgov and Savostyanov, 2014), but one might also consider a different projection for the left and the right-hand side of the operator (Petrov-Galerkin):

${\overline{\mathcal{V}}}_j^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,2}Y={\overline{\mathcal{V}}}_j^T{B}_{\text{TT}}.$

Here, ${\overline{\mathcal{V}}}_j$ should have the same dimensions as ${\mathcal{V}}_{j,2}$ to ensure that projected problem is square (and thus usually easier to solve). A possible non-symmetric approach is

${\overline{\mathcal{V}}}_{j}Z\approx {\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,2}\text{with}{\overline{\mathcal{V}}}_j^T{\overline{\mathcal{V}}}_j=I.$

This is only possible approximately if ${\overline{\mathcal{V}}}_j$ should have low rank in the TT format again. The solution of the projected problem then approximately minimizes the residual in the Frobenius norm (similar to GMRES). In this paper, we will not further discuss this approach as we focus on the performance of the operations involved, but other variants of projections are possible.

To solve the global problem, the MALS algorithm sweeps from the first two to the last two dimensions and back, see Alg. 3. This can be interpreted as a moving subspace correction as discussed in Oseledets et al. (2018). In contrast, the ALS algorithm only projects the problem onto the subspace of all but one sub-tensor. This yields smaller local problems but provides no mechanism to increase the TT-rank of the approximate solution. So in the simplest form, it can only converge for special initial guesses that already have the same TT-rank as the desired approximate solution. In Section 3.3, we discuss a possible way to avoid this problem.

3.2.1. Inner solver: TT-GMRES

The projected problem (17) again has a structure similar to (13) but with just two dimensions. More specifically, after some contractions, the projected operator has the form

${\mathcal{V}}_{j,2}^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,2}={\overline{A}}_{j-1,\text{left}}\bowtie A_j\bowtie A_{j+1}\bowtie {\overline{A}}_{j+2,\text{right}}$

with ${\overline{A}}_{j-1,\text{left}}\in {\mathbf{R}}^{r_{j-1}\times r_{j-1}\times r_{j-1}^A}$ and ${\overline{A}}_{j+2,\text{right}}\in {\mathbf{R}}^{r_{j+1}^A\times r_{j+1}\times r_{j+1}}$ . And the initial guess as well as the required solution is in the form:

$Y=X_j\bowtie X_{j+1}\in {\mathbf{R}}^{r_{j-1}\times n_j\times n_{j+1}\times r_{j+1}}.$

Using tensor contractions, we can also express the projected right-hand side as:

${\mathcal{V}}_{j,2}^T{B}_{\text{TT}}={\overline{B}}_{j,\text{left}}\bowtie {\overline{B}}_{j+1,\text{right}}.$

As long as the rank of Y during the iteration is much smaller than r_j−1n _j , respectively n_j+1r_j+1, it is usually beneficial to use the factorized form for Y. This results in an inner-outer scheme with an outer MALS and an inner TT-GMRES algorithm. From our observation, this also yields slightly smaller ranks in the outer MALS than using standard GMRES with the dense form of Y and a subsequent factorization.

Remark 1

In the first MALS sweeps, it is not necessary to solve the inner problem very accurately. So one can use a larger relative tolerance for the inner iteration than for the outer iteration. The same yields in the last MALS sweeps (close to the solution). Combining both aspects, we employ a relative tolerance of:

${\overline{ϵ}}_{\mathit{inner}}=\max \left({ϵ}_{\mathit{inner}},ϵ\frac{\Vert B_{\text{TT}}\Vert }{\Vert {\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-B_{\text{TT}}\Vert }\right),\text{with}{ϵ}_{\mathit{inner}}≔\sqrt{ϵ}.$

Remark 2

One might wonder if an inner-outer iteration scheme with outer (flexible) TT-GMRES and inner MALS (as preconditioner) might also work. In our experiments, this results in super-linear (up to exponential) growth of the TT-ranks in the Arnoldi iteration. This can be explained by the fact that the displacement structure of the linear operator is not retained through this form of a varying preconditioner.

For symmetric problems, our implementation switches to TT-MINRES and for positive definite problems, one can also employ a tensor-train variant of the CG algorithm.

3.3. AMEn method

For the MALS method above, always two sub-tensors are considered at once in the inner problem. One can also consider only one sub-tensor (ALS) at a time using the projection operator

$\begin{array}{ll}\hfill {\mathcal{V}}_{j,1}& :{\mathbf{R}}^{r_{j-1}\times n_j\times r_j}\to {\mathbf{R}}^{n_1\times \cdots \times n_d},\text{with}\hfill \\ \hfill {\mathcal{V}}_{j,1}Y& =X_1\bowtie \cdots \bowtie X_{j-1}\bowtie Y\bowtie X_{j+1}\bowtie \cdots \bowtie X_d.\hfill \end{array}$

(18)

This results in a smaller local operator. By contracting sub-tensors of ${\mathcal{V}}_{j,1}$ and ${\mathcal{A}}_{\text{TT}}$ , one obtains:

${\overline{\mathcal{A}}}_{\text{TT}}≔{\mathcal{V}}_{j,1}^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,1}={\overline{A}}_{j-1,\text{left}}\bowtie A_j\bowtie {\overline{A}}_{j+1,\text{right}}.$

(19)

However, one needs a way to adapt the rank of the approximate solution and to ensure convergence. An early approach from physics was just to increase the rank through adding random directions. A more sophisticated method is introduced in Dolgov and Savostyanov (2014) named AMEn (alternating minimal energy). The main idea is to enrich the subspace after each inner iteration with directions obtained from the current residual tensor. For this, the current residual tensor is projected onto the subspace of the left sub-tensors of the current approximation when sweeping left-to-right, respectively the right sub-tensors when sweeping right-to-left. The corresponding projection operators are given by:

$\begin{array}{llll}\hfill {\mathcal{V}}_{j,d-j}& :{\mathbf{R}}^{r_{j-1}\times n_j\times \cdots \times n_d}\to {\mathbf{R}}^{n_1\times \cdots \times n_d},\text{with}\hfill & \hfill {\mathcal{V}}_{j,d-j}Y& =X_1\bowtie \cdots \bowtie X_{j-1}\bowtie Y,\hfill \\ \hfill {\mathcal{V}}_{1,j}& :{\mathbf{R}}^{n_1\times \cdots \times n_j\times r_{j+1}}\to {\mathbf{R}}^{n_1\times \cdots \times n_d},\text{with}\hfill & \hfill {\mathcal{V}}_{1,j}Y& =Y\bowtie X_{j+1}\bowtie \cdots \bowtie X_d.\hfill \end{array}$

(20)

For s.p.d. operators, this results in a steepest descent step for which global convergence is shown in Dolgov and Savostyanov (2014). In practice, it is often sufficient to approximate a few directions of the residual tensor and add them to the current subspace to obtain fast convergence. The standard form of the AMEn algorithm is depicted in Alg. 4 (based on the SVD variant in Dolgov and Savostyanov, 2014). The step to update the residual $R_{\text{TT}}={\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-X_{\text{TT}}$ (line 9) requires saving all intermediate matrices from left- respectively right-orthogonalization of the sub-tensors R₁, …, R _d of R_TT. To calculate k suitable directions to enrich the subspace in the left-to-right sweep (line 15), we left-orthogonalize Z _j w.r.t. X _j , such that:

$US{V}^T=(I-{(X_j)}_{{\text{unfold}}_l}{(X_j)}_{{\text{unfold}}_l}^T){(R_j)}_{{\text{unfold}}_l},{(Z_j)}_{{\text{unfold}}_l}=U.$

This results in the method AMEn + SVD from Dolgov and Savostyanov (2014). To append the directions (line 16), we concatenate the tensors, such that ${\overline{r}}_j=r_j+k$ and ${({\overline{X}}_j)}_{:,:,r_j+1:r_j+k}=Z_{:,:,1:k}$ and ${({\overline{X}}_{j+1})}_{r_j+1:r_j+k,:,:}=0$ . The same is done in the right-to-left sweep with mirrored dimensions.

As can be seen from Alg. 4, calculating the required directions from the residual mostly involves updating the sub-tensors of the residual in every step of the sweep. This is not significantly more work than calculating the residual in the first place. Still, the complete algorithm is so cheap that the residual calculation accounts for a significant part of the overall runtime. That is why Dolgov and Savostyanov (2014) discusses several more heuristic ways to determine suitable enrichment directions.

3.3.1. TT-AMEn + ALS

The most promising variant from Dolgov and Savostyanov (2014) is based on an ALS-like approximation of the residual. As the resulting complete algorithm is not explicitly shown there, we illustrate the required steps Alg. 5 and discuss important implementation details.

First, to really decrease the work, one needs a way to check the error without using the global residual norm. In Dolgov and Savostyanov (2014), this is not further discussed but in the code used for the numerical experiments of Dolgov and Savostyanov (2014), the global residual error is estimated using the projected residuals (before solving the local problem), assuming that

$\frac{1}{2\sqrt{d-1}}{\max }_j\left(\left.{‖{\mathcal{V}}_{j,1}^T({\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-B_{\text{TT}})‖}_F\right)\right)\lesssim ‖{\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-B_{\text{TT}}‖.$

The factor 1/2 is just a heuristic way to ensure that the global residual norm is smaller than the tolerance.

Second, one needs a cheap way to enrich the subspace. For this, a rough approximation of the residual is usually sufficient which can be obtained by a fixed-rank ALS iteration (ALS (l)). This results in the following update of the jth subtensor of the current approximation of the residual ${\tilde{R}}_{\text{TT}}$ :

$\begin{array}{ll}\hfill {\tilde{R}}_j^{\prime }& ≔{\mathcal{W}}_{j,1}^T\left({\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-B_{\text{TT}}\right)={\mathcal{W}}_{j,1}^T{\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-{\mathcal{W}}_{j,1}^T{B}_{\text{TT}},\text{where}\hfill \end{array}$

(21)

$\begin{array}{ll}\hfill {\mathcal{W}}_{j,1}& :{\mathbf{R}}^{l\times n_j\times l}\to {\mathbf{R}}^{n_1\times \cdots \times n_d},{\mathcal{W}}_{j,1}Y={\tilde{R}}_1\bowtie \cdots {\tilde{R}}_{j-1}\bowtie Y\bowtie {\tilde{R}}_{j+1}\bowtie \cdots \bowtie {\tilde{R}}_d.\hfill \end{array}$

(22)

Here, ${\tilde{R}}_1,\dots ,{\tilde{R}}_{j-1}$ must be left-orthogonal and ${\tilde{R}}_{j+1},\dots ,{\tilde{R}}_d$ right-orthogonal. We remark that the approximation ${\tilde{R}}_{\text{TT}}$ cannot be used to check convergence as $\Vert {\tilde{R}}_{\text{TT}}{\Vert }_F\le \Vert R_{\text{TT}}{\Vert }_F$ . As for the projected local problem, the required additional terms can be updated successively in each step of the sweep. In our implementation, we use B_TT as initial guess for the approximate residual ${\tilde{R}}_{\text{TT}}$ and truncate it, respectively extend it by random directions to obtain the desired rank l. In Section 5.5, we show numerical experiments with both the full AMEn and the AMEn + ALS algorithm.

3.4. Preconditioning

For iterative solvers of linear systems, it is a common approach to employ a preconditioner to obtain much faster convergence. Of course, we can precondition all previously discussed algorithms. However, some additional aspects should be considered when preconditioning linear solvers in the TT format. On the one hand, for MALS and AMEn, we can employ local preconditioners for the projected operators ${\mathcal{V}}_j^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_j$ with ${\mathcal{V}}_j={\mathcal{V}}_{j,2}$ or ${\mathcal{V}}_j={\mathcal{V}}_{j,1}$ . In this case,

• one needs to calculate a different preconditioner in every step of the sweep,

• often only a few local iterations are performed,

• one may need to contract, for example, ${\mathcal{V}}_{j,1}^T{\mathcal{A}}_{\text{TT}}{\mathcal{V}}_{j,1}={\overline{A}}_{j-1,\text{left}}\bowtie A_j\bowtie {\overline{A}}_{j+1,\text{right}}$ (costly).

On the other hand, we could employ a global preconditioner which is

• either applied to tensor-train “vectors” (TT-GMRES),

• or directly to the tensor-train operator ${\mathcal{A}}_{\text{TT}}$ (TT-GMRES, MALS, AMEn),

• but is not tailored to the local problems for MALS/AMEn.

Furthermore, a global preconditioner (and also local preconditioners for MALS) should retain the low-rank of the solution and right-hand side. To emphasize this: just the fact that $\kappa (\mathcal{P}{\mathcal{A}}_{\text{TT}})\ll \kappa ({\mathcal{A}}_{\text{TT}})$ does not imply that ${\mathrm{rank}}_{\text{TT}}({\mathcal{P}\mathcal{B}}_{\text{TT}})\le {\mathrm{rank}}_{\text{TT}}(B_{\text{TT}})$ for left-preconditioning with $\mathcal{P}$ . And similarly the same problem occurs for the approximate solution X_TT for right-preconditioning.

3.4.1. TT-rank-1 preconditioner

Considering the desired properties, we suggest a simple, global rank-1 preconditioner that usually is cheap to calculate (for operators of sufficiently low rank). The two-sided variant can be constructed as follows: First, approximate the operator with a tensor-train of rank 1 using TT-truncation: ${\tilde{\mathcal{A}}}_{\text{TT}}={\tilde{A}}_1\otimes {\tilde{A}}_2\otimes \cdots \otimes {\tilde{A}}_d={\mathrm{trunc}}_{r=1}({\mathcal{A}}_{\text{TT}})$ . Then, perform an SVD for each sub-matrix ( $U_j{S}_j{V}_j^T={\tilde{A}}_j$ , j = 1, …, d) to construct the TT operators for the left and the right:

${\mathcal{P}}_{\text{TT,left}}=(S_1^{-1/2}{U}_1^T)\otimes \cdots \otimes (S_d^{-1/2}{U}_d^T),{\mathcal{P}}_{\text{TT,right}}=(V_1{S}_1^{-1/2})\otimes \cdots \otimes (V_d{S}_d^{-1/2}).$

This yields the preconditioned system:

$\left({\mathcal{P}}_{\text{TT,left}}{\mathcal{A}}_{\text{TT}}{\mathcal{P}}_{\text{TT,right}}\right)Y_{\text{TT}}={\mathcal{P}}_{\text{TT,left}}{B}_{\text{TT}}\text{with}{X}_{\text{TT}}={\mathcal{P}}_{\text{TT,right}}{Y}_{\text{TT}}.$

For a symmetric operator ${\mathcal{A}}_{\text{TT}}$ where for each sub-tensor ${(A_j)}_{:,k,l,:}={(A_j)}_{:,l,k,:}$ , the preconditioned operator is still symmetric. And as the preconditioner has rank one, it does not lead to higher ranks for the right-hand side or the exact solution X_TT. However, if the operator has a displacement structure, this is not preserved for the preconditioned operator.

4. Comparison of algorithms

In the following we present numerical experiments for linear systems in TT format. We use the pitts library (Röhrig-Zöllner and Becklas, 2024) which also contains the setup and output for all results shown in this paper. For illustrating the numerical behavior, we consider a simple multidimensional convection-diffusion equation. It is discretized using a finite difference stencil, for example, in 1d the operator is

$\frac{\text{tridiag}(-\mathrm{1,2},-1)}{h_j^2}+\frac{c}{\sqrt{d}}\frac{\text{tridiag}(\mathrm{0,1},-1)}{h_j},\text{for\hspace{0.17em}}{h}_j=\frac{1}{n_j+1}.$

Here, c denotes a convection constant and the convection direction is diagonal through all dimensions. As right-hand side, we use a vector of all ones (rank 1) or just a tensor-train with chosen rank and random sub-tensors. All cases use double-precision calculations and a desired residual tolerance of ϵ = 10⁻⁸.

4.1. Behavior of tensor-train ranks in the calculation

First, to understand the computational complexity of the different methods, we show the behavior of the TT ranks during the calculation. As shown in Figure 1(a), the improved orthogonalization in the TT-GMRES algorithms (SIMGS, Alg. 2) results in slightly slower rank growth in the Krylov basis than standard MGS; Both show the expected linear growth (for operators with a displacement structure). A naive MGS implementation (all truncations performed with tolerance δ _i in Alg. 1) results in exploding ranks after a few iterations. Furthermore, the calculated solutions of the MGS variants are not within the desired residual tolerance (with $\Vert {\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-B_{\text{TT}}{\Vert }_F\approx 10ϵ$ ). With the TT-rank-1 preconditioner, the algorithm converges after less than half the number of iterations, but higher ranks occur (as the preconditioned operator does not have a displacement structure). Similarly, a GMRES restart also results in higher TT ranks. Overall, the TT-GMRES has a computational complexity that is cubic in the maximal TT rank (r_max) of the Krylov basis: $O(dn{r}_{\text{max}}^3+d{n}^2{r}_{\text{max}}^2{(r^A)}^2)$ . For the MALS method, we observe a (linearly) growing rank of the approximation solution in the iteration up to the rank (r) of the resulting approximate solution, see Figure 1(b). However, in the inner TT-GMRES iterations, the Krylov basis ranks behave like after a GMRES restart after the first sweep. So significantly higher ranks occur in the inner iteration. The effect on the computational complexity is only quadratic for MALS (only two sub-tensors in the inner problem): $O(dnr{r}_{\text{max}}^2+d{n}^{2}r{r}_{\text{max}}{(r^A)}^2)$ . The behavior of the rank of the approximate solution for AMEn is similar to MALS. But as the local problem consists only of one sub-tensor, no higher intermediate ranks can occur in the inner iteration. So only $O(dn{r}^3+d{n}^2{r}^2{(r^A)}^2)$ operations are needed assuming that the enrichment rank k is chosen appropriately.

Figure 1.

Tensor-train ranks for the Krylov basis, respectively the approximate solution for a 20¹⁰ convection-diffusion problem (c = 10) and RHS B_TT of ones. For TT-GMRES (left), both MGS variants lead to inaccurate solutions that are not within the desired residual tolerance in contrast to all cases with SIMGS. Overall, more accurate orthogonalization (SIMGS) without restart and preconditioning features the lowest maximal ranks during the calculation. For MALS (right), the solution ranks only increase slowly with each sweep (as intended), but the Krylov basis vectors of the inner iteration again yield higher ranks. (a) Krylov basis ranks for TT-GMRES. (b) Ranks for precond. MALS with inner TT-GMRES.

4.2. Computational complexity of the different methods

The TT ranks explain the results in Figure 2: TT-GMRES needs 10–100× more operations than MALS. MALS needs 10–100× more operations than AMEn. For larger individual dimensions (n _j ), MALS with inner TT-GMRES needs fewer operations than with inner MALS. In addition, applying the operator becomes more costly, and it is beneficial to use a sparse matrix format for the sub-tensors of the operator (sparse TT-Op variant in Figure 2(a)). We also measured the behavior using the quantics TT format (QTT, Khoromskij, 2011) where we just convert the operator to a $2^{\overline{d}}$ tensor. For our test case here, QTT-MALS needs fewer operations than TT-MALS (ranks in inner iteration can grow at most by a factor of 2). However, QTT-AMEn needs more operations than TT-AMEn here. Looking at the approximate solution, it has TT ranks (…, r, 2r, 3r, 4r, …, 4r, 3r, 2r, r, …) where r is the rank of the approximate solution without transformation to QTT format. So the QTT format is not beneficial here as the solution is not well approximable with small QTT ranks.

Figure 2.

Number of floating-point operations measured using likwid (Treibig et al., 2010) for a convection-diffusion problem (c = 10). Dashed lines use the TT-rank-1 preconditioner. Dotted lines first transform the problem to the QTT format (Khoromskij, 2011). In all cases, AMEn requires orders of magnitude fewer operations than MALS and TT-GMRES. (a) Varying dimensions (RHS B_TT of ones) (b) Random RHS B_TT with varying ranks (dimension 20¹⁰).

5. Performance of algorithmic building blocks

In the following, we discuss the required basic tensor-train operations (“building blocks”) for the algorithms from Section 3. We focus on the node-level performance on a single multicore CPU with some remarks for a distributed parallel implementation. We only consider those operations where we see a significant improvement over the standard implementation as introduced in Oseledets (2011). These are in particular left-/right-orthogonalization with/without truncation, TT addition with subsequent truncation (TT-axpby + trunc), and faster contractions.

5.1. Replacing costly SVDs and pivoted QR decompositions with faster but less accurate alternatives

In Röhrig-Zöllner et al. (2022), the authors present a significantly faster implementation for decomposing a large dense tensor in the TT format (TT-SVD) using a Q-less tall-skinny QR (TSQR) algorithm. For the TT-SVD, one starts with a decomposition of the form:

$\underset{B,Q^T}{\min }{‖{(X)}_{{\text{unfold}}_{d-1}}-B{Q}^T‖}_F\text{with}{Q}^{T}Q=I.$

${(X)}_{{\text{unfold}}_{d-1}}$ is a tall-skinny matrix, and the standard algorithm uses a truncated SVD to build B and Q:

${(X)}_{{\text{unfold}}_{d-1}}=US{V}^T\text{and}B=U_{:,1:r_1}{S}_{1:r_1,1:r_1},Q=V_{:,1:r_1}.$

The optimized algorithm uses the following steps instead (grayed-out matrices are not calculated):

${(X)}_{{\text{unfold}}_{d-1}}=\overline{Q}R,R\approx \overline{U}S{V}^T,\text{and}B={(X)}_{{\text{unfold}}_{d-1}}{V}_{:,1:r_1},Q=V_{:,1:r_1}.$

As V has orthonormal columns, the matrix B can be calculated accurately this way.

For the left-/right-orthogonalization (see first loop in TT-rounding in Oseledets, 2011), we unfortunately need slightly different operations, for example, in the left-to-right QR sweep:

${(X_j)}_{\text{left}}=QB,\text{and}{(X_j^{\prime })}_{\text{left}}=Q,X_{j+1}^{\prime }=B\bowtie X_{j+1}.$

And very similarly for a left-to-right SVD sweep:

$\underset{Q,B}{\min }\Vert {(X_j)}_{\text{left}}-QB{\Vert }_F\text{and}{(X_j^{\prime })}_{\text{left}}=Q,X_{j+1}^{\prime }=B\bowtie X_{j+1}.$

The only difference is that the pure orthogonalization is exact up to the numerical rank whereas the truncation intentionally cuts off with a given tolerance. We can use the same trick as for the TT-SVD in a slightly different way here. For the orthogonalization, one obtains with pivoting matrix P:

${(X_j)}_{\text{left}}P=QR\text{and}{\tilde{X}}_j^{\prime }=X_j\bowtie (P{R}^{-1}),X_{j+1}^{\prime }=(R{P}^T)\bowtie X_{j+1}.$

(23)

And again similarly for the truncation:

${(X_j)}_{\text{left}}=QR,R=US{V}^T\text{and}{\tilde{X}}_j^{\prime }=X_j\bowtie (V{S}^{-1}),X_{j+1}^{\prime }=(S{V}^T)\bowtie X_{j+1}.$

(24)

There is a difference to the TT-SVD here: in both cases, the faster formulation might introduce a significant numerical error $(\Vert {\tilde{X}}_j^{\prime }-X_j^{\prime }{\Vert }_F)$ . For κ(R) ≫ 1, respectively κ(S) ≫ 1, applying R⁻¹, respectively VS⁻¹, can be numerically unstable. For the SVD-sweep, one can see that the error is only large in “unimportant” directions and should not affect the accuracy of the complete tensor-train: Assuming X₁, …, X_j−1 are left-orthogonal and X_j+1, …, X _d right-orthogonal, then:

$\Vert {\tilde{X}}_{\text{TT}}-X_{\text{TT}}{\Vert }_F=\Vert ({\tilde{X}}_j^{\prime }-X_j^{\prime })\bowtie X_{j+1}^{\prime }{\Vert }_F=\Vert ({\tilde{X}}_j^{\prime }-X_j^{\prime })\bowtie S{\Vert }_F.$

This is not ensured for the orthogonalization step. However, this shows that the error should be weighted by the singular values of the corresponding unfolding of the approximated tensor. So for the truncation, the left-to-right QR sweep is followed by a right-to-left SVD sweep with:

${\tilde{X}}_j^{\prime \prime }={\tilde{X}}_j^{\prime }\bowtie (\overline{U}\overline{S}),\text{for}\overline{U}\overline{S}{\overline{V}}^T={({\overline{X}}_{j+1})}_{\text{right}}.$

We suggest checking the error a posteriori using the orthogonalization error weighted by the singular values of the jth unfolding:

${‖{\overline{S}}^2-{({\tilde{X}}_j^{\prime \prime })}_{\text{left}}^T{({\tilde{X}}_j^{\prime \prime })}_{\text{left}}‖}_{\infty }\stackrel{?}{\approx }ϵ.$

In our numerical tests this was always the case even if $\Vert I-{({\tilde{X}}_j^{\prime })}_{\text{left}}^T{({\tilde{X}}_j^{\prime })}_{\text{left}}{\Vert }_{\infty }\gg ϵ$ . If this a posteriori check fails, one can still recalculate the orthogonalization with the standard algorithm.

5.2. Exploiting orthogonalities in TT-axpby + trunc

To add tensors in TT format (Z_TT = αX_TT + βY_TT = Z₁ ⋈ ⋯ ⋈ Z _d ), one combines their sub-tensors:

$\begin{array}{ll}\hfill {(Z_1)}_{1,:,:}& =\left(\begin{array}{cc}\hfill {(X_1)}_{1,:,:}\hfill & \hfill {(Y_1)}_{1,:,:}\hfill \end{array}\right),\hfill \\ \hfill {(Z_j)}_{:,i,:}& =\left(\begin{array}{cc}\hfill {(X_j)}_{:,i,:}\hfill & \hfill 0\hfill \\ \hfill 0\hfill & \hfill {(Y_j)}_{:,i,:}\hfill \end{array}\right),\text{for}i=1,\dots ,n_j,j=2,\dots ,d-1\hfill \\ \hfill {(Z_d)}_{:,:,1}& =\left(\begin{array}{c}\hfill \alpha {(X_d)}_{:,:,1}\hfill \\ \hfill \beta {(Y_d)}_{:,:,1}\hfill \end{array}\right).\hfill \end{array}$

This operation is needed in the TT-GMRES algorithm as well as in MALS and the standard AMEn. For AMEn, the operation is performed step-by-step in each step of a sweep. In all cases, the sub-tensors of X_TT and Y_TT are already left- or right-orthogonal. And after the addition, one performs a left- or right-orthogonalization of the Z_TT for a subsequent truncation step. We can exploit the block non-zero structure and pre-existing orthogonalities. In the following, we assume rank_TT (X_TT) ≥ rank_TT (Y_TT) and that X₁, …, X_d−1 are left-orthogonal. Then, we can compute left-orthogonal sub-tensors for ${\overline{Z}}_1,\dots ,{\overline{Z}}_{d-1}$ with smaller QR-decompositions Q₁R₁, …, Q_d−1R_d−1 than for the standard algorithm. For the first sub-tensor, we obtain:

$\begin{array}{llllllll}\hfill & \hfill & \hfill {(Z_1)}_{1,:,:}& =\underset{{({\overline{Z}}_1)}_{1,:,:}}{\underbrace{\left(\begin{array}{cc}\hfill {(X_1)}_{1,:,:}\hfill & \hfill Q_1\hfill \end{array}\right)}}\left(\begin{array}{cc}\hfill I\hfill & \hfill M_1\hfill \\ \hfill 0\hfill & \hfill R_1\hfill \end{array}\right)\hfill & \hfill & \text{with}\hfill & \hfill M_1& ={(X_1)}_{\text{left}}^T{(Y_1)}_{\text{left}},\hfill \\ \hfill & \hfill & \hfill & \hfill & \hfill & \hfill & \hfill Q_1{R}_1& =(I-{(X_1)}_{\text{left}}{(X_1)}_{\text{left}}^T){(Y_1)}_{\text{left}}.\hfill \end{array}$

(25)

For the next sub-tensors j = 2, …, d − 1, we obtain:

$\begin{array}{llllllll}\hfill & \hfill & \hfill {(Z_j^{\prime })}_{:,i,:}& =\left(\begin{array}{cc}\hfill I\hfill & \hfill M_{j-1}\hfill \\ \hfill 0\hfill & \hfill R_{j-1}\hfill \end{array}\right)\left(\begin{array}{cc}\hfill {(X_j)}_{:,i,:}\hfill & \hfill 0\hfill \\ \hfill 0\hfill & \hfill {(Y_j)}_{:,i,:}\hfill \end{array}\right)\hfill & \hfill & \text{for}\hfill & \hfill i& =1,\dots ,n_j\hfill \\ \hfill & \hfill & \hfill & =\underset{{({\overline{Z}}_j)}_{:,i,:}}{\underbrace{\left(\begin{array}{cc}\hfill \begin{array}{c}\hfill {(X_j)}_{:,i,:}\hfill \\ \hfill 0\hfill \end{array}\hfill & \hfill {(Q_j)}_{:,i,:}\hfill \end{array}\right)}}\left(\begin{array}{cc}\hfill I\hfill & \hfill M_j\hfill \\ \hfill 0\hfill & \hfill R_j\hfill \end{array}\right)\hfill & \hfill & \text{with}\hfill & \hfill M_j& ={\overline{X}}_j^T{\overline{Y}}_j,\hfill & \hfill Q_j{R}_j& =(I-{\overline{X}}_j{\overline{X}}_j^T){\overline{Y}}_j,\hfill \end{array}$

(26)

where we simplified the notation by introducing:

${\overline{X}}_j≔{\left(\left(\begin{array}{c}\hfill I\hfill \\ \hfill 0\hfill \end{array}\right)\bowtie X_j\right)}_{\text{left}},{\overline{Y}}_j={\left(\left(\begin{array}{c}\hfill M_{j-1}\hfill \\ \hfill R_{j-1}\hfill \end{array}\right)\bowtie Y_j\right)}_{\text{left}}.$

The last sub-tensor simply results in:

${({\overline{Z}}_d)}_{:,:,1}=\left(\begin{array}{cc}\hfill I\hfill & \hfill M_{d-1}\hfill \\ \hfill 0\hfill & \hfill R_{d-1}\hfill \end{array}\right)\left(\begin{array}{c}\hfill \alpha {(X_d)}_{:,:,1}\hfill \\ \hfill \beta {(Y_d)}_{:,:,1}\hfill \end{array}\right).$

(27)

5.2.1. Stable residual calculation with inaccurate orthogonalization

For (standard) AMEn and MALS, we update a left- respectively right-orthogonal representation of ${\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}≕Y_{\text{TT}}$ in each step and calculate the residual B_TT − Y_TT from it reusing the orthogonality. This is susceptible to numerical errors as for the solution all directions should cancel out. More specifically, we observed relative errors of up to ϵ = ∼10⁻³ when the residual norm was close to the machine precision.

In general, when combining inaccurate orthogonalization with the optimized TT-axpby + trunc algorithm, we need to adjust the formulas in (25) and (26) to compensate for the loss of orthogonality. Assuming ${\overline{X}}_j^T{\overline{X}}_j\approx I$ , we suggest to employ a Cholesky decomposition to correct the projection in (26) (and similarly in (25)) by using:

$Q_j^{\prime }{R}_j^{\prime }=\left(I-{\overline{X}}_j{(L_j{L}_j^T)}^{-1}{\overline{X}}_j^T\right){\overline{Y}}_j,\text{with}{L}_j{L}_j^T={\overline{X}}_j^T{\overline{X}}_j.$

(28)

This emulates re-orthogonalization but is faster if ${\overline{Y}}_j$ has fewer columns than ${\overline{X}}_j$ as re-orthogonalization would require to update the matrix ${\overline{X}}_j^{\prime }={\overline{X}}_j{(L_j)}^{-T}$ instead.

5.3. Faster contractions: inner iteration of AMEn

In AMEn, one needs to apply the projected TT operator from (19) to a dense tensor:

$Z={\overline{\mathcal{A}}}_{\text{TT},j}Y\text{with}{\overline{\mathcal{A}}}_{\text{TT},j}={\overline{A}}_{j-1,\text{left}}\bowtie A_j\bowtie {\overline{A}}_{j+1,\text{right}},Y,Z\in {\mathbf{R}}^{r_{j-1}\times n_j\times r_j}.$

We can calculate this with the following contractions assuming that the ordering of dimensions denoted by : is equal on the left- and right-hand sides of the equations:

$\begin{array}{ll}\hfill {(Y^{\prime })}_{:,:,:,:}& \leftarrow \sum _i{({\overline{A}}_{j+1,\text{right}})}_{:,:,i}{Y}_{:,:,i},\hfill \\ \hfill {(Y^{\prime \prime })}_{:,:,:,:}& \leftarrow \sum _{i_1,i_2}{(A_j)}_{:,:,i_1,i_2}{(Y^{\prime })}_{i_2,:,:,i_1},\hfill \\ \hfill Z_{:,:,:}& \leftarrow \sum _{i_1,i_2}{({\overline{A}}_{j-1,\text{left}})}_{:,i_1,i_2}{(Y^{\prime \prime })}_{i_2,:,:,i_1}.\hfill \end{array}$

If we reorder the array dimensions of the sub-tensors of ${\overline{\mathcal{A}}}_{\text{TT}}$ appropriately and let $({\widehat{A}}_1,{\widehat{A}}_2,{\widehat{A}}_3)$ denote the reordered tensors, we can use the following steps instead:

$\begin{array}{llllllll}\hfill & \hfill & \hfill {({\widehat{Y}}^{\prime })}_{:,:,:,k}& \leftarrow \sum _i{({\widehat{A}}_3)}_{:,i,k}{Y}_{:,:,i}\hfill & \hfill & \text{for}\hfill & \hfill k& =1,\dots ,r_j^A,\hfill \\ \hfill & \hfill & \hfill {({\widehat{Y}}^{\prime \prime })}_{:,:,:,k}& \leftarrow \sum _{i_1,i_2}{({\widehat{A}}_2)}_{:,i_1,i_2,k}{({\widehat{Y}}^{\prime })}_{:,:,i_1,i_2}\hfill & \hfill & \text{for}\hfill & \hfill k& =1,\dots ,r_{j-1}^A,\hfill \\ \hfill & \hfill & \hfill Z_{:,:,:}& \leftarrow \sum _{i_1,i_2}{({\widehat{A}}_1)}_{:,i_1,i_2}{({\widehat{Y}}^{\prime \prime })}_{:,:,i_1,i_2}.\hfill & \hfill & \hfill & \hfill \end{array}$

The benefit of this reordering is that we can combine the contracted and “free” dimensions ((i₁, i₂) and (:, :)) to obtain fewer large dimensions which reduces overhead in the implementation. In addition, we employ a column-major storage with a padding to obtain array strides that are multiples of the cache line length but not high powers of 2 to avoid cache thrashing. Another approach consists in using optimized tensor contractions as discussed in detail in Springer and Bientinesi (2018) which may reduce the overhead due to several small dimensions.

5.4. Resulting building block performance

Here, we show experiments on a single CPU socket of an AMD EPYC 7773X (“Zen 3 V-Cache”) with 64 cores and the Intel oneMKL (Intel, 2023) as underlying BLAS/LAPACK library. ¹ We use the Q-less TSQR implementation discussed in Röhrig-Zöllner et al. (2022).

Figure 3(a) shows timings for different variants of the TT-axpby + trunc operation. We observe significant speedup through replacing SVDs and pivoted QR decompositions by a fast Q-less TSQR implementation: For the QR orthogonalization sweep alone, the runtime improves by a factor of ∼4.5 for the largest case (difference between colored an black lines). For the subsequent SVD sweep, the runtime improves by a factor of ∼2.5. Overall speedup is about ∼3.5. We only see a small effect through reusing the orthogonality in the example here. It uses random tensor-trains as inputs, so the rank of the result is the sum of the individual ranks. In practice, for example, for the Arnoldi iteration or for calculating the residual in AMEn, the rank of the resulting tensor-train is often smaller which leads to less work in the SVD sweep and thus a bigger effect through reusing orthogonalities in the preceding QR sweep.

Figure 3.

Effect of building block optimizations: For adding two tensors in the tensor-train format (left), we obtain a speedup of ∼3.5 by mapping the calculation onto faster linear algebra operations as explained in Section 5.1 and Section 5.2. For applying the linear operator of the inner problem in AMEn (right), we obtain a speedup of ∼3 through directly calling optimized BLAS routines and through reordering array dimensions. (a) Timings for the truncated addition of two tensor-trains (dim. 5010, ranks r _X = 50 and r _Y = 1, …, 700). (b) Performance of the contractions for multiplying a 3d TT-operator with a dense tensor (dim.: r × 50 × r).

For the contraction in the inner iteration of AMEn, we illustrate the performance of different variants in Figure 3(b). The operation consists of 3 tensor contractions, one of them is memory-bound and the other two are compute-bound for the chosen dimensions (for r ≳ 100). The dotted line shows the performance of 3 equivalent matrix-matrix multiplications (GEMM). In particular for small dimension r, there is a significant improvement through reordering and combining the dimensions as discussed in Section 5.3. Without appropriate padding, there are performance drops whenever the r is close to a number dividable by, for example, 128 (not shown here). Overall, there is also a significant difference between using Eigen (Guennebaud and Jacob, 2010) with MKL as backend and directly calling MKL through the cblas interface, possibly because Eigen explicitly initializes the result to zero. The standard implementation used here loops over a lot of possibly small GEMM operations. We also show results using the function torch.einsum in pytorch (Paszke et al., 2019, version 2.4.0) with opt_einsum (Smith and Gray, 2018) for the unoptimized ordering of dimensions with MKL 2022.1 as backend. The function torch.einsum results in similar performance to the unoptimized variant with Eigen and MKL. This underlines that we choose the optimal contraction order and that the suggested optimization of the data layout speeds up the computation further compared to common tensor libraries such as pytorch.

5.5. Complete TT-AMEn algorithm

For the complete algorithm for solving a linear system, we focus on the AMEn method as it needs at least an order of magnitude fewer operations than the other methods. As shown in Figure 4, the full AMEn variant needs approximately twice the time of the ALS variant. This is due to calculating the global residual ${\mathcal{A}}_{\text{TT}}{X}_{\text{TT}}-B_{\text{TT}}$ which is almost as costly as calculating X_TT itself. We also show timings obtained with the ttpy implementation (ALS variant) from Dolgov and Savostyanov (2014) for which we linked with Intel MKL for the underlying operations. Through using optimized building blocks and the suggested TT-rank-1 preconditioner, we speed up the calculation by a factor of ∼5 for both the full and the ALS variants. In our tests with the ALS variant, we obtain better convergence and time-to-solution by using a slightly better approximation of the residual (cases with ALS (k + 5) where k denotes the AMEn enrichment rank and k = rank (B_TT)). The TT-rank-1 preconditioner reduces the required number of inner GMRES iterations by about a factor of two here: from ∼1530 to ∼750 for the SVD variant and from ∼1580 to ∼850 for the ALS variant for the largest case. However, the number of outer iterations (sweeps) stays the same and a significant part of the runtime is spent in the outer iteration. Therefore, the preconditioner only speeds up the total runtime by about a factor of 1.2–1.5.

Figure 4.

Timings for TT-AMEn for solving a linear system from a 50¹⁰ convection-diffusion problem (c = 10) and random RHS B_TT with varying ranks. Dashed lines use the TT-rank-1 preconditioner. Dotted black lines illustrate the asymptotic complexity using the formula c (0.35 (r/700)³ + 0.65 (r/700)²). The heuristic ALS variant (right) is about twice as fast as the full variant (left). For both variants, the time-to-solution is reduced by a factor of ∼5 by combining all suggested optimizations. (a) Full variant (AMEn + SVD) (b) ALS variant (AMEn + ALS).

6. Conclusion and future work

In this paper, we discussed the complexity and the performance of linear solvers in tensor-train format. In particular, we considered three different common methods, namely TT-GMRES, MALS (DMRG approach for linear systems) and AMEn, and tested their behavior for a simple, non-symmetric discretization of a convection-diffusion equation. Concerning the complexity in terms of floating-point operations, we illustrated that AMEn can be about 100× faster than MALS, which in turn can be about 100× faster than TT-GMRES. These results already include an optimized orthogonalization scheme for the Arnoldi iteration in the TT-GMRES method which is also used as inner iteration of the MALS method.

Concerning the performance, we focussed on the required building blocks on a many-core CPU. We suggested three improvements over the standard implementation: (a) exploiting orthogonalities in the TT-addition with subsequent truncation, (b) using a Q-less tall-skinny QR (TSQR) implementation to speed up costly singular value and QR decompositions, and (c) optimizing the memory layout/ordering of required tensor-contraction sequences for applying the tensor-train operator in the inner iteration. As improvements (a) and (b) lead to less robust underlying linear algebra operations, we discussed their accuracy in the context of the required tensor-train operations. In addition, we presented a simple generic preconditioner based on a tensor-train rank-1 approximation of the operator. Overall, we obtained a speedup of about 5× over the reference implementation on a 64-core CPU.

For future work, we want to investigate building block improvements for other tensor-network algorithms which are often based on very similar underlying operations.