TT Manifold and Riemannian Optimization

Overview

This notebook demonstrates Riemannian optimization on the tensor train (TT) manifold using torchtt.

Tensor trains of fixed ranks form a smooth Riemannian manifold \(\mathcal{M}_{\mathbf{r}}\), meaning we can perform optimization directly on this manifold rather than in the ambient space of all tensors. This approach is particularly powerful for problems where we want to constrain our solution to have a specific TT-rank structure.

Riemannian Gradient via Automatic Differentiation

The key insight is that torchtt can compute the Riemannian gradient using automatic differentiation (AD). Given a smooth function \(f: \mathcal{M}_{\mathbf{r}} \to \mathbb{R}\) and a point \(\mathcal{X}\) on the manifold:

  1. The Euclidean gradient \(\nabla f(\mathcal{X})\) is computed via PyTorch’s AD

  2. This gradient is then projected onto the tangent space \(T_{\mathcal{X}}\mathcal{M}_{\mathbf{r}}\) to obtain the Riemannian gradient

The function torchtt.manifold.riemannian_gradient(x, func) handles both steps automatically, making manifold optimization as simple as standard gradient descent.

Why Riemannian Optimization?

Compared to standard gradient descent on TT cores:

  • Preserves rank structure: The iterates stay on the manifold of fixed-rank TTs

  • Better conditioning: Avoids issues with redundant parameterization

  • Faster convergence: Often requires fewer iterations and larger step sizes

This notebook shows two examples:

  1. Function minimization: Finding a TT tensor that minimizes a given objective

  2. Tensor completion: Reconstructing a low-rank tensor from sparse observations

[1]:

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

Example 1: Minimizing a Quadratic Function

We create a random target tensor and define a quadratic loss function \(f(\mathcal{X}) = \frac{1}{2}\|\mathcal{X} - \mathcal{T}\|^2\) where \(\mathcal{T}\) is the target. The goal is to find a TT tensor of fixed rank that minimizes this objective.

[2]:

N = [10,11,12,13,14]
Rt = [1,3,4,5,6,1]
Rx = [1,6,6,6,6,1]
target = tntt.randn(N,Rt).round(0)
func = lambda x: 0.5*(x-target).norm(True)

Riemannian Gradient Descent

The optimization loop below uses torchtt.manifold.riemannian_gradient() to compute the Riemannian gradient at each step. Notice how:

  • The gradient is computed automatically via AD — no manual derivation needed

  • A simple retraction (rounding) projects the updated tensor back onto the manifold

  • Convergence is extremely fast, reaching machine precision in just a few iterations

[3]:

x0 = tntt.randn(N,Rx)
x =x0.clone()
for i in range(20):
    # compute riemannian gradient using AD
    gr = tntt.manifold.riemannian_gradient(x,func)

    #stepsize length
    alpha = 1.0

    # update step
    x = (x-alpha*gr).round(0,Rx)
    print('Value ' , func(x).numpy())

Value  198043.7594408848
Value  161702.02598821177
Value  71389.60114411256
Value  23698.57390652948
Value  2460.4626114669804
Value  0.9337952947342707
Value  1.815510743155269e-09
Value  9.409020444391612e-25
Value  1.479598700080797e-24
Value  1.5548277884655298e-24
Value  2.969944199121008e-24
Value  7.059492707508358e-25
Value  1.7110194779310687e-24
Value  8.010955421507427e-25
Value  1.1671818921822796e-24
Value  1.1501295103843825e-24
Value  6.074575267741516e-25
Value  6.251762467634565e-25
Value  8.579854638729817e-25
Value  7.778999496013356e-25

Comparison: Standard Gradient Descent on TT Cores

As a comparison, we perform conventional gradient descent with respect to the TT cores (the ambient space parameterization). Notice the dramatic difference:

  • Much smaller step size required for stability (α = 0.00001 vs α = 1.0)

  • Slower convergence: Even after 1000 iterations, the loss is orders of magnitude higher

  • The Riemannian approach avoids the ill-conditioning inherent in the TT core parameterization

