hessian-eigenthings 1.0.0a1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (76) hide show
  1. hessian_eigenthings-1.0.0a1/.claude/hooks/check-docs-affected.sh +2 -0
  2. hessian_eigenthings-1.0.0a1/.claude/hooks/check_docs_affected.py +86 -0
  3. hessian_eigenthings-1.0.0a1/.claude/settings.json +15 -0
  4. hessian_eigenthings-1.0.0a1/.github/workflows/ci.yml +98 -0
  5. hessian_eigenthings-1.0.0a1/.github/workflows/docs.yml +25 -0
  6. hessian_eigenthings-1.0.0a1/.gitignore +104 -0
  7. hessian_eigenthings-1.0.0a1/LICENSE +21 -0
  8. hessian_eigenthings-1.0.0a1/PKG-INFO +166 -0
  9. hessian_eigenthings-1.0.0a1/README.md +112 -0
  10. hessian_eigenthings-1.0.0a1/docs/concepts/ggn-vs-fisher-vs-hessian.md +76 -0
  11. hessian_eigenthings-1.0.0a1/docs/concepts/numerical-stability.md +56 -0
  12. hessian_eigenthings-1.0.0a1/docs/concepts/spectral-density.md +78 -0
  13. hessian_eigenthings-1.0.0a1/docs/concepts/top-k-eigenvalues.md +81 -0
  14. hessian_eigenthings-1.0.0a1/docs/concepts/trace-estimation.md +67 -0
  15. hessian_eigenthings-1.0.0a1/docs/concepts/what-is-the-hessian.md +44 -0
  16. hessian_eigenthings-1.0.0a1/docs/concepts/why-hvp-not-full-h.md +46 -0
  17. hessian_eigenthings-1.0.0a1/docs/getting-started/installation.md +20 -0
  18. hessian_eigenthings-1.0.0a1/docs/getting-started/quickstart.md +51 -0
  19. hessian_eigenthings-1.0.0a1/docs/getting-started/transformers-quickstart.md +50 -0
  20. hessian_eigenthings-1.0.0a1/docs/how-to/analyze-a-huggingface-model.md +65 -0
  21. hessian_eigenthings-1.0.0a1/docs/how-to/analyze-with-transformer-lens.md +62 -0
  22. hessian_eigenthings-1.0.0a1/docs/how-to/custom-curvature-operators.md +119 -0
  23. hessian_eigenthings-1.0.0a1/docs/how-to/custom-loss-functions.md +122 -0
  24. hessian_eigenthings-1.0.0a1/docs/how-to/distributed-ddp.md +59 -0
  25. hessian_eigenthings-1.0.0a1/docs/how-to/per-layer-hessian.md +83 -0
  26. hessian_eigenthings-1.0.0a1/docs/index.md +12 -0
  27. hessian_eigenthings-1.0.0a1/docs/reference/algorithms.md +3 -0
  28. hessian_eigenthings-1.0.0a1/docs/reference/api.md +3 -0
  29. hessian_eigenthings-1.0.0a1/docs/reference/loss_fns.md +3 -0
  30. hessian_eigenthings-1.0.0a1/examples/README.md +17 -0
  31. hessian_eigenthings-1.0.0a1/examples/huggingface_tiny_gpt2.py +66 -0
  32. hessian_eigenthings-1.0.0a1/examples/supervised_mlp.py +66 -0
  33. hessian_eigenthings-1.0.0a1/examples/transformer_lens_attention_only.py +64 -0
  34. hessian_eigenthings-1.0.0a1/hessian_eigenthings/__init__.py +3 -0
  35. hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/__init__.py +35 -0
  36. hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/lanczos.py +184 -0
  37. hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/power_iteration.py +151 -0
  38. hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/result.py +16 -0
  39. hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/spectral_density.py +135 -0
  40. hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/trace.py +151 -0
  41. hessian_eigenthings-1.0.0a1/hessian_eigenthings/batching.py +61 -0
  42. hessian_eigenthings-1.0.0a1/hessian_eigenthings/linalg/__init__.py +3 -0
  43. hessian_eigenthings-1.0.0a1/hessian_eigenthings/linalg/backend.py +69 -0
  44. hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/__init__.py +31 -0
  45. hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/huggingface.py +59 -0
  46. hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/standard.py +51 -0
  47. hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/transformer_lens.py +47 -0
  48. hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/__init__.py +19 -0
  49. hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/base.py +70 -0
  50. hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/distributed/__init__.py +3 -0
  51. hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/distributed/ddp.py +85 -0
  52. hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/fisher.py +149 -0
  53. hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/ggn.py +147 -0
  54. hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/hessian.py +220 -0
  55. hessian_eigenthings-1.0.0a1/hessian_eigenthings/param_utils.py +75 -0
  56. hessian_eigenthings-1.0.0a1/mkdocs.yml +93 -0
  57. hessian_eigenthings-1.0.0a1/pyproject.toml +111 -0
  58. hessian_eigenthings-1.0.0a1/tests/conftest.py +27 -0
  59. hessian_eigenthings-1.0.0a1/tests/cross_library/__init__.py +0 -0
  60. hessian_eigenthings-1.0.0a1/tests/cross_library/test_against_curvlinops.py +213 -0
  61. hessian_eigenthings-1.0.0a1/tests/test_batching.py +63 -0
  62. hessian_eigenthings-1.0.0a1/tests/test_ddp_hessian.py +88 -0
  63. hessian_eigenthings-1.0.0a1/tests/test_eigendecomp_on_hessian.py +86 -0
  64. hessian_eigenthings-1.0.0a1/tests/test_fisher_operator.py +101 -0
  65. hessian_eigenthings-1.0.0a1/tests/test_ggn_operator.py +162 -0
  66. hessian_eigenthings-1.0.0a1/tests/test_hessian_finite_difference.py +134 -0
  67. hessian_eigenthings-1.0.0a1/tests/test_hessian_operator.py +182 -0
  68. hessian_eigenthings-1.0.0a1/tests/test_lanczos_random_matrix.py +124 -0
  69. hessian_eigenthings-1.0.0a1/tests/test_linalg_backend.py +91 -0
  70. hessian_eigenthings-1.0.0a1/tests/test_loss_fns.py +107 -0
  71. hessian_eigenthings-1.0.0a1/tests/test_param_utils.py +91 -0
  72. hessian_eigenthings-1.0.0a1/tests/test_power_iteration_random_matrix.py +96 -0
  73. hessian_eigenthings-1.0.0a1/tests/test_smoke.py +5 -0
  74. hessian_eigenthings-1.0.0a1/tests/test_spectral_density.py +150 -0
  75. hessian_eigenthings-1.0.0a1/tests/test_trace.py +141 -0
  76. hessian_eigenthings-1.0.0a1/uv.lock +5090 -0
