Computing Linear Layer in Tensor/Outer-Product space in PyTorch is Very Slow
It turns out that PyTorch has torch.bilinear, which is backed up by CUDA and does exactly what I need. That's neat and very fast. It still leaves the case of higher-order tensorings. I don't see any torch.trilinear and so forth, but for now it's great.
I would like to make a PyTorch model that takes the outer product of the input with itself and then does a linear regression on that. As an example, consider the input vector [1,2,3], then I would like to compute w and b to optimize [1*1, 1*2, 1*3, 2*1, 2*2, 2*3, 3*1, 3*2, 3*3] @ w + b.
For a batch input with r rows and c columns, I can do this in PyTorch with
(input.reshape(r,c,1) @ input.reshape(r,1,c)).reshape(r,c**2) @ weigts + b
My problem is that it is extraordinarily slow. Like a factor 1000 times slower and more memory consumptious than Adding a fully connected c*c RELU layer, even though it has the same number of weights.
My question is why this happens?
Is reshape a very expensive operation for PyTorch? Could I reformulate it in a different way that would make things more efficient?
Another equivalent formulation I know is torch.diag(input @ weights @ input.T) + b, but now we are computing way more values than we need (r*r) just to throw them away again.