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

[biaogeng]

Is your feature request related to a problem? Please describe.
Yes. My general goal is to inversely determine the material parameters of a finite element model. The algorithm includes solving the eigenproblem posed by the mass and stiffness matrices of the finite element model. The overall steps are something like

1. material_parameters = neural_network_model(inputs)
2. Kmatrix, Mmatrix = assemble(mesh, material_parameters)
3. eigen_solution = xitorch.linalg.symeig(A=Kmat, neig=100, M=Mmat)
4. loss = loss_function(eigen_solution)
5. loss.backward()

The above algorithm with xitorch works flawlessly. The only issue is that once the problem size is non-trivial (with more than 10,000 degrees of freedom), it takes too long to solve the eigenproblem (~ 3min) that training becomes impractical. I noticed that the solution speed is independent of the number of requested eigenpairs.

Describe the solution you'd like
I wonder if it's possible to speed up the eigen solution since I only need a few of the lowest eigenpairs. The matrices in the finite element system are very sparse. However, I noticed that they are converted back when used to solve the eigenproblem.

Describe alternatives you've considered
Right now, the only eigen solution package that I found that supports both general eigenproblem and backward is the xitorch.linalg.symeig() (so many thanks for this amazing work!). The eigen solver in the scipy package is fast for a subset of eigenpairs, but it doesn't support backward().

Additional context
I'm not familiar with the eigen solution algorithms, but any suggestion is appreciated. And thank you again for this amazing work!

[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.