@@ -0,0 +1,2 @@
1
+ #!/usr/bin/env bash
2
+ exec python3 "$(dirname "$0")/check_docs_affected.py"
@@ -0,0 +1,86 @@
1
+ """PostToolUse hook for Claude Code.
2
+
3
+ When a source file under hessian_eigenthings/ is edited, print a reminder to
4
+ stderr listing the related docs page(s). The reminder appears in Claude's
5
+ context, nudging it to keep the docs in sync. Soft-only — does not block tools.
6
+
7
+ Reads the tool input as JSON on stdin (Claude Code's hook protocol).
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import json
13
+ import os
14
+ import sys
15
+
16
+ # Map source files (paths relative to repo root) to canonical docs pages.
17
+ DOCS_MAP: dict[str, list[str]] = {
18
+ "hessian_eigenthings/operators/hessian.py": [
19
+ "docs/concepts/what-is-the-hessian.md",
20
+ "docs/concepts/why-hvp-not-full-h.md",
21
+ ],
22
+ "hessian_eigenthings/operators/ggn.py": [
23
+ "docs/concepts/ggn-vs-fisher-vs-hessian.md",
24
+ ],
25
+ "hessian_eigenthings/operators/fisher.py": [
26
+ "docs/concepts/ggn-vs-fisher-vs-hessian.md",
27
+ ],
28
+ "hessian_eigenthings/operators/distributed/ddp.py": [
29
+ "docs/how-to/distributed-ddp.md",
30
+ ],
31
+ "hessian_eigenthings/algorithms/lanczos.py": [
32
+ "docs/concepts/top-k-eigenvalues.md",
33
+ ],
34
+ "hessian_eigenthings/algorithms/power_iteration.py": [
35
+ "docs/concepts/top-k-eigenvalues.md",
36
+ ],
37
+ "hessian_eigenthings/algorithms/trace.py": [
38
+ "docs/concepts/trace-estimation.md",
39
+ ],
40
+ "hessian_eigenthings/algorithms/spectral_density.py": [
41
+ "docs/concepts/spectral-density.md",
42
+ ],
43
+ "hessian_eigenthings/loss_fns/standard.py": [
44
+ "docs/how-to/custom-loss-functions.md",
45
+ ],
46
+ "hessian_eigenthings/loss_fns/huggingface.py": [
47
+ "docs/how-to/analyze-a-huggingface-model.md",
48
+ ],
49
+ "hessian_eigenthings/loss_fns/transformer_lens.py": [
50
+ "docs/how-to/analyze-with-transformer-lens.md",
51
+ ],
52
+ "hessian_eigenthings/param_utils.py": [
53
+ "docs/how-to/per-layer-hessian.md",
54
+ ],
55
+ }
56
+
57
+
58
+ def main() -> int:
59
+ try:
60
+ payload = json.load(sys.stdin)
61
+ except json.JSONDecodeError:
62
+ return 0
63
+
64
+ file_path = (payload.get("tool_input") or {}).get("file_path", "")
65
+ if not file_path:
66
+ return 0
67
+
68
+ cwd = os.getcwd()
69
+ rel_path = file_path[len(cwd) + 1 :] if file_path.startswith(cwd) else file_path
70
+
71
+ related = DOCS_MAP.get(rel_path)
72
+ if not related:
73
+ return 0
74
+
75
+ print(f"[doc-drift hook] Edited {rel_path}. Related docs:", file=sys.stderr)
76
+ for doc in related:
77
+ print(f" - {doc}", file=sys.stderr)
78
+ print(
79
+ "[doc-drift hook] If the public API or behavior changed, update those pages.",
80
+ file=sys.stderr,
81
+ )
82
+ return 0
83
+
84
+
85
+ if __name__ == "__main__":
86
+ sys.exit(main())
@@ -0,0 +1,15 @@
1
+ {
2
+ "hooks": {
3
+ "PostToolUse": [
4
+ {
5
+ "matcher": "Edit|Write|MultiEdit",
6
+ "hooks": [
7
+ {
8
+ "type": "command",
9
+ "command": ".claude/hooks/check-docs-affected.sh"
10
+ }
11
+ ]
12
+ }
13
+ ]
14
+ }
15
+ }
@@ -0,0 +1,98 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [master, v1]
6
+ pull_request:
7
+ branches: [master, v1]
8
+
9
+ concurrency:
10
+ group: ${{ github.workflow }}-${{ github.ref }}
11
+ cancel-in-progress: true
12
+
13
+ jobs:
14
+ lint:
15
+ runs-on: ubuntu-latest
16
+ steps:
17
+ - uses: actions/checkout@v4
18
+ - uses: astral-sh/setup-uv@v5
19
+ with:
20
+ enable-cache: true
21
+ - run: uv sync --group dev
22
+ - run: uv run ruff check .
23
+ - run: uv run black --check .
24
+
25
+ typecheck:
26
+ runs-on: ubuntu-latest
27
+ steps:
28
+ - uses: actions/checkout@v4
29
+ - uses: astral-sh/setup-uv@v5
30
+ with:
31
+ enable-cache: true
32
+ - run: uv sync --group dev
33
+ - run: uv run mypy
34
+
35
+ test:
36
+ runs-on: ubuntu-latest
37
+ strategy:
38
+ fail-fast: false
39
+ matrix:
40
+ python-version: ["3.10", "3.11", "3.12"]
41
+ steps:
42
+ - uses: actions/checkout@v4
43
+ - uses: astral-sh/setup-uv@v5
44
+ with:
45
+ enable-cache: true
46
+ - run: uv sync --group dev --python ${{ matrix.python-version }}
47
+ - run: uv run pytest --cov=hessian_eigenthings --cov-report=term-missing
48
+
49
+ test-cross-library:
50
+ runs-on: ubuntu-latest
51
+ steps:
52
+ - uses: actions/checkout@v4
53
+ - uses: astral-sh/setup-uv@v5
54
+ with:
55
+ enable-cache: true
56
+ - run: uv sync --group dev --extra curvlinops
57
+ # Tolerate exit code 5 ("no tests collected") so the job stays green
58
+ # if all curvlinops-marked tests are removed in some future refactor.
59
+ - run: uv run pytest -m curvlinops || [ $? -eq 5 ]
60
+
61
+ test-extras:
62
+ runs-on: ubuntu-latest
63
+ steps:
64
+ - uses: actions/checkout@v4
65
+ - uses: astral-sh/setup-uv@v5
66
+ with:
67
+ enable-cache: true
68
+ - run: uv sync --group dev --extra transformers --extra transformer-lens
69
+ - run: uv run pytest -m "transformers or transformer_lens" || [ $? -eq 5 ]
70
+
71
+ test-examples:
72
+ runs-on: ubuntu-latest
73
+ steps:
74
+ - uses: actions/checkout@v4
75
+ - uses: astral-sh/setup-uv@v5
76
+ with:
77
+ enable-cache: true
78
+ - run: uv sync --group dev --extra transformers --extra transformer-lens
79
+ - name: supervised_mlp
80
+ run: timeout 60 uv run python examples/supervised_mlp.py
81
+ - name: huggingface_tiny_gpt2
82
+ run: timeout 120 uv run python examples/huggingface_tiny_gpt2.py
83
+ - name: transformer_lens_attention_only
84
+ run: timeout 180 uv run python examples/transformer_lens_attention_only.py
85
+
86
+ test-docs-codeblocks:
87
+ runs-on: ubuntu-latest
88
+ steps:
89
+ - uses: actions/checkout@v4
90
+ - uses: astral-sh/setup-uv@v5
91
+ with:
92
+ enable-cache: true
93
+ - run: uv sync --group dev --extra transformers --extra transformer-lens
94
+ # Only the quickstart pages have fully self-contained, runnable code.
95
+ # Concept/how-to pages use placeholder identifiers (model, loader, ...)
96
+ # to illustrate API shape; those snippets are checked by example scripts
97
+ # in test-examples instead.
98
+ - run: uv run pytest --markdown-docs docs/getting-started/
@@ -0,0 +1,25 @@
1
+ name: Docs
2
+
3
+ on:
4
+ push:
5
+ branches: [master, v1]
6
+ pull_request:
7
+ branches: [master, v1]
8
+
9
+ permissions:
10
+ contents: write
11
+
12
+ jobs:
13
+ build:
14
+ runs-on: ubuntu-latest
15
+ steps:
16
+ - uses: actions/checkout@v4
17
+ with:
18
+ fetch-depth: 0
19
+ - uses: astral-sh/setup-uv@v5
20
+ with:
21
+ enable-cache: true
22
+ - run: uv sync --group docs
23
+ - run: uv run mkdocs build --strict
24
+ - if: github.event_name == 'push' && github.ref == 'refs/heads/master'
25
+ run: uv run mkdocs gh-deploy --force
@@ -0,0 +1,104 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ *.egg-info/
24
+ .installed.cfg
25
+ *.egg
26
+ MANIFEST
27
+
28
+ # PyInstaller
29
+ # Usually these files are written by a python script from a template
30
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
31
+ *.manifest
32
+ *.spec
33
+
34
+ # Installer logs
35
+ pip-log.txt
36
+ pip-delete-this-directory.txt
37
+
38
+ # Unit test / coverage reports
39
+ htmlcov/
40
+ .tox/
41
+ .coverage
42
+ .coverage.*
43
+ .cache
44
+ nosetests.xml
45
+ coverage.xml
46
+ *.cover
47
+ .hypothesis/
48
+ .pytest_cache/
49
+
50
+ # Translations
51
+ *.mo
52
+ *.pot
53
+
54
+ # Django stuff:
55
+ *.log
56
+ local_settings.py
57
+ db.sqlite3
58
+
59
+ # Flask stuff:
60
+ instance/
61
+ .webassets-cache
62
+
63
+ # Scrapy stuff:
64
+ .scrapy
65
+
66
+ # Sphinx documentation
67
+ docs/_build/
68
+
69
+ # PyBuilder
70
+ target/
71
+
72
+ # Jupyter Notebook
73
+ .ipynb_checkpoints
74
+
75
+ # pyenv
76
+ .python-version
77
+
78
+ # celery beat schedule file
79
+ celerybeat-schedule
80
+
81
+ # SageMath parsed files
82
+ *.sage.py
83
+
84
+ # Environments
85
+ .env
86
+ .venv
87
+ env/
88
+ venv/
89
+ ENV/
90
+ env.bak/
91
+ venv.bak/
92
+
93
+ # Spyder project settings
94
+ .spyderproject
95
+ .spyproject
96
+
97
+ # Rope project settings
98
+ .ropeproject
99
+
100
+ # mkdocs documentation
101
+ /site
102
+
103
+ # mypy
104
+ .mypy_cache/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2018 Noah Golmant
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,166 @@
1
+ Metadata-Version: 2.4
2
+ Name: hessian-eigenthings
3
+ Version: 1.0.0a1
4
+ Summary: Iterative eigendecomposition of curvature operators (Hessian, GGN, Fisher) for PyTorch models.
5
+ Project-URL: Homepage, https://github.com/noahgolmant/pytorch-hessian-eigenthings
6
+ Project-URL: Documentation, https://noahgolmant.github.io/pytorch-hessian-eigenthings
7
+ Project-URL: Issues, https://github.com/noahgolmant/pytorch-hessian-eigenthings/issues
8
+ Author: Noah Golmant
9
+ License: MIT License
10
+
11
+ Copyright (c) 2018 Noah Golmant
12
+
13
+ Permission is hereby granted, free of charge, to any person obtaining a copy
14
+ of this software and associated documentation files (the "Software"), to deal
15
+ in the Software without restriction, including without limitation the rights
16
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
17
+ copies of the Software, and to permit persons to whom the Software is
18
+ furnished to do so, subject to the following conditions:
19
+
20
+ The above copyright notice and this permission notice shall be included in all
21
+ copies or substantial portions of the Software.
22
+
23
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
24
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
25
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
26
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
27
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
28
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
29
+ SOFTWARE.
30
+ License-File: LICENSE
31
+ Classifier: Development Status :: 3 - Alpha
32
+ Classifier: Intended Audience :: Science/Research
33
+ Classifier: License :: OSI Approved :: MIT License
34
+ Classifier: Programming Language :: Python :: 3
35
+ Classifier: Programming Language :: Python :: 3.10
36
+ Classifier: Programming Language :: Python :: 3.11
37
+ Classifier: Programming Language :: Python :: 3.12
38
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
39
+ Requires-Python: >=3.10
40
+ Requires-Dist: numpy>=1.24
41
+ Requires-Dist: scipy>=1.11
42
+ Requires-Dist: torch>=2.3
43
+ Requires-Dist: tqdm>=4.65
44
+ Provides-Extra: curvlinops
45
+ Requires-Dist: curvlinops-for-pytorch>=3.0; extra == 'curvlinops'
46
+ Provides-Extra: pyhessian
47
+ Requires-Dist: pyhessian>=0.1; extra == 'pyhessian'
48
+ Provides-Extra: transformer-lens
49
+ Requires-Dist: transformer-lens>=2.0; extra == 'transformer-lens'
50
+ Provides-Extra: transformers
51
+ Requires-Dist: datasets>=2.14; extra == 'transformers'
52
+ Requires-Dist: transformers>=4.40; extra == 'transformers'
53
+ Description-Content-Type: text/markdown
54
+
55
+ # pytorch-hessian-eigenthings
56
+
57
+ The `hessian-eigenthings` module provides an efficient (and scalable!) way to compute the eigendecomposition of the Hessian, along with other curvature matrices like the Generalized Gauss-Newton and empirical Fisher, for an arbitrary PyTorch model. It uses PyTorch's Hessian-vector product and your choice of (a) the Lanczos method or (b) stochastic power iteration with deflation to compute the top eigenvalues and eigenvectors. There's also Hutch++ for trace estimation and Stochastic Lanczos Quadrature for the spectral density.
58
+
59
+ > **v1.0.0a1**: alpha release. The original `0.x` API has been removed; pin `hessian-eigenthings==0.0.2` if you depend on it.
60
+
61
+ ## Why use this?
62
+
63
+ The eigenvalues and eigenvectors of the Hessian have been implicated in many generalization properties of neural networks. People hypothesize that "flat minima" with lower eigenvalues generalize better, that the Hessians of large models are very low-rank, that certain optimizers lead to flatter or sharper minima, and so on. But computing and storing the full Hessian requires memory that is quadratic in the number of parameters, which is infeasible for anything but toy models.
64
+
65
+ Iterative methods like Lanczos and power iteration can find the eigendecomposition of arbitrary linear operators given just a matrix-vector multiplication function. The Hessian-vector product (HVP) is exactly that, and it can be computed with linear memory by taking the derivative of the inner product between the gradient and the vector $v$. This library combines the HVP with iterative algorithms to compute the eigendecomposition without the quadratic memory bottleneck.
66
+
67
+ You can use this for HVP computation, the more general iterative algorithms on any linear operator you provide, or the conjunction of the two for Hessian spectrum analysis on real models, including HuggingFace and TransformerLens transformers.
68
+
69
+ ## Installation
70
+
71
+ ```bash
72
+ pip install hessian-eigenthings
73
+ ```
74
+
75
+ If you want the HuggingFace or TransformerLens helpers:
76
+
77
+ ```bash
78
+ pip install "hessian-eigenthings[transformers,transformer-lens]"
79
+ ```
80
+
81
+ ## Usage
82
+
83
+ The pattern is: build a `CurvatureOperator` from your model, run any algorithm against it. Most of the time you want the Hessian:
84
+
85
+ ```python
86
+ import torch
87
+ from torch import nn
88
+
89
+ from hessian_eigenthings.algorithms import lanczos, trace, spectral_density
90
+ from hessian_eigenthings.loss_fns import supervised_loss
91
+ from hessian_eigenthings.operators import HessianOperator
92
+
93
+ model = nn.Sequential(nn.Linear(20, 32), nn.Tanh(), nn.Linear(32, 1)).to(torch.float64)
94
+ x, y = torch.randn(128, 20, dtype=torch.float64), torch.randn(128, 1, dtype=torch.float64)
95
+ data = [(x[i:i+32], y[i:i+32]) for i in range(0, 128, 32)]
96
+
97
+ H = HessianOperator(model, data, supervised_loss(nn.functional.mse_loss))
98
+
99
+ eig = lanczos(H, k=5, seed=0)
100
+ print("top eigenvalues:", eig.eigenvalues)
101
+
102
+ t = trace(H, num_matvecs=99, seed=0) # Hutch++ by default
103
+ print(f"trace ≈ {t.estimate:.3f} ± {t.stderr:.3f}")
104
+
105
+ density = spectral_density(H, num_runs=8, lanczos_steps=40, seed=0)
106
+ # density.density is a probability density over density.grid; integrates to ~1
107
+ ```
108
+
109
+ If you'd rather have the GGN (PSD by construction, often what people mean by "the Hessian" on classification losses), use `GGNOperator`. For per-sample-gradient outer products, `EmpiricalFisherOperator`. They all share the same interface, so the algorithms above work on any of them.
110
+
111
+ There's also a finite-difference HVP path (`HessianOperator(method="finite_difference")`) for cases where double-backward is impractical. It's useful with FSDP and similar setups where the autograd graph gets detached during gradient communication.
112
+
113
+ You can restrict any operator to a subset of parameters with `param_filter`, e.g. `param_filter=match_names("blocks.*.attn.*")` to compute the Hessian of just the attention layers. Useful for per-block sharpness studies.
114
+
115
+ ## Examples
116
+
117
+ The [`examples/`](examples/) directory has three runnable scripts:
118
+
119
+ - `supervised_mlp.py`: top-k eigenvalues, Hutch++ trace, and SLQ density on a small synthetic-data MLP. No downloads, runs in seconds.
120
+ - `huggingface_tiny_gpt2.py`: full and attention-only Hessian of a HuggingFace causal LM (`sshleifer/tiny-gpt2`).
121
+ - `transformer_lens_attention_only.py`: per-block (attention vs MLP) Hessian on a TransformerLens model.
122
+
123
+ ## Documentation
124
+
125
+ Full docs at <https://noahgolmant.github.io/pytorch-hessian-eigenthings>. There are concept pages explaining the math behind each algorithm, how-to recipes for common workflows, and an auto-generated API reference. The [GGN vs Fisher vs Hessian](https://noahgolmant.github.io/pytorch-hessian-eigenthings/concepts/ggn-vs-fisher-vs-hessian/) page is worth reading before deciding which operator to instantiate. They're easy to conflate.
126
+
127
+ ## Working on the library
128
+
129
+ This project uses [`uv`](https://docs.astral.sh/uv/):
130
+
131
+ ```bash
132
+ git clone https://github.com/noahgolmant/pytorch-hessian-eigenthings
133
+ cd pytorch-hessian-eigenthings
134
+ uv sync --group dev --group docs --extra transformers --extra transformer-lens --extra curvlinops
135
+ uv run pytest
136
+ uv run mkdocs serve
137
+ ```
138
+
139
+ ## Citing this work
140
+
141
+ If you find this repo useful and would like to cite it in a publication (as [others](https://scholar.google.com/scholar?oi=bibs&hl=en&cites=18039594054930134223) have done, thank you!), here is a BibTeX entry:
142
+
143
+ ```bibtex
144
+ @misc{hessian-eigenthings,
145
+ author = {Noah Golmant and Zhewei Yao and Amir Gholami and Michael Mahoney and Joseph Gonzalez},
146
+ title = {pytorch-hessian-eigenthings: efficient PyTorch Hessian eigendecomposition},
147
+ month = oct,
148
+ year = 2018,
149
+ version = {1.0},
150
+ url = {https://github.com/noahgolmant/pytorch-hessian-eigenthings}
151
+ }
152
+ ```
153
+
154
+ ## Acknowledgements
155
+
156
+ The original 2018 implementation was written in collaboration with Zhewei Yao, Amir Gholami, Michael Mahoney, and Joseph Gonzalez at UC Berkeley's [RISELab](https://rise.cs.berkeley.edu).
157
+
158
+ The deflated power iteration routine is based on code in the [HessianFlow](https://github.com/amirgholami/HessianFlow) repository, described in: Z. Yao, A. Gholami, Q. Lei, K. Keutzer, M. Mahoney. *"Hessian-based Analysis of Large Batch Training and Robustness to Adversaries"*, NeurIPS 2018 ([arXiv:1802.08241](https://arxiv.org/abs/1802.08241)).
159
+
160
+ The accelerated stochastic power iteration is based on: C. De Sa, B. He, I. Mitliagkas, C. Ré, P. Xu. *"Accelerated Stochastic Power Iteration"*, PMLR 2017 ([arXiv:1707.02670](https://arxiv.org/abs/1707.02670)).
161
+
162
+ The v1 refresh borrows ideas from [PyHessian](https://github.com/amirgholami/PyHessian), [curvlinops](https://github.com/f-dangel/curvlinops), and [HessFormer](https://github.com/PureStrength-AI/HessFormer).
163
+
164
+ ## License
165
+
166
+ MIT.
@@ -0,0 +1,112 @@
1
+ # pytorch-hessian-eigenthings
2
+
3
+ The `hessian-eigenthings` module provides an efficient (and scalable!) way to compute the eigendecomposition of the Hessian, along with other curvature matrices like the Generalized Gauss-Newton and empirical Fisher, for an arbitrary PyTorch model. It uses PyTorch's Hessian-vector product and your choice of (a) the Lanczos method or (b) stochastic power iteration with deflation to compute the top eigenvalues and eigenvectors. There's also Hutch++ for trace estimation and Stochastic Lanczos Quadrature for the spectral density.
4
+
5
+ > **v1.0.0a1**: alpha release. The original `0.x` API has been removed; pin `hessian-eigenthings==0.0.2` if you depend on it.
6
+
7
+ ## Why use this?
8
+
9
+ The eigenvalues and eigenvectors of the Hessian have been implicated in many generalization properties of neural networks. People hypothesize that "flat minima" with lower eigenvalues generalize better, that the Hessians of large models are very low-rank, that certain optimizers lead to flatter or sharper minima, and so on. But computing and storing the full Hessian requires memory that is quadratic in the number of parameters, which is infeasible for anything but toy models.
10
+
11
+ Iterative methods like Lanczos and power iteration can find the eigendecomposition of arbitrary linear operators given just a matrix-vector multiplication function. The Hessian-vector product (HVP) is exactly that, and it can be computed with linear memory by taking the derivative of the inner product between the gradient and the vector $v$. This library combines the HVP with iterative algorithms to compute the eigendecomposition without the quadratic memory bottleneck.
12
+
13
+ You can use this for HVP computation, the more general iterative algorithms on any linear operator you provide, or the conjunction of the two for Hessian spectrum analysis on real models, including HuggingFace and TransformerLens transformers.
14
+
15
+ ## Installation
16
+
17
+ ```bash
18
+ pip install hessian-eigenthings
19
+ ```
20
+
21
+ If you want the HuggingFace or TransformerLens helpers:
22
+
23
+ ```bash
24
+ pip install "hessian-eigenthings[transformers,transformer-lens]"
25
+ ```
26
+
27
+ ## Usage
28
+
29
+ The pattern is: build a `CurvatureOperator` from your model, run any algorithm against it. Most of the time you want the Hessian:
30
+
31
+ ```python
32
+ import torch
33
+ from torch import nn
34
+
35
+ from hessian_eigenthings.algorithms import lanczos, trace, spectral_density
36
+ from hessian_eigenthings.loss_fns import supervised_loss
37
+ from hessian_eigenthings.operators import HessianOperator
38
+
39
+ model = nn.Sequential(nn.Linear(20, 32), nn.Tanh(), nn.Linear(32, 1)).to(torch.float64)
40
+ x, y = torch.randn(128, 20, dtype=torch.float64), torch.randn(128, 1, dtype=torch.float64)
41
+ data = [(x[i:i+32], y[i:i+32]) for i in range(0, 128, 32)]
42
+
43
+ H = HessianOperator(model, data, supervised_loss(nn.functional.mse_loss))
44
+
45
+ eig = lanczos(H, k=5, seed=0)
46
+ print("top eigenvalues:", eig.eigenvalues)
47
+
48
+ t = trace(H, num_matvecs=99, seed=0) # Hutch++ by default
49
+ print(f"trace ≈ {t.estimate:.3f} ± {t.stderr:.3f}")
50
+
51
+ density = spectral_density(H, num_runs=8, lanczos_steps=40, seed=0)
52
+ # density.density is a probability density over density.grid; integrates to ~1
53
+ ```
54
+
55
+ If you'd rather have the GGN (PSD by construction, often what people mean by "the Hessian" on classification losses), use `GGNOperator`. For per-sample-gradient outer products, `EmpiricalFisherOperator`. They all share the same interface, so the algorithms above work on any of them.
56
+
57
+ There's also a finite-difference HVP path (`HessianOperator(method="finite_difference")`) for cases where double-backward is impractical. It's useful with FSDP and similar setups where the autograd graph gets detached during gradient communication.
58
+
59
+ You can restrict any operator to a subset of parameters with `param_filter`, e.g. `param_filter=match_names("blocks.*.attn.*")` to compute the Hessian of just the attention layers. Useful for per-block sharpness studies.
60
+
61
+ ## Examples
62
+
63
+ The [`examples/`](examples/) directory has three runnable scripts:
64
+
65
+ - `supervised_mlp.py`: top-k eigenvalues, Hutch++ trace, and SLQ density on a small synthetic-data MLP. No downloads, runs in seconds.
66
+ - `huggingface_tiny_gpt2.py`: full and attention-only Hessian of a HuggingFace causal LM (`sshleifer/tiny-gpt2`).
67
+ - `transformer_lens_attention_only.py`: per-block (attention vs MLP) Hessian on a TransformerLens model.
68
+
69
+ ## Documentation
70
+
71
+ Full docs at <https://noahgolmant.github.io/pytorch-hessian-eigenthings>. There are concept pages explaining the math behind each algorithm, how-to recipes for common workflows, and an auto-generated API reference. The [GGN vs Fisher vs Hessian](https://noahgolmant.github.io/pytorch-hessian-eigenthings/concepts/ggn-vs-fisher-vs-hessian/) page is worth reading before deciding which operator to instantiate. They're easy to conflate.
72
+
73
+ ## Working on the library
74
+
75
+ This project uses [`uv`](https://docs.astral.sh/uv/):
76
+
77
+ ```bash
78
+ git clone https://github.com/noahgolmant/pytorch-hessian-eigenthings
79
+ cd pytorch-hessian-eigenthings
80
+ uv sync --group dev --group docs --extra transformers --extra transformer-lens --extra curvlinops
81
+ uv run pytest
82
+ uv run mkdocs serve
83
+ ```
84
+
85
+ ## Citing this work
86
+
87
+ If you find this repo useful and would like to cite it in a publication (as [others](https://scholar.google.com/scholar?oi=bibs&hl=en&cites=18039594054930134223) have done, thank you!), here is a BibTeX entry:
88
+
89
+ ```bibtex
90
+ @misc{hessian-eigenthings,
91
+ author = {Noah Golmant and Zhewei Yao and Amir Gholami and Michael Mahoney and Joseph Gonzalez},
92
+ title = {pytorch-hessian-eigenthings: efficient PyTorch Hessian eigendecomposition},
93
+ month = oct,
94
+ year = 2018,
95
+ version = {1.0},
96
+ url = {https://github.com/noahgolmant/pytorch-hessian-eigenthings}
97
+ }
98
+ ```
99
+
100
+ ## Acknowledgements
101
+
102
+ The original 2018 implementation was written in collaboration with Zhewei Yao, Amir Gholami, Michael Mahoney, and Joseph Gonzalez at UC Berkeley's [RISELab](https://rise.cs.berkeley.edu).
103
+
104
+ The deflated power iteration routine is based on code in the [HessianFlow](https://github.com/amirgholami/HessianFlow) repository, described in: Z. Yao, A. Gholami, Q. Lei, K. Keutzer, M. Mahoney. *"Hessian-based Analysis of Large Batch Training and Robustness to Adversaries"*, NeurIPS 2018 ([arXiv:1802.08241](https://arxiv.org/abs/1802.08241)).
105
+
106
+ The accelerated stochastic power iteration is based on: C. De Sa, B. He, I. Mitliagkas, C. Ré, P. Xu. *"Accelerated Stochastic Power Iteration"*, PMLR 2017 ([arXiv:1707.02670](https://arxiv.org/abs/1707.02670)).
107
+
108
+ The v1 refresh borrows ideas from [PyHessian](https://github.com/amirgholami/PyHessian), [curvlinops](https://github.com/f-dangel/curvlinops), and [HessFormer](https://github.com/PureStrength-AI/HessFormer).
109
+
110
+ ## License
111
+
112
+ MIT.