curvlinops-for-pytorch 3.0.0

scipy Linear operators for curvature matrices in PyTorch

Linear Operators for Curvature Matrices in PyTorch

This library provides linear operators---a unified interface for matrix-free computation---for deep learning curvature matrices in PyTorch. curvlinops is inspired by SciPy's sparse.linalg.LinearOperator interface and implements a PyTorch version.

You can read our position paper to know more about why combining linear operators with curvature matrices might be a good idea.

Main features:

• Broad support of curvature matrices. curvlinops supports many common curvature matrices and approximations thereof, such as the Hessian, Fisher, generalized Gauss-Newton, and K-FAC (overview, visual tour).

• Unified interface. All linear operators share the same interface, making it easy to switch between curvature matrices.

• Purely PyTorch. All computations can run on a GPU.

• SciPy export. You can export a curvlinops linear operator to a SciPy LinearOperator with .to_scipy(). This allows plugging it into scipy, while carrying out the heavy lifting (matrix-vector multiplies) in PyTorch on GPU. My favorite example is scipy.sparse.linalg.eigsh that lets you compute a subset of eigen-pairs (example).

• Randomized estimation algorithms. curvlinops offers functionality to estimate properties the matrix represented by a linear operators, like its spectral density (example), inverse (example), trace & diagonal (example).

Installation

pip install curvlinops-for-pytorch

Useful Links

• Basic usage

• Advanced examples

• Documentation: https://curvlinops.readthedocs.io/en/latest/

• Bug reports & feature requests: https://github.com/f-dangel/curvlinops/issues

Citation

If you find curvlinops useful for your work, consider citing our position paper

@article{dangel2025position,
  title =        {Position: Curvature Matrices Should Be Democratized via Linear
                  Operators},
  author =       {Dangel, Felix and Eschenhagen, Runa and Ormaniec, Weronika and
                  Fernandez, Andres and Tatzel, Lukas and Kristiadi, Agustinus},
  journal =      {arXiv 2501.19183},
  year =         2025,
}

Future ideas

• Refactor the back-end for curvature-matrix multiplication into pure functions to improve recycle-ability and ease the use of torch.compile.

• Multi-GPU support.

• Include more curvature matrices

  • E.g. the GGN's hierarchical decomposition

Logo mage credits

• PyTorch logo: https://github.com/soumith, CC BY-SA 4.0, via Wikimedia Commons