f-dangel / vivit

[TMLR 2022] Curvature access through the generalized Gauss-Newton's low-rank structure: Eigenvalues, eigenvectors, directional derivatives & Newton steps

ViViT: Curvature access through the generalized Gauss-Newton's low-rank structure

ViViT is a collection of numerical tricks to efficiently access curvature from the generalized Gauss-Newton (GGN) matrix based on its low-rank structure. Provided functionality includes computing

• GGN eigenvalues (basic example)
• GGN eigenpairs (eigenvalues + eigenvector, basic example)
• 1ˢᵗ- and 2ⁿᵈ-order directional derivatives along GGN eigenvectors (basic example)
• Directionally damped Newton steps (basic example)

These operations can also further approximate the GGN to reduce cost via sub-sampling, Monte-Carlo approximation, and block-diagonal approximation.

• Documentation: https://vivit.readthedocs.io/en/latest/
• Bug reports & feature requests: https://github.com/f-dangel/vivit/issues

How does it work? ViViT uses and extends BackPACK for PyTorch. The described functionality is realized through a combination of existing and new BackPACK extensions and hooks into its backpropagation.

Installation

pip install vivit-for-pytorch

Examples

Basic and advanced demos can be found in the documentation.

How to cite

If you are using ViViT, consider citing the paper

@article{dangel2022vivit,
  title =        {Vi{V}i{T}: Curvature Access Through The Generalized
                  Gauss-Newton{\textquoteright}s Low-Rank Structure},
  author =       {Felix Dangel and Lukas Tatzel and Philipp Hennig},
  journal =      {Transactions on Machine Learning Research (TMLR)},
  year =         2022,
}