Speeding up tensor contractions and multiplication

Question

Consider a tensor $T\in\mathbb{R}^{N\times N\times N\times M}$ and two vectors $x,y\in\mathbb{R}^N$. I want to compute the $N\times M$ vector defined by $X_{ij}=\operatorname{tr}(x^\top T_{:,:,i,j}y)=\operatorname{tr}_{12}(yx^\top T)$ efficiently.

I tried this in two different ways:

TCtable[x_, T_, y_] := ParallelTable[Module[{Tslice},
  Tslice = T[[;; , ;; , i, j]];
  Tr[x\[Transpose].Tslice.y] ], {i, 1, Length[x]}, {j,1,Last[Dimensions[T]]}];

TCtrace[x_, T_, y_] := TensorContract[y.x\[Transpose].T, {1, 2}];

My tensors have the very nice property that $T_{:,:,i,j}$ is very sparse for all $i,j$ (so I am representing my tensor in Mathematica as a sparse array).

With $N=500,M=3$ the parallel table method takes about 1 second, while the explicit tensor multiplication and partial trace takes about 20 seconds. Are there other clever ways to speed this up? I want to compute this for many different tensors and vectors of the same size so if there's a way to compile the code or amortize the complexity that would also be great!

Answer 1

You seem to be coming from Matlab as you try to transpose a vector, a concept that is not that useful in Mathematica. We will see in a second why that is.

Dense tensor example

n = 40;
m = 200;
T = RandomReal[{-1, 1}, {n, n, n, m}];
x = RandomReal[{-1, 1}, {n, 1}];
y = RandomReal[{-1, 1}, {n, 1}];

I defined x and y as $n\times 1$-matrices because that is how Matlab represents vectors.

That's the timings of your methods:

a = TCtable[x, T, y]; // AbsoluteTiming // First
b = TCtrace[x, T, y]; // AbsoluteTiming // First

3.0256

0.092169

Now in Mathematica, vectors are pure lists of numbers (tensors of rank 1). So, let's flatten out x and y and find a quicker way to compute your result:

x1 = Flatten[x];
y1 = Flatten[y];
c = y1.(x1.T); // AbsoluteTiming // First

0.009084

Let's also check the errors:

Max[Abs[a - b]]
Max[Abs[a - c]]

1.77636*10^-14

1.42109*10^-14

Sparse tensor example

n = 500;
m = 3;
k = 1000000;
pat = Join[
   RandomInteger[{1, n}, {k, 3}],
   RandomInteger[{1, m}, {k, 1}],
   2
   ];
vals = RandomReal[{-1, 1}, k];
T = SparseArray[pat -> vals, {n, n, n, m}, 0.];
x = RandomReal[{-1, 1}, {n, 1}];
y = RandomReal[{-1, 1}, {n, 1}];
x1 = Flatten[x];
y1 = Flatten[y];

Timings:

a = TCtable[x, T, y]; // AbsoluteTiming // First
b = TCtrace[x, T, y]; // AbsoluteTiming // First
c = y1.(x1.T); // AbsoluteTiming // First

4.77773

6.60386

0.37733

Errors:

Max[Abs[a - b]]
Max[Abs[a - c]]

2.84217*10^-14

5.32907*10^-15

The higher-rank SparseArrays in Mathematica are stored such that tensor-vector multipilcation from the right is faster than from the left. So if you can arrange to assemple the tensor in that way, that might give you a further speed-up:

S = SparseArray[pat[[All, {3, 4, 2, 1}]] -> vals, {n, m, n, n}, 0.];
d = (S.x1).y1; // AbsoluteTiming // First
Max[Abs[c - d]]

0.113273

0.

This is more than 42 times faster than OP's code.