NNGeometry allows you to:

• compute Fisher Information Matrices (FIM) or derivates, using efficient approximations such as low-rank matrices, KFAC, diagonal and so on.
• compute finite-width Neural Tangent Kernels (Gram matrices), even for multiple output functions.
• compute per-examples jacobians of the loss w.r.t network parameters, or of any function such as the network's output.
• easily and efficiently compute linear algebra operations involving these matrices regardless of their approximation.

In the Elastic Weight Consolidation continual learning technique, you want to compute . It can be achieved with a diagonal approximation for the FIM using:

F = FIM(model=model,
        loader=loader,
        representation=PMatDiag,
        n_output=10)

regularizer = F.vTMv(w - w_a)

If diagonal is not sufficiently accurate then you could instead choose a KFAC approximation, by just changing PMatDiag to PMatKFAC in the above. Note that it internally involves very different operations, depending on the chosen representation (e.g. KFAC, EKFAC, ...).

You can visit the documentation at https://nngeometry.readthedocs.io.

More example usage are available in the repository https://github.com/tfjgeorge/nngeometry-examples.

We welcome any feature request or bug report in the issue tracker.

We also welcome contributions, please submit your PRs!

If you use NNGeometry in a published project, please cite our work using the following bibtex entry

@software{george_nngeometry,
  author       = {Thomas George},
  title        = {{NNGeometry: Easy and Fast Fisher Information 
                   Matrices and Neural Tangent Kernels in PyTorch}},
  month        = feb,
  year         = 2021,
  publisher    = {Zenodo},
  version      = {v0.2.1},
  doi          = {10.5281/zenodo.4532597},
  url          = {https://doi.org/10.5281/zenodo.4532597}
}