6
$\begingroup$

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!

$\endgroup$
  • 1
    $\begingroup$ The first method needs $2N\times(N\times M)$ multiplications. The second method needs $2 N\times N\times N\times N\times M$. The first method is fast not because it uses a ParallelTable, but because the second one needs times-$N^2$ operations. $\endgroup$ – yarchik Jul 7 at 19:34
  • $\begingroup$ I'm not surprised that the second method is slow. Really, what I'm interested in is further speeding up the first method! Something I noticed is that slicing the tensor along the last two indices is much slow than slicing along the first two... is there an easy way to reorganize T so that accessing the matrix slices is more efficient? $\endgroup$ – user51761 Jul 7 at 19:37
  • $\begingroup$ Also for your observation there is an explanation. It has to do with how the data is physically stored. It is faster to access data elements that are located close to each other in the memory. $\endgroup$ – yarchik Jul 7 at 20:19
  • 1
    $\begingroup$ By the way, I do not think it is possible to further speed up the first method significantly. $\endgroup$ – yarchik Jul 7 at 20:21
  • 1
    $\begingroup$ @yarchik I do thinks so... ;) $\endgroup$ – Henrik Schumacher Jul 12 at 23:51
10
$\begingroup$

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.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.