pytorch-hessian-eigenthings
pytorch-hessian-eigenthings copied to clipboard
Add spectral density computation using lanczos vectors
Implement Stochastic Lanczos Quadrature using scipy or a custom Lanczos implementation. Integrates the approach from this repo.
Hi @noahgolmant,
I would like to work on this issue. Do you have some starting points?
Thanks
Hello! Algorithm 1 from this paper has a mathematical description of the approach, and the authors then implemented it in the repo I linked above. I think the jax implementation is nice and clear. Most of the jax code should be one-to-one with PyTorch. The main steps would be:
- Port the lanczos implementation to PyTorch. I am currently using scipy's built-in Lanczos method, but this only gives access to the final eigenvalues and eigenvectors, not the tridiagonal matrices.
- Port the code for density estimation from tridiagonal matrices.
- Hook up the code to use the existing
HVPOperator
class with this to go straight from model/dataloader/loss to eigenvalue density. - Add some basic precision tests with random matrices.
Would love to see this happen!