Find Only Largest Eigenvector of Symmetric Matrix

I am looking to find only the primary eigenvector of a real, positive-definite symmetric 20x2000x2000 tensor:

K = torch.random.randn(20, 2000, 2000)
K_sym = K+K.transpose(1,2)/2
eig, ev = torch.symeig(K_sym, eigenvectors=True)
ev = ev[:,1999,:]

ev.shape

> torch.Size[20, 2000]

This implementation is very slow, even on the GPU,

Can anyone suggest a faster way of finding the first eigenvector?

For example, power iteration may be a solution. I have based this on the solution found in torch.nn.utils.spectral_norm :

def TopEigenvector(self, K, n_power_iterations=400, dim=1, eps=1e-12):
        v = torch.ones(K.shape[0],K.shape[1],1).to(self.device)
        with torch.no_grad():
            for _ in range(n_power_iterations):
                v = normalize(torch.bmm(K, v),dim=dim, eps=eps, out=v)
                if(n_power_iterations > 0):
                    v = v.clone(memory_format=torch.contiguous_format)
        return v

However, this relies on the dreaded for-loop, and also does not have a gradient for obvious reasons.

I have tested backprop using this function and, so far, it does not complain. Is there, however, a way that is gradient friendly?

EDIT:

The function is differentiable using the following:

    def TopEigenvector(self, K, n_power_iterations=400, dim=1):
        v = torch.ones(K.shape[0],K.shape[1],1).to(self.device)
        for _ in range(n_power_iterations):
            m = torch.bmm(K, v)
            n = torch.norm(m, dim=1).unsqueeze(1)
            v = m/n
        return v

Slow, but way faster than the alternative.

Hi Jack!

You might find that the Lanczos method runs faster.

Also, it may even make sense to bump over into scipy (a numpy
relative), and use an established implementation, for example,
scipy.sparse.linalg.eigs. The careful implementation might beat
pytorch tensors on the gpu.

(I’m pretty sure that it becomes equivalent to Lanczos for your real
SPD use case.)

As for getting autograd / backward() to work, I don’t have anything
useful to say.

[Note to the Pytorch Mavens: I think I’ve seen questions like this
before. I don’t personally have a pytorch use case for the top-k
eigenvectors, but is this common enough that it might make sense
to implement an autograd / gpu-friendly top-k eigenvector function?]

Best.

K. Frank

1 Like

Hi Frank,

Thanks for the informative answer!

For now, this method seems to solve the speed issue - it’s blazingly fast by comparison. I would be happy using scipy, but how would I ensure the gradients are calculated correctly? Does there exist a backprop method for such things?

Hi Jack!

If you use scipy (or numpy, etc.) you would have to write your own
custom backward() for your operation in order to get autograd to
work. For this kind of eigenvalue computation, writing backward()
would involve some significant analysis and numerical programming.
(It should, be doable – just a lot more than a one-liner.)

Now that you’ve gotten rid of the with torch.no_grad():
block I think that the version you posted above should give you
legitimate gradients. (I’m not really clear on how autograd works
with loops like this, but I don’t see any reason it shouldn’t.)

If the loop conflicts with autograd (or if you think the loop is hurting
performance), you could unwrap the loop. (Although now that
you’ve increased the default value of n_power_iterations
to 400 from 5, fully unwrapping the loop might be unattractive.)

Best.

K. Frank