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.
- hessian_eigenthings-1.0.0a1/.claude/hooks/check-docs-affected.sh +2 -0
- hessian_eigenthings-1.0.0a1/.claude/hooks/check_docs_affected.py +86 -0
- hessian_eigenthings-1.0.0a1/.claude/settings.json +15 -0
- hessian_eigenthings-1.0.0a1/.github/workflows/ci.yml +98 -0
- hessian_eigenthings-1.0.0a1/.github/workflows/docs.yml +25 -0
- hessian_eigenthings-1.0.0a1/.gitignore +104 -0
- hessian_eigenthings-1.0.0a1/LICENSE +21 -0
- hessian_eigenthings-1.0.0a1/PKG-INFO +166 -0
- hessian_eigenthings-1.0.0a1/README.md +112 -0
- hessian_eigenthings-1.0.0a1/docs/concepts/ggn-vs-fisher-vs-hessian.md +76 -0
- hessian_eigenthings-1.0.0a1/docs/concepts/numerical-stability.md +56 -0
- hessian_eigenthings-1.0.0a1/docs/concepts/spectral-density.md +78 -0
- hessian_eigenthings-1.0.0a1/docs/concepts/top-k-eigenvalues.md +81 -0
- hessian_eigenthings-1.0.0a1/docs/concepts/trace-estimation.md +67 -0
- hessian_eigenthings-1.0.0a1/docs/concepts/what-is-the-hessian.md +44 -0
- hessian_eigenthings-1.0.0a1/docs/concepts/why-hvp-not-full-h.md +46 -0
- hessian_eigenthings-1.0.0a1/docs/getting-started/installation.md +20 -0
- hessian_eigenthings-1.0.0a1/docs/getting-started/quickstart.md +51 -0
- hessian_eigenthings-1.0.0a1/docs/getting-started/transformers-quickstart.md +50 -0
- hessian_eigenthings-1.0.0a1/docs/how-to/analyze-a-huggingface-model.md +65 -0
- hessian_eigenthings-1.0.0a1/docs/how-to/analyze-with-transformer-lens.md +62 -0
- hessian_eigenthings-1.0.0a1/docs/how-to/custom-curvature-operators.md +119 -0
- hessian_eigenthings-1.0.0a1/docs/how-to/custom-loss-functions.md +122 -0
- hessian_eigenthings-1.0.0a1/docs/how-to/distributed-ddp.md +59 -0
- hessian_eigenthings-1.0.0a1/docs/how-to/per-layer-hessian.md +83 -0
- hessian_eigenthings-1.0.0a1/docs/index.md +12 -0
- hessian_eigenthings-1.0.0a1/docs/reference/algorithms.md +3 -0
- hessian_eigenthings-1.0.0a1/docs/reference/api.md +3 -0
- hessian_eigenthings-1.0.0a1/docs/reference/loss_fns.md +3 -0
- hessian_eigenthings-1.0.0a1/examples/README.md +17 -0
- hessian_eigenthings-1.0.0a1/examples/huggingface_tiny_gpt2.py +66 -0
- hessian_eigenthings-1.0.0a1/examples/supervised_mlp.py +66 -0
- hessian_eigenthings-1.0.0a1/examples/transformer_lens_attention_only.py +64 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/__init__.py +3 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/__init__.py +35 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/lanczos.py +184 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/power_iteration.py +151 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/result.py +16 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/spectral_density.py +135 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/algorithms/trace.py +151 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/batching.py +61 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/linalg/__init__.py +3 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/linalg/backend.py +69 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/__init__.py +31 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/huggingface.py +59 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/standard.py +51 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/loss_fns/transformer_lens.py +47 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/__init__.py +19 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/base.py +70 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/distributed/__init__.py +3 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/distributed/ddp.py +85 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/fisher.py +149 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/ggn.py +147 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/operators/hessian.py +220 -0
- hessian_eigenthings-1.0.0a1/hessian_eigenthings/param_utils.py +75 -0
- hessian_eigenthings-1.0.0a1/mkdocs.yml +93 -0
- hessian_eigenthings-1.0.0a1/pyproject.toml +111 -0
- hessian_eigenthings-1.0.0a1/tests/conftest.py +27 -0
- hessian_eigenthings-1.0.0a1/tests/cross_library/__init__.py +0 -0
- hessian_eigenthings-1.0.0a1/tests/cross_library/test_against_curvlinops.py +213 -0
- hessian_eigenthings-1.0.0a1/tests/test_batching.py +63 -0
- hessian_eigenthings-1.0.0a1/tests/test_ddp_hessian.py +88 -0
- hessian_eigenthings-1.0.0a1/tests/test_eigendecomp_on_hessian.py +86 -0
- hessian_eigenthings-1.0.0a1/tests/test_fisher_operator.py +101 -0
- hessian_eigenthings-1.0.0a1/tests/test_ggn_operator.py +162 -0
- hessian_eigenthings-1.0.0a1/tests/test_hessian_finite_difference.py +134 -0
- hessian_eigenthings-1.0.0a1/tests/test_hessian_operator.py +182 -0
- hessian_eigenthings-1.0.0a1/tests/test_lanczos_random_matrix.py +124 -0
- hessian_eigenthings-1.0.0a1/tests/test_linalg_backend.py +91 -0
- hessian_eigenthings-1.0.0a1/tests/test_loss_fns.py +107 -0
- hessian_eigenthings-1.0.0a1/tests/test_param_utils.py +91 -0
- hessian_eigenthings-1.0.0a1/tests/test_power_iteration_random_matrix.py +96 -0
- hessian_eigenthings-1.0.0a1/tests/test_smoke.py +5 -0
- hessian_eigenthings-1.0.0a1/tests/test_spectral_density.py +150 -0
- hessian_eigenthings-1.0.0a1/tests/test_trace.py +141 -0
- hessian_eigenthings-1.0.0a1/uv.lock +5090 -0
|
@@ -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,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.
|