[4]:

y = x0.detach().clone()

for i in range(1000):
    tntt.grad.watch(y)
    fval = func(y)
    deriv = tntt.grad.grad(fval,y)
    alpha = 0.00001 # for stability
    y = tntt.TT([y.cores[i].detach()-alpha*deriv[i] for i in range(len(deriv))])
    if (i+1)%100 == 0: print(func(y))

tensor(199489.7611, dtype=torch.float64)
tensor(199176.1922, dtype=torch.float64)
tensor(193181.5858, dtype=torch.float64)
tensor(86946.0497, dtype=torch.float64)
tensor(18060.0565, dtype=torch.float64)
tensor(6145.5984, dtype=torch.float64)
tensor(3902.7198, dtype=torch.float64)
tensor(2690.0054, dtype=torch.float64)
tensor(1782.9748, dtype=torch.float64)
tensor(1109.8465, dtype=torch.float64)

Example 2: Manifold Tensor Completion

A powerful application of Riemannian optimization is tensor completion — reconstructing a low-rank tensor from sparse, potentially noisy observations.

Given:

  • A set of observed indices \(\Omega = \{(i_1^{(k)}, \ldots, i_d^{(k)})\}_{k=1}^{M}\)

  • Noisy observations \(y_k = \mathcal{T}_{i_1^{(k)}, \ldots, i_d^{(k)}} + \epsilon_k\)

We minimize the least-squares loss:

\[f(\mathcal{X}) = \sum_{k=1}^{M} \left(\mathcal{X}_{i_1^{(k)}, \ldots, i_d^{(k)}} - y_k\right)^2\]

The Riemannian gradient is again computed via AD using torchtt.manifold.riemannian_gradient(). This is particularly useful because:

  • The loss involves only sparse evaluations (not the full tensor)

  • The gradient projection onto the tangent space is handled automatically

  • We can recover tensors even when observing a tiny fraction of the entries

[5]:

N = 25
target = tntt.randn([N]*4,[1,2,3,3,1])
Xs = tntt.meshgrid([tn.linspace(0,1,N, dtype = tn.float64)]*4)
target = Xs[0]+1+Xs[1]+Xs[2]+Xs[3]+Xs[0]*Xs[1]+Xs[1]*Xs[2]+tntt.TT(tn.sin(Xs[0].full()))
target = target.round(1e-10)
print(target.R)

M = 15000 # number of observations
indices = tn.randint(0,N,(M,4))

# observations are considered to be noisy
sigma_noise = 0.00001
obs = tn.normal(target.apply_mask(indices), sigma_noise)

# define the loss function
loss = lambda x: (x.apply_mask(indices)-obs).norm()**2

#%% Manifold learning
print('Riemannian gradient descent\n')
# starting point
x = tntt.randn([N]*4,[1,4,4,4,1])

tme = datetime.datetime.now()
# iterations
for i in range(10000):
    # manifold gradient
    gr = tntt.manifold.riemannian_gradient(x,loss)

    step_size = 1.0
    R = x.R
    # step update
    x = (x - step_size * gr).round(0,R)

    # compute loss value
    if (i+1)%100 == 0:
        loss_value = loss(x)
        print('Iteration %4d loss value %e error %e tensor norm %e'%(i+1,loss_value.numpy(),(x-target).norm()/target.norm(), x.norm()**2))

tme = datetime.datetime.now() - tme
print('')
print('Time elapsed',tme)
print('Number of observations %d, tensor shape %s, percentage of entries observed %6.4f'%(M,str(x.N),100*M/np.prod(x.N)))
print('Number of unknowns %d, number of observations %d, DoF/observations %.6f'%(tntt.numel(x),M,tntt.numel(x)/M))

