TT decomposition in torchtt
This notebook is a tutorial on how to use the basic functionalities of the torchtt package.
Imports
[1]:
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
Decomposition of a full tensor in TT format
We now create a 4d torch.tensor which we will use later
[2]:
tens_full = tn.reshape(tn.arange(32*16*8*10, dtype = tn.float64),[32,16,8,10])
The TT approximation of a given tensor is \(\mathsf{x}_{i_1 i_2 ... i_d} \approx \sum_{r_1,...,r_{d-1}=1}^{R_1,...,R_{d-1}} \mathsf{g}^{(1)}_{1i_1r_1}\cdots\mathsf{g}^{(d)}_{r_{d-1}i_d1}\). Using the constructor torchtt.TT() a full tensor can be decomposed in the TT format.
[3]:
tens_tt = tntt.TT(tens_full)
The newly instantiated object contains the cores as a list, the mode sizes and the rank.
[4]:
print('TT cores', tens_tt.cores)
print('Mode size ', tens_tt.N)
print('TT rank ', tens_tt.R)
TT cores [tensor([[[-0.0048, -0.3454],
[-0.0144, -0.3288],
[-0.0240, -0.3122],
[-0.0335, -0.2956],
[-0.0431, -0.2791],
[-0.0527, -0.2625],
[-0.0622, -0.2459],
[-0.0718, -0.2293],
[-0.0814, -0.2127],
[-0.0909, -0.1961],
[-0.1005, -0.1796],
[-0.1101, -0.1630],
[-0.1196, -0.1464],
[-0.1292, -0.1298],
[-0.1388, -0.1132],
[-0.1483, -0.0966],
[-0.1579, -0.0800],
[-0.1675, -0.0635],
[-0.1770, -0.0469],
[-0.1866, -0.0303],
[-0.1962, -0.0137],
[-0.2057, 0.0029],
[-0.2153, 0.0195],
[-0.2249, 0.0360],
[-0.2344, 0.0526],
[-0.2440, 0.0692],
[-0.2536, 0.0858],
[-0.2632, 0.1024],
[-0.2727, 0.1190],
[-0.2823, 0.1355],
[-0.2919, 0.1521],
[-0.3014, 0.1687]]], dtype=torch.float64), tensor([[[ 2.4448e-01, 9.6477e-03],
[ 2.4521e-01, 8.3789e-03],
[ 2.4594e-01, 7.1101e-03],
[ 2.4667e-01, 5.8413e-03],
[ 2.4741e-01, 4.5725e-03],
[ 2.4814e-01, 3.3037e-03],
[ 2.4887e-01, 2.0348e-03],
[ 2.4960e-01, 7.6604e-04],
[ 2.5034e-01, -5.0277e-04],
[ 2.5107e-01, -1.7716e-03],
[ 2.5180e-01, -3.0404e-03],
[ 2.5253e-01, -4.3092e-03],
[ 2.5327e-01, -5.5780e-03],
[ 2.5400e-01, -6.8468e-03],
[ 2.5473e-01, -8.1156e-03],
[ 2.5546e-01, -9.3845e-03]],
[[-3.1962e-03, 2.5540e-01],
[-2.7735e-03, 2.5467e-01],
[-2.3509e-03, 2.5393e-01],
[-1.9282e-03, 2.5320e-01],
[-1.5056e-03, 2.5247e-01],
[-1.0829e-03, 2.5174e-01],
[-6.6029e-04, 2.5101e-01],
[-2.3764e-04, 2.5027e-01],
[ 1.8500e-04, 2.4954e-01],
[ 6.0765e-04, 2.4881e-01],
[ 1.0303e-03, 2.4808e-01],
[ 1.4529e-03, 2.4735e-01],
[ 1.8756e-03, 2.4661e-01],
[ 2.2982e-03, 2.4588e-01],
[ 2.7209e-03, 2.4515e-01],
[ 3.1435e-03, 2.4442e-01]]], dtype=torch.float64), tensor([[[-3.5310e-01, 7.8561e-04],
[-3.5323e-01, 5.6136e-04],
[-3.5336e-01, 3.3710e-04],
[-3.5349e-01, 1.1285e-04],
[-3.5362e-01, -1.1141e-04],
[-3.5375e-01, -3.3566e-04],
[-3.5388e-01, -5.5992e-04],
[-3.5401e-01, -7.8417e-04]],
[[ 2.6179e-04, 3.5401e-01],
[ 1.8703e-04, 3.5388e-01],
[ 1.1228e-04, 3.5375e-01],
[ 3.7521e-05, 3.5362e-01],
[-3.7234e-05, 3.5349e-01],
[-1.1199e-04, 3.5336e-01],
[-1.8674e-04, 3.5323e-01],
[-2.6150e-04, 3.5310e-01]]], dtype=torch.float64), tensor([[[ 1.5132e+06],
[ 1.5133e+06],
[ 1.5133e+06],
[ 1.5134e+06],
[ 1.5134e+06],
[ 1.5135e+06],
[ 1.5135e+06],
[ 1.5136e+06],
[ 1.5137e+06],
[ 1.5137e+06]],
[[ 1.4401e+02],
[ 1.1201e+02],
[ 8.0011e+01],
[ 4.8011e+01],
[ 1.6010e+01],
[-1.5991e+01],
[-4.7991e+01],
[-7.9992e+01],
[-1.1199e+02],
[-1.4399e+02]]], dtype=torch.float64)]
Mode size [32, 16, 8, 10]
TT rank [1, np.int64(2), np.int64(2), np.int64(2), 1]
Since the TT decomposition is not exact in most of the cases, an approximation is made. If the argument eps is provided to the torchtt.TT() function the decomposition can be performed upt to the given relative accuracy. Moreover the maximum rank can also be provided as the argument rmax.
[5]:
tens_full2 = tens_full+1e-5*tn.randn(tens_full.shape, dtype=tens_full.dtype)
tens_tt2 = tntt.TT(tens_full2, eps = 1e-4)
print(tens_tt2.R)
[1, np.int64(2), np.int64(2), np.int64(2), 1]
The original tensor can be recovered using the torchtt.TT.full() method (also check if it equals the original full tensor):
[6]:
tens_full_rec = tens_tt.full()
print(tn.linalg.norm(tens_full-tens_full_rec)/tn.linalg.norm(tens_full))
tensor(2.1452e-15, dtype=torch.float64)
Using the print() function, information about the newly created torchtt.TT instance can be displayed:
[7]:
print(tens_tt)
TT with sizes and ranks:
N = [32, 16, 8, 10]
R = [1, np.int64(2), np.int64(2), np.int64(2), 1]
Device: cpu, dtype: torch.float64
#entries 180 compression 0.00439453125
Tensor operators
As a generalization of the matrix vector algebra, one can define tensor operators that act on tensors. If the tensor is \(d\)-dimensional, the tensor operator will be \(2d\)-dimensional. The goal is to perform a product \(\mathsf{Ax}\in\mathbb{M_1\times \cdots \times M_d}\) between a tensor \(\mathsf{x}\in\mathbb{R}^{N_1\times \cdots \times N_d}\) and the operator \(\mathsf{A}\in \mathbb{R}^{(M_1\times \cdots \times M_d)\times(N_1\times \cdots \times N_d)}\). For the operators the following TT matrix format is used \(\mathsf{A}_{i_1...i_d,j_1...j_d}\approx \sum\limits_{r_1,...,r_{d-1}=1}^{R_1,...,R_{d-1}} \mathsf{g}^{(1)}_{1i_1j_1r_1}\cdots\mathsf{g}^{(d)}_{r_{d-1}i_dj_d1}\).
If a tensor operator needs to be decomposed from full, the additional argument shape of the torchtt.TT() constructor has to be used to provide the shape. If the tensor operator has the shape \((M_1\times \cdots \times M_d)\times(N_1\times \cdots \times N_d)\) the argument must be passed as [(M1,N1),(M2,N2),(M3,N2),...].
[8]:
A_full = tn.reshape(tn.arange(8*4*6*3*7*9, dtype = tn.float64),[8,4,6,3,7,9])
# create an instance of torchtt.TT
A_ttm = tntt.TT(A_full, eps = 1e-12, shape = [(8,3),(4,7),(6,9)])
Slicing
Slicing operation can be performed on a tensor in TT format. If all the dimensions are indexed with an integer and the multiindices are valid, a torch.tensor with the corresponding value is returned. Slices can be also used, however the returned object in this case is again a torchtt.TT instance.
[9]:
print(tens_tt[1,2,3,4])
print(tens_tt[1,1:4,2,:])
tensor(1474.0000, dtype=torch.float64)
TT with sizes and ranks:
N = [3, 10]
R = [1, 2, 1]
Device: cpu, dtype: torch.float64
#entries 26 compression 0.8666666666666667
TT rank rounding
In some cases the TT rank becomes too large and a reduction is desired. The goal is to perform a reduction of the rank while maintaining an accuracy. The problem statement of the rounding operation is: given a tensor \(\mathsf{x}\) in the TT format with the TT rank \(\mathbf{R}\) and an \(\epsilon>0\), find a tensor \(\tilde{\mathsf{x}}\) with TT rank \(\tilde{\mathbf{R}}\leq \mathbf{R}\) such that $
||:nbsphinx-math:mathsf{x}-\tilde{\mathsf{x}}||_F:nbsphinx-math:leq `:nbsphinx-math:epsilon || :nbsphinx-math:mathsf{x}` ||_F$. This is implemented using the member method of a TT object torchtt.TT.round(). The argument epsilon is passed to the function as well as the optional argument rmax which also restricts the rank of the rounding.
We will create a tensor of TT rank \((1,6,6,6,1)\) in the TT format.
[10]:
t1 = tntt.randn([10,20,30,40],[1,2,2,2,1])
t2 = tntt.randn([10,20,30,40],[1,2,2,2,1])
t3 = tntt.randn([10,20,30,40],[1,2,2,2,1])
t1, t2, t3 = t1/t1.norm(), t2/t2.norm(), t3/t3.norm()
tt = t1+1e-3*t2+1e-6*t3
t_full = tt.full()
print(tt)
TT with sizes and ranks:
N = [10, 20, 30, 40]
R = [1, 6, 6, 6, 1]
Device: cpu, dtype: torch.float64
#entries 2100 compression 0.00875
Rounding the tensor to a relative epsilon of 1e-5 yields.
[11]:
tt1 = tt.round(1e-5)
print(tt1)
print('Error ',tn.linalg.norm(tt1.full()-tt.full())/tn.linalg.norm(tt.full()))
TT with sizes and ranks:
N = [10, 20, 30, 40]
R = [1, 4, 4, 4, 1]
Device: cpu, dtype: torch.float64
#entries 1000 compression 0.004166666666666667
Error tensor(9.9712e-07, dtype=torch.float64)
This is equivalent to removing the t3 from tt and the error will be less than 1e-6. If a truncation with epsilon=1e-2 is done, the resulting tensor will have the rank [1,2,2,2,1].
[12]:
tt1 = tt.round(1e-2)
print(tt1)
print('Error ',tn.linalg.norm(tt1.full()-tt.full())/tn.linalg.norm(tt.full()))
TT with sizes and ranks:
N = [10, 20, 30, 40]
R = [1, 2, 2, 2, 1]
Device: cpu, dtype: torch.float64
#entries 300 compression 0.00125
Error tensor(0.0010, dtype=torch.float64)
The maximum rank of a truncation can also be provided as argument.
Special tensors
Some tensors can be directly constructed in the TT format: the one tensor, the zeros tensor, the identity tensor operator adn random tensors with a given rank.
The one tensor can be created directly in the TT format using torchtt.ones().
[13]:
print(tntt.ones([2,3,4]).full())
tensor([[[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]],
[[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]]], dtype=torch.float64)
The zero tensor ca be created in the TT format using torchtt.zeros().
[14]:
print(tntt.zeros([2,3,4]).full())
tensor([[[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.]],
[[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.]]], dtype=torch.float64)
The identity tensor operator is created using torchtt.eye().
[15]:
print(tntt.eye([10,20,30]))
TT-matrix with sizes and ranks:
M = [10, 20, 30]
N = [10, 20, 30]
R = [1, 1, 1, 1]
Device: cpu, dtype: torch.float64
#entries 1400 compression 3.888888888888889e-05
Tensors with random TT cores and a given rank can be created with torchtt.random().
[16]:
print(tntt.random([3,4,5,6,7],[1,2,5,5,2,1]))
print(tntt.random([(3,7),(4,6),(5,5),(6,10),(7,2)],[1,2,5,5,2,1]))
TT with sizes and ranks:
N = [3, 4, 5, 6, 7]
R = [1, 2, 5, 5, 2, 1]
Device: cpu, dtype: torch.float64
#entries 245 compression 0.09722222222222222
TT-matrix with sizes and ranks:
M = [3, 4, 5, 6, 7]
N = [7, 6, 5, 10, 2]
R = [1, 2, 5, 5, 2, 1]
Device: cpu, dtype: torch.float64
#entries 1535 compression 0.0001450302343159486
Random tensors with a given rank and random entries with expected value 0 and given variance can be created using torchtt.randn().
[17]:
# Variance 1.0
x = tntt.randn([30]*5,[1,8,16,16,8,1])
x_full = x.full()
print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 1.0)')
# Variance 4.0
x = tntt.randn([30]*5,[1,8,16,16,8,1],var = 4.0)
x_full = x.full()
print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 4.0)')
# Variance 0.01
x = tntt.randn([30]*5,[1,8,16,16,8,1], var = 0.001)
x_full = x.full()
print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 0.001)')
# Variance 1.0 (longer train)
x = tntt.randn([10]*7, [1,4,4,4,4,4,4,1], var = 1.0)
x_full = x.full()
print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 1.0)')
Var = 0.8806055022985024 (has to be comparable to 1.0)
Var = 3.1136021195814223 (has to be comparable to 4.0)
Var = 0.0008934094513731358 (has to be comparable to 0.001)
Var = 0.6374323197748297 (has to be comparable to 1.0)