Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/tfjgeorge/nngeometry
{KFAC,EKFAC,Diagonal,Implicit} Fisher Matrices and finite width NTKs in PyTorch
https://github.com/tfjgeorge/nngeometry
ek-fac ekfac fim fisher-information-matrix gram jacobian k-fac kfac neural-tangent-kernel ntk pytorch tangent
Last synced: 6 days ago
JSON representation
{KFAC,EKFAC,Diagonal,Implicit} Fisher Matrices and finite width NTKs in PyTorch
- Host: GitHub
- URL: https://github.com/tfjgeorge/nngeometry
- Owner: tfjgeorge
- License: mit
- Created: 2019-09-12T15:25:23.000Z (over 5 years ago)
- Default Branch: master
- Last Pushed: 2024-10-07T08:52:21.000Z (3 months ago)
- Last Synced: 2024-12-09T04:04:32.665Z (13 days ago)
- Topics: ek-fac, ekfac, fim, fisher-information-matrix, gram, jacobian, k-fac, kfac, neural-tangent-kernel, ntk, pytorch, tangent
- Language: Python
- Homepage: https://nngeometry.readthedocs.io
- Size: 25.1 MB
- Stars: 209
- Watchers: 7
- Forks: 21
- Open Issues: 1
-
Metadata Files:
- Readme: Readme.md
- License: LICENSE
Awesome Lists containing this project
- awesome-information-geometry - NNGeometry
README
# NNGeometry
 [](https://codecov.io/gh/tfjgeorge/nngeometry) [](https://zenodo.org/badge/latestdoi/208082966) [](https://badge.fury.io/py/nngeometry)
NNGeometry allows you to:
- compute Gauss-Newton or **Fisher Information Matrices** (FIM), as well as any matrix that is written as the covariance of gradients w.r.t. parameters, using efficient approximations such as low-rank matrices, KFAC, EKFAC, 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**.
- compute **implicit** operations on these matrices, that do not require explicitely storing large matrices that would not fit in memory.
It offers a high level abstraction over the parameter and function spaces described by neural networks. As a simple example, a parameter space vector `PVector` actually contains weight matrices, bias vectors, or convolutions kernels of the whole neural network (a set of tensors). Using NNGeometry's API, performing a step in parameter space (e.g. an update of your favorite optimization algorithm) is abstracted as a python addition: `w_next = w_previous + epsilon * delta_w`.
## Example
In the Elastic Weight Consolidation continual learning technique, you want to compute $`\left(\mathbf{w}-\mathbf{w}_{A}\right)^{\top}F\left(\mathbf{w}-\mathbf{w}_{A}\right)`$. It can be achieved with a diagonal approximation for the FIM using:
```python
F = FIM(model=model,
loader=loader,
representation=PMatDiag,
n_output=10)
regularizer = F.vTMv(w - w_a)
```
The first statement instantiates a diagonal matrix, and populates it with the diagonal coefficients of the FIM of the model `model` computed using the examples from the dataloader `loader`.
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, ...).
## Documentation
You can visit the documentation at https://nngeometry.readthedocs.io.
More example usage are available in the repository https://github.com/tfjgeorge/nngeometry-examples.
## Feature requests, bugs, contributions, or any kind of request
You are now many who are using NNGeometry in your work: do not hesitate to drop me a line ([email protected]) about your project so that I have a better understanding of your use cases or the current limitations of the library.
We welcome any feature request or bug report in the [issue tracker](https://github.com/tfjgeorge/nngeometry/issues).
We also welcome contributions, please submit your PRs!
## Citation
If you use NNGeometry in a published project, please cite our work using the following bibtex entry
```tex
@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.3},
doi = {10.5281/zenodo.4532597},
url = {https://doi.org/10.5281/zenodo.4532597}
}
```
## License
This project is distributed under the MIT license (see LICENSE file).
This project also includes code licensed under the BSD 3 clause as it borrows some code from https://github.com/owkin/grad-cnns.