Any way to improve the computational speed of the linalg.symeig() function?

[xitorch]

Any way to improve the computational speed of the linalg.symeig() function?

[mfkasim1]

Thanks for raising the issue. If you don't specify your symeig method, it will be defaulted to exact eigendecomposition, i.e. using the algorithm below:

xitorch/xitorch/_impls/linalg/symeig.py

Lines 27 to 44 in aaee1ed

	else:
	Mmatrix = M.fullmatrix() # (*BM, q, q)
	# M decomposition to make A symmetric
	# it is done this way to make it numerically stable in avoiding
	# complex eigenvalues for (near-)degenerate case
	L = torch.linalg.cholesky(Mmatrix) # (*BM, q, q)
	Linv = torch.inverse(L) # (*BM, q, q)
	LinvT = Linv.transpose(-2, -1).conj() # (*BM, q, q)
	A2 = torch.matmul(Linv, torch.matmul(Amatrix, LinvT)) # (*BAM, q, q)
	# calculate the eigenvalues and eigenvectors
	# (the eigvecs are normalized in M-space)
	# evals, evecs = torch.linalg.eigh(A2, eigenvectors=True) # (*BAM, q, q)
	evals, evecs = degen_symeig.apply(A2) # (*BAM, q, q)
	evals, evecs = _take_eigpairs(evals, evecs, neig, mode) # (*BAM, neig) and (*BAM, q, neig)
	evecs = torch.matmul(LinvT, evecs)
	return evals, evecs

Where it basically computes all the eigenvalues and just keep 100 of them.
There is another algorithm we implemented which is davidson, but not being used extensively, so it might fail.
You can try it by specifying xitorch.linalg.symeig(..., method="davidson").
If Davidson's method does not work, another option is to implement LOBPCG method, which we might be able to implement.

[biaogeng]

@mfkasim1 Thank you for the quick response! I did try to use the "davidson" method. However, the performance semms to be much worse than the default method out of the box. For a matrix size of 4500x4500, the default method took about 8 seconds, while the "davidson" method took close to 10 minutes for 20 eigenpairs (and 2.5 minutes for 1 eigenpair). I also noticed much higher CPU load when using the "davidson" method. I wonder if there's any tweaking that I can do to this method to make it faster...

For the LOBPCG method, do you think it can reach performance on par with packages like scipy, especially for sparse matrices? I'm curious if you have any plans to implement it "soon" ... I hope I could attempt to do it myself, but I assume it must be very difficult for someone who doesn't know the internals of PyTorch and is not very math-savvy.

[mfkasim1]

If you'd like, you can make a PR for it. Pytorch already has torch.lobpcg, so you even might be able to use it instead of xitorch (it just don't provide backward for non-zero M matrix). Another alternative is to add lobpcg in xitorch, but under the hood, it will use torch.lobpcg, and xitorch can provide the backward.