Batched eigenvector computation

Hey guys, I am trying to perform a batched computation of eigenvectors from a tensor of dimension (B, M, M) at each epoch. This computation will be used to compute a term in the loss function of the overall model.
B is the batch size, therefore I want to obtain B separate, individual results.
I am currently achieving this via a simple for loop over the batch dimensions as follows:

batch_eig_vals, batch_eig_vecs = [], []
for i in range(S_B.shape[0]):
    eig_vals, eig_vecs = torch.symeig(S_B[i], eigenvectors=True)
    batch_eig_vals.append(eig_vals)
    batch_eig_vecs.append(eig_vecs)

batch_eig_vals = torch.stack(batch_eig_vals)
batch_eig_vecs = torch.stack(batch_eig_vecs)

However, this is extremely slow: each epoch runs about 5x slower than its counterpart without this particular term.
Are there any clever tricks I can apply to speed this up?

Hey,

Note that symeig accepts batch of inputs. So you can actually give it your whole S_B at once.
Other than that, I don’t think there is much to do. eigenvalue decomposition is quite expensive. Note that if you’re on the GPU and using small Tensors, it might be better to send them back to the CPU before doing the symeig.

1 Like

Thank you. Having been using torch.eig() I had missed that.
You mention sending Tensors to the CPU, but what performance advantages would I have doing so?

Linear algebra on the GPU is not great (more branching). So especially for small Tensors, doing these computations on the CPU might be better. But it depends a lot on your task, you’ll have to try.