diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index e0bd51233d..47af62b2e3 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -223,8 +223,8 @@ jobs:
fail-fast: false
matrix:
notebook:
+ - "Attribution_Patching_Demo"
- "Activation_Patching_in_TL_Demo"
- # - "Attribution_Patching_Demo"
- "ARENA_Content"
- "BERT"
- "Exploratory_Analysis_Demo"
diff --git a/demos/Activation_Patching_in_TL_Demo.ipynb b/demos/Activation_Patching_in_TL_Demo.ipynb
index a848029de9..d71858f6d1 100644
--- a/demos/Activation_Patching_in_TL_Demo.ipynb
+++ b/demos/Activation_Patching_in_TL_Demo.ipynb
@@ -562,6 +562,8 @@
}
],
"source": [
+ "# 2,160 forward passes of GPT-2; too slow for CI without GPU\n",
+ "# NBVAL_SKIP\n",
"ALL_HEAD_LABELS = [f\"L{i}H{j}\" for i in range(model.cfg.n_layers) for j in range(model.cfg.n_heads)]\n",
"if DO_SLOW_RUNS:\n",
" attn_head_out_act_patch_results = patching.get_act_patch_attn_head_out_by_pos(model, corrupted_tokens, clean_cache, ioi_metric)\n",
@@ -1283,4 +1285,4 @@
},
"nbformat": 4,
"nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/demos/Attribution_Patching_Demo.ipynb b/demos/Attribution_Patching_Demo.ipynb
index bd9b6c7072..f110f59ee7 100644
--- a/demos/Attribution_Patching_Demo.ipynb
+++ b/demos/Attribution_Patching_Demo.ipynb
@@ -46,7 +46,14 @@
{
"cell_type": "code",
"execution_count": 1,
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:17:47.592530Z",
+ "iopub.status.busy": "2026-03-06T20:17:47.592314Z",
+ "iopub.status.idle": "2026-03-06T20:17:47.623795Z",
+ "shell.execute_reply": "2026-03-06T20:17:47.623529Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -54,19 +61,10 @@
"text": [
"Running as a Jupyter notebook - intended for development only!\n"
]
- },
- {
- "name": "stderr",
- "output_type": "stream",
- "text": [
- "/var/folders/m3/z6c6rcdj1rbb2jh9vqpgvxg40000gn/T/ipykernel_25358/2480103146.py:24: DeprecationWarning: `magic(...)` is deprecated since IPython 0.13 (warning added in 8.1), use run_line_magic(magic_name, parameter_s).\n",
- " ipython.magic(\"load_ext autoreload\")\n",
- "/var/folders/m3/z6c6rcdj1rbb2jh9vqpgvxg40000gn/T/ipykernel_25358/2480103146.py:25: DeprecationWarning: `magic(...)` is deprecated since IPython 0.13 (warning added in 8.1), use run_line_magic(magic_name, parameter_s).\n",
- " ipython.magic(\"autoreload 2\")\n"
- ]
}
],
"source": [
+ "# NBVAL_IGNORE_OUTPUT\n",
"# Janky code to do different setup when run in a Colab notebook vs VSCode\n",
"import os\n",
"\n",
@@ -100,8 +98,15 @@
},
{
"cell_type": "code",
- "execution_count": null,
- "metadata": {},
+ "execution_count": 2,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:17:47.641002Z",
+ "iopub.status.busy": "2026-03-06T20:17:47.640902Z",
+ "iopub.status.idle": "2026-03-06T20:17:48.042243Z",
+ "shell.execute_reply": "2026-03-06T20:17:48.041991Z"
+ }
+ },
"outputs": [],
"source": [
"# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh\n",
@@ -117,20 +122,15 @@
{
"cell_type": "code",
"execution_count": 3,
- "metadata": {},
- "outputs": [
- {
- "ename": "ModuleNotFoundError",
- "evalue": "No module named 'torchtyping'",
- "output_type": "error",
- "traceback": [
- "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
- "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)",
- "Cell \u001b[0;32mIn[3], line 15\u001b[0m\n\u001b[1;32m 12\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mplotly\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mexpress\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpx\u001b[39;00m\n\u001b[1;32m 13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DataLoader\n\u001b[0;32m---> 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorchtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m TensorType \u001b[38;5;28;01mas\u001b[39;00m TT\n\u001b[1;32m 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m List, Union, Optional, Callable\n\u001b[1;32m 17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfunctools\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m partial\n",
- "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'torchtyping'"
- ]
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:17:48.043320Z",
+ "iopub.status.busy": "2026-03-06T20:17:48.043256Z",
+ "iopub.status.idle": "2026-03-06T20:17:49.916214Z",
+ "shell.execute_reply": "2026-03-06T20:17:49.915938Z"
}
- ],
+ },
+ "outputs": [],
"source": [
"# Import stuff\n",
"import torch\n",
@@ -146,7 +146,9 @@
"import plotly.express as px\n",
"from torch.utils.data import DataLoader\n",
"\n",
- "from torchtyping import TensorType as TT\n",
+ "from typing import Any\n",
+ "class TT: # torchtyping stub\n",
+ " def __class_getitem__(cls, _): return Any\n",
"from typing import List, Union, Optional, Callable\n",
"from functools import partial\n",
"import copy\n",
@@ -156,27 +158,30 @@
"from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer\n",
"import dataclasses\n",
"import datasets\n",
- "from IPython.display import HTML, Markdown"
+ "from IPython.display import HTML, Markdown\n",
+ "\n",
+ "import circuitsvis as cv\n"
]
},
{
"cell_type": "code",
- "execution_count": 5,
- "metadata": {},
+ "execution_count": 4,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:17:49.917420Z",
+ "iopub.status.busy": "2026-03-06T20:17:49.917312Z",
+ "iopub.status.idle": "2026-03-06T20:17:50.070729Z",
+ "shell.execute_reply": "2026-03-06T20:17:50.070436Z"
+ }
+ },
"outputs": [],
"source": [
"import transformer_lens\n",
"import transformer_lens.utils as utils\n",
- "from transformer_lens.hook_points import (\n",
- " HookedRootModule,\n",
- " HookPoint,\n",
- ") # Hooking utilities\n",
"from transformer_lens import (\n",
- " HookedTransformer,\n",
- " HookedTransformerConfig,\n",
- " FactoredMatrix,\n",
" ActivationCache,\n",
- ")"
+ ")\n",
+ "from transformer_lens.model_bridge import TransformerBridge"
]
},
{
@@ -188,17 +193,38 @@
},
{
"cell_type": "code",
- "execution_count": 6,
- "metadata": {},
+ "execution_count": 5,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:17:50.071943Z",
+ "iopub.status.busy": "2026-03-06T20:17:50.071881Z",
+ "iopub.status.idle": "2026-03-06T20:17:50.092786Z",
+ "shell.execute_reply": "2026-03-06T20:17:50.092534Z"
+ }
+ },
"outputs": [],
"source": [
- "from neel_plotly import line, imshow, scatter"
+ "try:\n",
+ " from neel_plotly import line, imshow, scatter\n",
+ "except ImportError:\n",
+ " # neel_plotly is an optional visualization dependency.\n",
+ " # Define no-op stubs so patching computations still run without it.\n",
+ " def line(*args, **kwargs): pass\n",
+ " def imshow(*args, **kwargs): pass\n",
+ " def scatter(*args, **kwargs): pass\n"
]
},
{
"cell_type": "code",
- "execution_count": 7,
- "metadata": {},
+ "execution_count": 6,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:17:50.093774Z",
+ "iopub.status.busy": "2026-03-06T20:17:50.093722Z",
+ "iopub.status.idle": "2026-03-06T20:17:50.107523Z",
+ "shell.execute_reply": "2026-03-06T20:17:50.107307Z"
+ }
+ },
"outputs": [],
"source": [
"import transformer_lens.patching as patching"
@@ -214,33 +240,153 @@
},
{
"cell_type": "code",
- "execution_count": 8,
- "metadata": {},
+ "execution_count": 7,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:17:50.108499Z",
+ "iopub.status.busy": "2026-03-06T20:17:50.108449Z",
+ "iopub.status.idle": "2026-03-06T20:18:11.048982Z",
+ "shell.execute_reply": "2026-03-06T20:18:11.048621Z"
+ }
+ },
"outputs": [
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "85f0ab0d989a42989ffb354440fe312a",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "config.json: 0%| | 0.00/665 [00:00= 2 and v.shape[0] > 1:\n",
+ " cache.cache_dict[\"hook_pos_embed\"] = pe.expand(v.shape[0], -1, -1)\n",
+ " break\n",
+ "\n",
+ "\n",
+ "def _fix_head_result_cache(cache, model, is_grad=False):\n",
+ " \"\"\"Fix hook_result shape for TransformerBridge compatibility.\n",
"\n",
- " def forward_cache_hook(act, hook):\n",
- " cache[hook.name] = act.detach()\n",
+ " TransformerBridge's hook_result captures [batch, pos, d_model] (post-head-sum)\n",
+ " but stack_head_results expects [batch, pos, n_heads, d_model] (pre-sum).\n",
"\n",
- " model.add_hook(filter_not_qkv_input, forward_cache_hook, \"fwd\")\n",
+ " For forward caches: remove entries so compute_head_results recomputes from z + W_O.\n",
+ " For grad caches: expand gradient from [batch, pos, d_model] to\n",
+ " [batch, pos, n_heads, d_model] since d(loss)/d(result_h) = d(loss)/d(out) for all h.\n",
+ " \"\"\"\n",
+ " n_heads = model.cfg.n_heads\n",
+ " for layer in range(model.cfg.n_layers):\n",
+ " key = f\"blocks.{layer}.attn.hook_result\"\n",
+ " if key not in cache.cache_dict:\n",
+ " continue\n",
+ " val = cache.cache_dict[key]\n",
+ " if not isinstance(val, torch.Tensor) or val.ndim != 3:\n",
+ " continue # Already correct shape or not a tensor\n",
+ " if is_grad:\n",
+ " # Gradient of sum is identity: d(loss)/d(result_h) = d(loss)/d(out)\n",
+ " cache.cache_dict[key] = val.unsqueeze(-2).expand(\n",
+ " *val.shape[:-1], n_heads, val.shape[-1]\n",
+ " )\n",
+ " else:\n",
+ " # Remove so compute_head_results can recompute from z + W_O\n",
+ " del cache.cache_dict[key]\n",
"\n",
+ "\n",
+ "def get_cache_fwd_and_bwd(model, tokens, metric):\n",
+ " model.reset_hooks()\n",
+ " device = model.cfg.device\n",
" grad_cache = {}\n",
"\n",
" def backward_cache_hook(act, hook):\n",
@@ -417,13 +632,38 @@
"\n",
" model.add_hook(filter_not_qkv_input, backward_cache_hook, \"bwd\")\n",
"\n",
- " value = metric(model(tokens))\n",
+ " # run_with_cache handles forward caching with proper alias resolution\n",
+ " output, fwd_cache = model.run_with_cache(\n",
+ " tokens, names_filter=filter_not_qkv_input\n",
+ " )\n",
+ "\n",
+ " # Fix forward cache: move from CPU back to model device, expand pos_embed\n",
+ " _fix_cache_device_and_pos_embed(fwd_cache, device)\n",
+ "\n",
+ " value = metric(output)\n",
" value.backward()\n",
" model.reset_hooks()\n",
+ "\n",
+ " # Add alias entries for grad_cache (backward hooks use hook.name, not dict key).\n",
+ " # Only add aliases that pass our filter to avoid KeyErrors downstream.\n",
+ " for key, hp in model.hook_dict.items():\n",
+ " if hp.name != key and filter_not_qkv_input(key):\n",
+ " if hp.name in grad_cache and key not in grad_cache:\n",
+ " grad_cache[key] = grad_cache[hp.name]\n",
+ "\n",
+ " # Fix grad cache pos_embed batch dimension too\n",
+ " grad_act_cache = ActivationCache(grad_cache, model)\n",
+ " _fix_cache_device_and_pos_embed(grad_act_cache, device)\n",
+ "\n",
+ " # Fix hook_result shape: TransformerBridge captures [batch, pos, d_model]\n",
+ " # but stack_head_results expects [batch, pos, n_heads, d_model]\n",
+ " _fix_head_result_cache(fwd_cache, model, is_grad=False)\n",
+ " _fix_head_result_cache(grad_act_cache, model, is_grad=True)\n",
+ "\n",
" return (\n",
" value.item(),\n",
- " ActivationCache(cache, model),\n",
- " ActivationCache(grad_cache, model),\n",
+ " fwd_cache,\n",
+ " grad_act_cache,\n",
" )\n",
"\n",
"\n",
@@ -438,7 +678,7 @@
")\n",
"print(\"Corrupted Value:\", corrupted_value)\n",
"print(\"Corrupted Activations Cached:\", len(corrupted_cache))\n",
- "print(\"Corrupted Gradients Cached:\", len(corrupted_grad_cache))"
+ "print(\"Corrupted Gradients Cached:\", len(corrupted_grad_cache))\n"
]
},
{
@@ -454,8 +694,15 @@
},
{
"cell_type": "code",
- "execution_count": 14,
- "metadata": {},
+ "execution_count": 13,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:18:11.625801Z",
+ "iopub.status.busy": "2026-03-06T20:18:11.625732Z",
+ "iopub.status.idle": "2026-03-06T20:18:11.643613Z",
+ "shell.execute_reply": "2026-03-06T20:18:11.643388Z"
+ }
+ },
"outputs": [],
"source": [
"def create_attention_attr(\n",
@@ -480,8 +727,15 @@
},
{
"cell_type": "code",
- "execution_count": 15,
- "metadata": {},
+ "execution_count": 14,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:18:11.644629Z",
+ "iopub.status.busy": "2026-03-06T20:18:11.644574Z",
+ "iopub.status.idle": "2026-03-06T20:18:11.658123Z",
+ "shell.execute_reply": "2026-03-06T20:18:11.657927Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -515,8 +769,15 @@
},
{
"cell_type": "code",
- "execution_count": 16,
- "metadata": {},
+ "execution_count": 15,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:18:11.659129Z",
+ "iopub.status.busy": "2026-03-06T20:18:11.659081Z",
+ "iopub.status.idle": "2026-03-06T20:18:12.156846Z",
+ "shell.execute_reply": "2026-03-06T20:18:12.156025Z"
+ }
+ },
"outputs": [
{
"data": {
@@ -533,96 +794,18 @@
{
"data": {
"text/html": [
- "\n",
- " \n",
- "\n",
- " \n",
- " \n",
- " \n",
- " \n",
- " "
+ "\n",
+ " "
],
"text/plain": [
- ""
+ ""
]
},
"metadata": {},
@@ -643,96 +826,18 @@
{
"data": {
"text/html": [
- "\n",
- " \n",
- "\n",
- " \n",
- " \n",
- " \n",
- " \n",
- " "
+ "\n",
+ " "
],
"text/plain": [
- ""
+ ""
]
},
"metadata": {},
@@ -771,10 +876,10 @@
" if title:\n",
" display(Markdown(\"### \" + title))\n",
" display(\n",
- " pysvelte.AttentionMulti(\n",
+ " cv.attention.attention_heads(\n",
" tokens=model.to_str_tokens(tokens),\n",
- " attention=attention_attr_signed.permute(1, 2, 0)[:, :, :top_k],\n",
- " head_labels=head_labels[:top_k],\n",
+ " attention=attention_attr_signed[:top_k],\n",
+ " attention_head_names=head_labels[:top_k],\n",
" )\n",
" )\n",
"\n",
@@ -810,8 +915,15 @@
},
{
"cell_type": "code",
- "execution_count": 17,
- "metadata": {},
+ "execution_count": 16,
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-03-06T20:18:12.159003Z",
+ "iopub.status.busy": "2026-03-06T20:18:12.158829Z",
+ "iopub.status.idle": "2026-03-06T20:18:12.248543Z",
+ "shell.execute_reply": "2026-03-06T20:18:12.248254Z"
+ }
+ },
"outputs": [
{
"data": {
@@ -820,9 +932,9 @@
"\n",
"\n",
" \n",
+ " }) }; \n",
"\n",
"