AMEN and DMRG for fast TT operations

The torchtt package includes DMRG and AMEN schemes for fast matrix vector product and elementwise inversion in the TT format.

Imports:

import torch as tn
import torch as tn
try:
    import torchtt as tntt
except:
    print('Installing torchTT...')
    %pip install git+https://github.com/ion-g-ion/torchTT
    import torchtt as tntt
import datetime

Efficient matrix vector product

When performing the multiplication between a a TT matrix and a TT tensor the rank of the result is the product of the ranks of the inputs. Therefore rank rounding has to be performed. This increases the complexity to \(\mathcal{O}(Ndr^6)\). In order to overcome this, Oseledets proposed in “DMRG Approach to Fast Linear Algebra in the TT-Format” the DMRG optimization scheme to reduce the complexity. This feature is implemented in torchtt by the member function fast_matvec() of the TT class. An example is showed in the following.

Create a random TT object and a TT matrix.

n = 32 # mode size
A = tntt.random([(n,n)]*8,[1]+7*[6]+[1]) # random array
x = tntt.random([n]*8,[1]+7*[5]+[1]) # random tensor

Increase the rank without adding redundant information. The multiplication performed in this case is actually equivalent to \(32\mathbf{\mathsf{Ax}}\).

A = A + A + A + A - A + A - A + A
x = x + x + x + x + x + x + x + x - x + x - x + x
print(A)
print(x)

TT-matrix with sizes and ranks:
M = [32, 32, 32, 32, 32, 32, 32, 32]
N = [32, 32, 32, 32, 32, 32, 32, 32]
R = [1, 48, 48, 48, 48, 48, 48, 48, 1]
Device: cpu, dtype: torch.float64
#entries 14254080 compression 1.179069862577986e-17

TT with sizes and ranks:
N = [32, 32, 32, 32, 32, 32, 32, 32]
R = [1, 60, 60, 60, 60, 60, 60, 60, 1]

Device: cpu, dtype: torch.float64
#entries 695040 compression 6.32135197520256e-07

Perform the TT matvec directly and round the result. The runtime is reported.

tme = datetime.datetime.now()
y = (A @ x).round(1e-12)
tme = datetime.datetime.now() - tme
print('Time classic ', tme)

Time classic  0:01:16.990921

This time run the fast matvec routine.

tme = datetime.datetime.now()
yf = A.fast_matvec(x)
tme = datetime.datetime.now() - tme
print('Time DMRG    ', tme)

Time DMRG     0:00:03.134253

Check if the error is the same (debugging purpose).

print('Relative error ',(y-yf).norm().numpy()/y.norm().numpy())

Relative error  1.762527613735165e-14

A second routine is the torchtt.fast_mv(). The method is described in https://arxiv.org/pdf/2410.19747. This works well for tensors in QTT.

A = tntt.random([(2,2)]*8,[1]+7*[6]+[1]) # random array
x = tntt.random([2]*8,[1]+7*[5]+[1]) # random tensor
for _ in range(8): A+=A
for _ in range(8): x+=x

tme = datetime.datetime.now()
yf2 = tntt.fast_mv(A, x)
tme = datetime.datetime.now() - tme
print('Time fast 2  ', tme)

Time fast 2   0:00:01.307426

Elementwise division in the TT format

One other basic linear algebra function that cannot be done without optimization is the elementwise division of two tensors in the TT format. In contrast to the elemntwise multiplication (where the resulting TT cores can be explicitly computed), the elementwise inversion has to be solved by means of an optimization problem (the method of choice is AMEN).

The operator “/” can be used for elemntwise division between tensors. Moreover one can use “/” between a scalar and a torchtt.TT instance.

Create 2 tensors:

• \(\mathsf{x}_{i_1i_2i_3i_4} = 2 + i_1\)

• \(\mathsf{y}_{i_1i_2i_3i_4} = i_1^2+i_2+i_3+1\)

and express them in the TT format. For both of them a TT decomposition of the elemmentwise inverse cannot be explicitly formed.

N = [32,50,44,64]
I = tntt.meshgrid([tn.arange(n,dtype = tn.float64) for n in N])
x = 2+I[0]
x = x.round(1e-15)
y = I[0]*I[0]+I[1]+I[2]+I[3]+1
y = y.round(1e-15)

Perform \(\mathsf{z}_{\mathbf{i}} = \frac{\mathsf{x}_{\mathbf{i}}}{\mathsf{z}_{\mathbf{i}}}\) and report the relative error.

z = x/y
print('Relative error', tn.linalg.norm(z.full()-x.full()/y.full())/tn.linalg.norm(z.full()))

Relative error tensor(3.3539e-11, dtype=torch.float64)

Perform \(\mathsf{u}_{\mathbf{i}} = \frac{1}{\mathsf{z}_{\mathbf{i}}}\) and report the relative error.

u = 1/y
print('Relative error', tn.linalg.norm(u.full()-1/y.full())/tn.linalg.norm(u.full()))

Relative error tensor(1.1624e-11, dtype=torch.float64)

Following are also possible:

• scalar (float, int) divided elementwise by a tensor in the TT format.

• torch.tensor with 1 element divided elementwise by a tensor in the TT format.

w = 1.0/y
a = tn.tensor(1.0)/y