print('Rank after rounding',x.round(1e-6))

[1, 3, 3, 2, 1]
Riemannian gradient descent

Iteration  100 loss value 2.835414e+02 error 7.095354e-02 tensor norm 6.438727e+06
Iteration  200 loss value 1.564026e+02 error 5.716689e-02 tensor norm 6.570423e+06
Iteration  300 loss value 7.340437e+01 error 4.943336e-02 tensor norm 6.582769e+06
Iteration  400 loss value 4.401695e+01 error 4.601466e-02 tensor norm 6.586063e+06
Iteration  500 loss value 3.480826e+01 error 4.460205e-02 tensor norm 6.586760e+06
Iteration  600 loss value 3.336573e+01 error 4.413909e-02 tensor norm 6.587039e+06
Iteration  700 loss value 3.175818e+01 error 4.368099e-02 tensor norm 6.587212e+06
Iteration  800 loss value 2.895179e+01 error 4.290056e-02 tensor norm 6.587260e+06
Iteration  900 loss value 2.442120e+01 error 4.159506e-02 tensor norm 6.586853e+06
Iteration 1000 loss value 1.837452e+01 error 3.981241e-02 tensor norm 6.585507e+06
Iteration 1100 loss value 1.143451e+01 error 3.807197e-02 tensor norm 6.583487e+06
Iteration 1200 loss value 6.860995e+00 error 3.674023e-02 tensor norm 6.581799e+06
Iteration 1300 loss value 3.933710e+00 error 3.574334e-02 tensor norm 6.580529e+06
Iteration 1400 loss value 2.447444e+00 error 3.506999e-02 tensor norm 6.579593e+06
Iteration 1500 loss value 1.509968e+00 error 3.457828e-02 tensor norm 6.578953e+06
Iteration 1600 loss value 8.928927e-01 error 3.427316e-02 tensor norm 6.578572e+06
Iteration 1700 loss value 5.786926e-01 error 3.407428e-02 tensor norm 6.578382e+06
Iteration 1800 loss value 3.306242e-01 error 3.393152e-02 tensor norm 6.578289e+06
Iteration 1900 loss value 1.931639e-01 error 3.383940e-02 tensor norm 6.578253e+06
Iteration 2000 loss value 1.021285e-01 error 3.378516e-02 tensor norm 6.578264e+06
Iteration 2100 loss value 5.196498e-02 error 3.375590e-02 tensor norm 6.578287e+06
Iteration 2200 loss value 2.782952e-02 error 3.373978e-02 tensor norm 6.578301e+06
Iteration 2300 loss value 1.667637e-02 error 3.372975e-02 tensor norm 6.578304e+06
Iteration 2400 loss value 1.121847e-02 error 3.372262e-02 tensor norm 6.578300e+06
Iteration 2500 loss value 8.206944e-03 error 3.371708e-02 tensor norm 6.578290e+06
Iteration 2600 loss value 6.315916e-03 error 3.371260e-02 tensor norm 6.578279e+06
Iteration 2700 loss value 5.002699e-03 error 3.370891e-02 tensor norm 6.578268e+06
Iteration 2800 loss value 4.029615e-03 error 3.370584e-02 tensor norm 6.578257e+06
Iteration 2900 loss value 3.280340e-03 error 3.370329e-02 tensor norm 6.578246e+06
Iteration 3000 loss value 2.690207e-03 error 3.370117e-02 tensor norm 6.578237e+06
Iteration 3100 loss value 2.218880e-03 error 3.369939e-02 tensor norm 6.578229e+06
Iteration 3200 loss value 1.838926e-03 error 3.369790e-02 tensor norm 6.578221e+06
Iteration 3300 loss value 1.530572e-03 error 3.369665e-02 tensor norm 6.578215e+06
Iteration 3400 loss value 1.279022e-03 error 3.369559e-02 tensor norm 6.578209e+06
Iteration 3500 loss value 1.072939e-03 error 3.369470e-02 tensor norm 6.578204e+06
Iteration 3600 loss value 9.034909e-04 error 3.369395e-02 tensor norm 6.578199e+06
Iteration 3700 loss value 7.637210e-04 error 3.369331e-02 tensor norm 6.578195e+06
Iteration 3800 loss value 6.480980e-04 error 3.369276e-02 tensor norm 6.578191e+06
Iteration 3900 loss value 5.521954e-04 error 3.369229e-02 tensor norm 6.578188e+06
Iteration 4000 loss value 4.724499e-04 error 3.369189e-02 tensor norm 6.578185e+06
Iteration 4100 loss value 4.059799e-04 error 3.369154e-02 tensor norm 6.578182e+06
Iteration 4200 loss value 3.504455e-04 error 3.369124e-02 tensor norm 6.578180e+06
Iteration 4300 loss value 3.039398e-04 error 3.369098e-02 tensor norm 6.578178e+06
Iteration 4400 loss value 2.649045e-04 error 3.369075e-02 tensor norm 6.578176e+06
Iteration 4500 loss value 2.320622e-04 error 3.369054e-02 tensor norm 6.578175e+06
Iteration 4600 loss value 2.043636e-04 error 3.369037e-02 tensor norm 6.578173e+06
Iteration 4700 loss value 1.809446e-04 error 3.369021e-02 tensor norm 6.578172e+06
Iteration 4800 loss value 1.610926e-04 error 3.369006e-02 tensor norm 6.578171e+06
Iteration 4900 loss value 1.442184e-04 error 3.368994e-02 tensor norm 6.578170e+06
Iteration 5000 loss value 1.298341e-04 error 3.368982e-02 tensor norm 6.578169e+06
Iteration 5100 loss value 1.175355e-04 error 3.368972e-02 tensor norm 6.578168e+06
Iteration 5200 loss value 1.069866e-04 error 3.368962e-02 tensor norm 6.578168e+06
Iteration 5300 loss value 9.790816e-05 error 3.368953e-02 tensor norm 6.578167e+06
Iteration 5400 loss value 9.006754e-05 error 3.368945e-02 tensor norm 6.578167e+06
Iteration 5500 loss value 8.327080e-05 error 3.368938e-02 tensor norm 6.578166e+06
Iteration 5600 loss value 7.735601e-05 error 3.368931e-02 tensor norm 6.578166e+06
Iteration 5700 loss value 7.218779e-05 error 3.368925e-02 tensor norm 6.578166e+06
Iteration 5800 loss value 6.765283e-05 error 3.368919e-02 tensor norm 6.578165e+06
Iteration 5900 loss value 6.365614e-05 error 3.368913e-02 tensor norm 6.578165e+06
Iteration 6000 loss value 6.011804e-05 error 3.368908e-02 tensor norm 6.578165e+06
Iteration 6100 loss value 5.697154e-05 error 3.368902e-02 tensor norm 6.578165e+06
Iteration 6200 loss value 5.416030e-05 error 3.368898e-02 tensor norm 6.578165e+06
Iteration 6300 loss value 5.163683e-05 error 3.368893e-02 tensor norm 6.578165e+06
Iteration 6400 loss value 4.936106e-05 error 3.368889e-02 tensor norm 6.578165e+06
Iteration 6500 loss value 4.729913e-05 error 3.368884e-02 tensor norm 6.578165e+06
Iteration 6600 loss value 4.542238e-05 error 3.368880e-02 tensor norm 6.578165e+06
Iteration 6700 loss value 4.370652e-05 error 3.368877e-02 tensor norm 6.578165e+06
Iteration 6800 loss value 4.213089e-05 error 3.368873e-02 tensor norm 6.578165e+06
Iteration 6900 loss value 4.067795e-05 error 3.368869e-02 tensor norm 6.578165e+06
Iteration 7000 loss value 3.933271e-05 error 3.368866e-02 tensor norm 6.578165e+06
Iteration 7100 loss value 3.808240e-05 error 3.368862e-02 tensor norm 6.578165e+06
Iteration 7200 loss value 3.691605e-05 error 3.368859e-02 tensor norm 6.578166e+06
Iteration 7300 loss value 3.582426e-05 error 3.368856e-02 tensor norm 6.578166e+06
Iteration 7400 loss value 3.479897e-05 error 3.368853e-02 tensor norm 6.578166e+06
Iteration 7500 loss value 3.383319e-05 error 3.368850e-02 tensor norm 6.578166e+06
Iteration 7600 loss value 3.292090e-05 error 3.368847e-02 tensor norm 6.578166e+06
Iteration 7700 loss value 3.205688e-05 error 3.368844e-02 tensor norm 6.578166e+06
Iteration 7800 loss value 3.123657e-05 error 3.368842e-02 tensor norm 6.578167e+06
Iteration 7900 loss value 3.045602e-05 error 3.368839e-02 tensor norm 6.578167e+06
Iteration 8000 loss value 2.971177e-05 error 3.368837e-02 tensor norm 6.578167e+06
Iteration 8100 loss value 2.900076e-05 error 3.368834e-02 tensor norm 6.578167e+06
Iteration 8200 loss value 2.832031e-05 error 3.368832e-02 tensor norm 6.578167e+06
Iteration 8300 loss value 2.766807e-05 error 3.368829e-02 tensor norm 6.578168e+06
Iteration 8400 loss value 2.704194e-05 error 3.368827e-02 tensor norm 6.578168e+06
Iteration 8500 loss value 2.644004e-05 error 3.368825e-02 tensor norm 6.578168e+06
Iteration 8600 loss value 2.586071e-05 error 3.368822e-02 tensor norm 6.578168e+06
Iteration 8700 loss value 2.530246e-05 error 3.368820e-02 tensor norm 6.578169e+06
Iteration 8800 loss value 2.476395e-05 error 3.368818e-02 tensor norm 6.578169e+06
Iteration 8900 loss value 2.424398e-05 error 3.368816e-02 tensor norm 6.578169e+06
Iteration 9000 loss value 2.374144e-05 error 3.368814e-02 tensor norm 6.578169e+06
Iteration 9100 loss value 2.325534e-05 error 3.368812e-02 tensor norm 6.578169e+06
Iteration 9200 loss value 2.278479e-05 error 3.368810e-02 tensor norm 6.578170e+06
Iteration 9300 loss value 2.232895e-05 error 3.368808e-02 tensor norm 6.578170e+06
Iteration 9400 loss value 2.188707e-05 error 3.368806e-02 tensor norm 6.578170e+06
Iteration 9500 loss value 2.145846e-05 error 3.368805e-02 tensor norm 6.578170e+06
Iteration 9600 loss value 2.104246e-05 error 3.368803e-02 tensor norm 6.578170e+06
Iteration 9700 loss value 2.063850e-05 error 3.368801e-02 tensor norm 6.578171e+06
Iteration 9800 loss value 2.024601e-05 error 3.368799e-02 tensor norm 6.578171e+06
Iteration 9900 loss value 1.986450e-05 error 3.368798e-02 tensor norm 6.578171e+06
Iteration 10000 loss value 1.949349e-05 error 3.368796e-02 tensor norm 6.578171e+06

Time elapsed 0:01:34.523436
Number of observations 15000, tensor shape [25, 25, 25, 25], percentage of entries observed 3.8400
Number of unknowns 1000, number of observations 15000, DoF/observations 0.066667
Rank after rounding TT with sizes and ranks:
N = [25, 25, 25, 25]
R = [1, 4, 4, 4, 1]

Device: cpu, dtype: torch.float64
#entries 1000 compression 0.00256