Skip to content

Analysis Tools Review & Re-contextualization #1907

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jlarson4 merged 2 commits into from
Oct 15, 2026
Merged

Analysis Tools Review & Re-contextualization #1907

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aac55b1
Demoted some analysis tools to demos or sub-methods of other features
jlarson4 Oct 15, 2026
5fe2df9
Merge remote-tracking branch 'origin/dev' into chore/analysis-tools-s…
jlarson4 Oct 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1,191 changes: 386 additions & 805 deletions demos/Backward_Lens_Demo.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

75 changes: 47 additions & 28 deletions demos/Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,9 @@
"prompts behind it.\n",
"\n",
"**This notebook loads a frozen artifact and never calls the model.** The artifact was\n",
"generated once, out-of-band, by running the benchmark's CLI entry point; see the load cell\n",
"below for the exact command. Re-running the sweep is a separate, explicit step, not something\n",
"this notebook triggers.\n"
"generated once, out-of-band, by `scripts/jacobian_lens_causal_swap_benchmark.py`; the load\n",
"cell below names the exact command. Re-running the sweep is a separate, explicit step, not\n",
"something this notebook triggers.\n"
]
},
{
Expand All @@ -59,10 +59,10 @@
"id": "8a272a14",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-24T12:26:45.620923Z",
"iopub.status.busy": "2026-09-24T12:26:45.620839Z",
"iopub.status.idle": "2026-09-24T12:26:45.623972Z",
"shell.execute_reply": "2026-09-24T12:26:45.623559Z"
"iopub.execute_input": "2026-10-14T18:27:57.152894Z",
"iopub.status.busy": "2026-10-14T18:27:57.152687Z",
"iopub.status.idle": "2026-10-14T18:27:57.158584Z",
"shell.execute_reply": "2026-10-14T18:27:57.157853Z"
}
},
"outputs": [
Expand All @@ -88,10 +88,10 @@
"id": "bc924b1c",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-24T12:26:45.625566Z",
"iopub.status.busy": "2026-09-24T12:26:45.625485Z",
"iopub.status.idle": "2026-09-24T12:26:49.109284Z",
"shell.execute_reply": "2026-09-24T12:26:49.108728Z"
"iopub.execute_input": "2026-10-14T18:27:57.160521Z",
"iopub.status.busy": "2026-10-14T18:27:57.160357Z",
"iopub.status.idle": "2026-10-14T18:27:58.846420Z",
"shell.execute_reply": "2026-10-14T18:27:58.845982Z"
}
},
"outputs": [
Expand All @@ -111,14 +111,18 @@
}
],
"source": [
"import hashlib\n",
"import json\n",
"from pathlib import Path\n",
"\n",
"import numpy as np\n",
"import pandas as pd\n",
"\n",
"from transformer_lens.tools.analysis import load_artifact, success_rate_ci\n",
"from transformer_lens.utilities import success_rate_ci\n",
"\n",
"ARTIFACT_NAME = \"jacobian_lens_causal_swap_benchmark_gpt2.json\"\n",
"SCHEMA_VERSION = 2\n",
"RESULT_BLOCKS = (\"trials\", \"excluded_baselines\", \"real_success_ci\", \"control_success_ci\")\n",
"\n",
"\n",
"def _find_artifact() -> Path:\n",
Expand All @@ -137,15 +141,30 @@
" \"directory. This notebook requires a repository checkout: it loads a frozen \"\n",
" \"artifact that is not packaged in the transformer_lens wheel. Clone the repository \"\n",
" \"and run the notebook from within it, or regenerate the artifact with \"\n",
" \"`uv run python -m transformer_lens.tools.analysis.\"\n",
" \"jacobian_lens_causal_swap_benchmark_cli`.\"\n",
" \"`uv run python scripts/jacobian_lens_causal_swap_benchmark.py`.\"\n",
" )\n",
"\n",
"\n",
"def _fingerprint(payload) -> str:\n",
" \"\"\"SHA-256 over canonical JSON, the same recipe the generation script uses.\"\"\"\n",
" canonical = json.dumps(payload, sort_keys=True, separators=(\",\", \":\"))\n",
" return hashlib.sha256(canonical.encode()).hexdigest()\n",
"\n",
"\n",
"ARTIFACT_PATH = _find_artifact()\n",
"REPO_ROOT = ARTIFACT_PATH.parents[2]\n",
"artifact = load_artifact(ARTIFACT_PATH)\n",
"artifact = json.loads(ARTIFACT_PATH.read_text())\n",
"manifest = artifact[\"protocol_manifest\"]\n",
"# Recomputed rather than trusted, so a hand-edited artifact or one from another schema fails\n",
"# here instead of rendering as if the script had produced it.\n",
"expected = {\n",
" \"schema_version\": SCHEMA_VERSION,\n",
" \"protocol_fingerprint\": _fingerprint(manifest),\n",
" \"results_fingerprint\": _fingerprint({key: artifact[key] for key in RESULT_BLOCKS}),\n",
"}\n",
"for field, value in expected.items():\n",
" if artifact[field] != value:\n",
" raise ValueError(f\"artifact {field} is {artifact[field]!r}, expected {value!r}\")\n",
"print(f\"artifact={ARTIFACT_PATH.relative_to(REPO_ROOT)}\")\n",
"print(f\"schema_version={artifact['schema_version']}\")\n",
"print(f\"protocol_fingerprint={artifact['protocol_fingerprint']}\")\n",
Expand Down Expand Up @@ -188,10 +207,10 @@
"id": "1feb4158",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-24T12:26:49.110784Z",
"iopub.status.busy": "2026-09-24T12:26:49.110543Z",
"iopub.status.idle": "2026-09-24T12:26:49.123206Z",
"shell.execute_reply": "2026-09-24T12:26:49.122757Z"
"iopub.execute_input": "2026-10-14T18:27:58.847684Z",
"iopub.status.busy": "2026-10-14T18:27:58.847570Z",
"iopub.status.idle": "2026-10-14T18:27:58.860441Z",
"shell.execute_reply": "2026-10-14T18:27:58.860108Z"
}
},
"outputs": [
Expand Down Expand Up @@ -391,10 +410,10 @@
"id": "c74c1711",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-24T12:26:49.124146Z",
"iopub.status.busy": "2026-09-24T12:26:49.124073Z",
"iopub.status.idle": "2026-09-24T12:26:49.486825Z",
"shell.execute_reply": "2026-09-24T12:26:49.486520Z"
"iopub.execute_input": "2026-10-14T18:27:58.861362Z",
"iopub.status.busy": "2026-10-14T18:27:58.861303Z",
"iopub.status.idle": "2026-10-14T18:27:59.241197Z",
"shell.execute_reply": "2026-10-14T18:27:59.240779Z"
}
},
"outputs": [
Expand Down Expand Up @@ -516,10 +535,10 @@
"id": "dedef2ac",
"metadata": {
"execution": {
"iopub.execute_input": "2026-09-24T12:26:49.488012Z",
"iopub.status.busy": "2026-09-24T12:26:49.487879Z",
"iopub.status.idle": "2026-09-24T12:26:49.494153Z",
"shell.execute_reply": "2026-09-24T12:26:49.493735Z"
"iopub.execute_input": "2026-10-14T18:27:59.242250Z",
"iopub.status.busy": "2026-10-14T18:27:59.242130Z",
"iopub.status.idle": "2026-10-14T18:27:59.248415Z",
"shell.execute_reply": "2026-10-14T18:27:59.248094Z"
}
},
"outputs": [
Expand Down Expand Up @@ -764,7 +783,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.14"
"version": "3.12.12"
}
},
"nbformat": 4,
Expand Down
8 changes: 4 additions & 4 deletions demos/Jacobian_Lens_Demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,9 +334,9 @@
"(pseudoinverse), and write back $h \\leftarrow h + \\alpha\\, V(\\sigma(c) - c)$ where $\\sigma$\n",
"exchanges the two coordinates — leaving the orthogonal complement of the activation untouched.\n",
"The paper clamps the exchanged clean-run coordinates at every token position across an\n",
"intermediate-layer band (layers 10–24 ≈ the paper's 38–92% workspace band). Unlike\n",
"`swap_hooks`, which re-reads and swaps the live coordinates at each layer,\n",
"`swap_clamp_hooks` holds every layer at the same clean-run counterfactual target."
"intermediate-layer band (layers 10–24 ≈ the paper's 38–92% workspace band). By default\n",
"`swap_hooks` re-reads and swaps the live coordinates at each layer; passing the clean cache\n",
"as `clamp_to` instead holds every layer at the same clean-run counterfactual target."
]
},
{
Expand Down Expand Up @@ -393,7 +393,7 @@
" tokens = model.to_tokens(prompt)\n",
" with torch.no_grad():\n",
" _, clean_cache = model.run_with_cache(tokens)\n",
" hooks = lens.swap_clamp_hooks(model, source, target, layers, clean_cache)\n",
" hooks = lens.swap_hooks(model, source, target, layers, clamp_to=clean_cache)\n",
" with model.hooks(fwd_hooks=hooks):\n",
" return model(tokens)[0, -1].float()\n",
"\n",
Expand Down
Loading
Loading