diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index 1a4e35ccd..819cb6539 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -5,24 +5,24 @@ on:
branches:
- main
paths:
- - '**' # Include all files by default
- - '!.devcontainer/**'
- - '!.vscode/**'
- - '!.git*'
- - '!*.md'
- - '!.github/**'
- - '.github/workflows/checks.yml' # Still include current workflow
+ - "**" # Include all files by default
+ - "!.devcontainer/**"
+ - "!.vscode/**"
+ - "!.git*"
+ - "!*.md"
+ - "!.github/**"
+ - ".github/workflows/checks.yml" # Still include current workflow
pull_request:
branches:
- main
paths:
- - '**'
- - '!.devcontainer/**'
- - '!.vscode/**'
- - '!.git*'
- - '!*.md'
- - '!.github/**'
- - '.github/workflows/checks.yml'
+ - "**"
+ - "!.devcontainer/**"
+ - "!.vscode/**"
+ - "!.git*"
+ - "!*.md"
+ - "!.github/**"
+ - ".github/workflows/checks.yml"
# Allow this workflow to be called from other workflows
workflow_call:
inputs:
@@ -36,8 +36,8 @@ permissions:
contents: write
jobs:
- checks:
- name: Checks
+ compatibility-checks:
+ name: Compatibility Checks
runs-on: ubuntu-latest
strategy:
matrix:
@@ -45,7 +45,6 @@ jobs:
- "3.8"
- "3.9"
- "3.10"
- - "3.11"
steps:
- uses: actions/checkout@v3
- name: Install Poetry
@@ -67,20 +66,20 @@ jobs:
run: |
poetry lock --check
poetry install --with dev
- - name: Check format
- run: make check-format
- - name: Unit test
+ - name: Unit Test
run: make unit-test
- - name: Docstring test
- run: make docstring-test
- # - name: Type check
- # run: poetry run mypy transformer_lens
+ - name: Acceptance Test
+ run: make acceptance-test
- name: Build check
run: poetry build
-
- # Acceptance tests are run in parallel with unit checks.
- acceptance-tests:
- name: Acceptance Tests
+ - name: Upload Coverage Report Artifact
+ uses: actions/upload-artifact@v3
+ with:
+ name: documentation
+ path: htmlcov
+
+ code-checks:
+ name: Code Checks
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v3
@@ -103,8 +102,21 @@ jobs:
run: |
poetry lock --check
poetry install --with dev
- - name: Acceptance test
- run: make acceptance-test
+ - name: Check format
+ run: make check-format
+ - name: Docstring test
+ run: make docstring-test
+ - name: Type check
+ run: poetry run mypy .
+ - name: Test Suite with Coverage Report
+ run: make coverage-report-test
+ - name: Build check
+ run: poetry build
+ - name: Upload Coverage Report Artifact
+ uses: actions/upload-artifact@v3
+ with:
+ name: test-coverage
+ path: htmlcov
notebook-checks:
name: Notebook Checks
@@ -135,3 +147,61 @@ jobs:
- name: Check Notebook Output Consistency
# Note: currently only checks notebooks we have specifically setup for this
run: make notebook-test
+
+
+ build-docs:
+ # When running on a PR, this just checks we can build the docs without errors
+ # When running on merge to main, it builds the docs and then another job deploys them
+ name: ${{ github.event_name == 'pull_request' && 'Check Build Docs' || 'Build Docs' }}
+ runs-on: ubuntu-latest
+ needs: code-checks
+ steps:
+ - uses: actions/checkout@v4
+ - name: Install Poetry
+ uses: snok/install-poetry@v1
+ - name: Set up Python
+ uses: actions/setup-python@v4
+ with:
+ python-version: "3.11"
+ cache: "poetry"
+ - name: Install pandoc
+ uses: awalsh128/cache-apt-pkgs-action@latest
+ with:
+ packages: pandoc
+ version: 1.0
+ - name: Install dependencies
+ run: poetry install --with docs
+ - name: Download Test Coverage Artifact
+ uses: actions/download-artifact@v3
+ with:
+ name: test-coverage
+ path: docs/source/coverage
+ - name: Build Docs
+ run: HF_TOKEN="$HF_TOKEN" poetry run build-docs
+ env:
+ HF_TOKEN: "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"
+ - name: Upload Docs Artifact
+ uses: actions/upload-artifact@v3
+ with:
+ name: documentation
+ path: docs/build
+
+ deploy-docs:
+ name: Deploy Docs
+ runs-on: ubuntu-latest
+ # Only run if merging a PR into main
+ if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+ needs: build-docs
+ steps:
+ - uses: actions/checkout@v4
+ - name: Download Docs Artifact
+ uses: actions/download-artifact@v3
+ with:
+ name: documentation
+ path: docs/build
+ - name: Upload to GitHub Pages
+ uses: JamesIves/github-pages-deploy-action@v4
+ with:
+ folder: docs/build
+ clean-exclude: |
+ *.*.*/
\ No newline at end of file
diff --git a/.github/workflows/gh-pages.yml b/.github/workflows/gh-pages.yml
deleted file mode 100644
index a5f6927a7..000000000
--- a/.github/workflows/gh-pages.yml
+++ /dev/null
@@ -1,61 +0,0 @@
-name: Docs
-on:
- push:
- branches:
- - main
- pull_request:
- branches:
- - '*'
-
-permissions:
- contents: write
-
-jobs:
- build-docs:
- # When running on a PR, this just checks we can build the docs without errors
- # When running on merge to main, it builds the docs and then another job deploys them
- name: ${{ github.event_name == 'pull_request' && 'Check Build Docs' || 'Build Docs' }}
- runs-on: ubuntu-latest
- steps:
- - uses: actions/checkout@v4
- - name: Install Poetry
- uses: snok/install-poetry@v1
- - name: Set up Python
- uses: actions/setup-python@v4
- with:
- python-version: "3.11"
- cache: "poetry"
- - name: Install pandoc
- uses: awalsh128/cache-apt-pkgs-action@latest
- with:
- packages: pandoc
- version: 1.0
- - name: Install dependencies
- run: poetry install --with docs
- - name: Build Docs
- run: poetry run build-docs
- - name: Upload Docs Artifact
- uses: actions/upload-artifact@v3
- with:
- name: documentation
- path: docs/build
-
- deploy-docs:
- name: Deploy Docs
- runs-on: ubuntu-latest
- # Only run if merging a PR into main
- if: github.event_name == 'push' && github.ref == 'refs/heads/main'
- needs: build-docs
- steps:
- - uses: actions/checkout@v4
- - name: Download Docs Artifact
- uses: actions/download-artifact@v3
- with:
- name: documentation
- path: docs/build
- - name: Upload to GitHub Pages
- uses: JamesIves/github-pages-deploy-action@v4
- with:
- folder: docs/build
- clean-exclude: |
- *.*.*/
diff --git a/.gitignore b/.gitignore
index 131658efd..61589404d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,3 +18,4 @@ docs/build
.Ds_Store
.pylintrc
docs/source/generated
+**.orig
diff --git a/.vscode/cspell.json b/.vscode/cspell.json
index 738113240..eea97769b 100644
--- a/.vscode/cspell.json
+++ b/.vscode/cspell.json
@@ -1,29 +1,38 @@
{
"language": "en,en-GB",
"words": [
+ "accum",
"adriĆ ",
"aengus",
+ "allclose",
"alonso",
"arange",
"argmax",
+ "argmaxy",
"autodiff",
"autoregressive",
"barez",
- "Beartype",
+ "beartype",
"belrose",
"bertsimas",
"biderman",
"bilal",
+ "bincount",
"caxis",
"checkpointed",
"chughtai",
"circuitsvis",
- "Codespaces",
+ "Codeparrot",
+ "codespaces",
"colab",
"collectstart",
"colour",
"conmy",
"cooney",
+ "crfm",
+ "cumsum",
+ "datapoint",
+ "dictmodel",
"dimitris",
"disconfirm",
"dmitrii",
@@ -31,11 +40,14 @@
"doctest",
"doctree",
"dtype",
+ "dtypes",
"einops",
"elhage",
+ "endoftext",
"eqnarray",
"esben",
"evals",
+ "explictly",
"fazl",
"firstpage",
"fspath",
@@ -49,21 +61,26 @@
"howpublished",
"huggingface",
"icml",
+ "idxs",
"imshow",
"interp",
"interpretability",
"ioannis",
"ipynb",
+ "isin",
"isort",
"janiak",
+ "Janky",
"jaxtyping",
"jett",
+ "kaiming",
"keepdim",
"kissane",
"konstas",
"kran",
"lastpage",
"layernorm",
+ "ldim",
"lieberum",
"logits",
"logsumexp",
@@ -71,31 +88,50 @@
"maxdepth",
"mingpt",
"nanda",
+ "ndarray",
+ "ndim",
"neel",
"neox",
"nitpicky",
+ "occurences",
"olah",
+ "openwebtext",
+ "overcomplete",
+ "Overriden",
"pagename",
"pauly",
"pretrained",
"probs",
+ "producting",
"pycln",
"pypi",
"pytest",
+ "randn",
+ "rdim",
+ "relu",
"resid",
+ "rprint",
"rtml",
+ "rtol",
+ "shortformer",
"softmax",
+ "softmaxing",
"solu",
+ "stas",
"templatedir",
"templatename",
"toctree",
"topk",
+ "tqdm",
"transformerlens",
+ "tril",
+ "triu",
"troitskii",
"unembed",
"unembedded",
"unembedding",
"unigram",
+ "unsqueeze",
"virtualenvs",
"visualisation",
"xaxis",
diff --git a/.vscode/extensions.json b/.vscode/extensions.json
index 0eb6d4710..2d47fb339 100644
--- a/.vscode/extensions.json
+++ b/.vscode/extensions.json
@@ -17,8 +17,9 @@
"ms-toolsai.jupyter",
"richie5um2.vscode-sort-json",
"stkb.rewrap",
+ "streetsidesoftware.code-spell-checker-british-english",
"streetsidesoftware.code-spell-checker",
- "yzhang.markdown-all-in-one",
- "streetsidesoftware.code-spell-checker-british-english"
+ "tamasfe.even-better-toml",
+ "yzhang.markdown-all-in-one"
]
}
\ No newline at end of file
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 7092a6781..63e6e310a 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,11 +1,37 @@
{
- "editor.formatOnSave": true,
- "editor.codeActionsOnSave": {
- "source.organizeImports": true
- },
"[python]": {
"editor.defaultFormatter": "ms-python.black-formatter"
},
+ "[toml]": {
+ "editor.defaultFormatter": "tamasfe.even-better-toml"
+ },
+ "editor.codeActionsOnSave": {
+ "source.organizeImports": "explicit"
+ },
+ "editor.formatOnSave": true,
+ "editor.rulers": [100],
+ "evenBetterToml.formatter.allowedBlankLines": 1,
+ "evenBetterToml.formatter.arrayAutoCollapse": true,
+ "evenBetterToml.formatter.arrayAutoExpand": true,
+ "evenBetterToml.formatter.arrayTrailingComma": true,
+ "evenBetterToml.formatter.columnWidth": 100,
+ "evenBetterToml.formatter.compactArrays": true,
+ "evenBetterToml.formatter.compactEntries": true,
+ "evenBetterToml.formatter.compactInlineTables": true,
+ "evenBetterToml.formatter.indentEntries": true,
+ "evenBetterToml.formatter.indentString": " ",
+ "evenBetterToml.formatter.indentTables": true,
+ "evenBetterToml.formatter.inlineTableExpand": true,
+ "evenBetterToml.formatter.reorderArrays": true,
+ "evenBetterToml.formatter.reorderKeys": true,
+ "evenBetterToml.formatter.trailingNewline": true,
+ "evenBetterToml.schema.enabled": true,
+ "evenBetterToml.schema.links": true,
+ "evenBetterToml.syntax.semanticTokens": false,
+ "mypy-type-checker.importStrategy": "fromEnvironment",
+ "notebook.formatOnCellExecution": true,
+ "notebook.formatOnSave.enabled": true,
+ "pylint.importStrategy": "fromEnvironment",
"python.testing.pytestArgs": [
"transformer_lens",
],
@@ -13,11 +39,9 @@
"rewrap.autoWrap.enabled": true,
"rewrap.reformat": true,
"rewrap.wrappingColumn": 100,
- "mypy-type-checker.importStrategy": "fromEnvironment",
- "pylint.importStrategy": "fromEnvironment",
- "notebook.formatOnCellExecution": true,
- "notebook.formatOnSave.enabled": true,
- "cSpell.words": [
- "accum"
+ "mypy.runUsingActiveInterpreter": true,
+ "editor.defaultFormatter": "ms-python.black-formatter",
+ "black-formatter.args": [
+ "-l 100"
],
}
\ No newline at end of file
diff --git a/README.md b/README.md
index 3f65c881d..ab7d11396 100644
--- a/README.md
+++ b/README.md
@@ -24,7 +24,7 @@ TransformerLens lets you load in 50+ different open source language models, and
activations of the model to you. You can cache any internal activation in the model, and add in
functions to edit, remove or replace these activations as the model runs.
-~~ [OCTOBER SURVEY HERE](https://forms.gle/bw7U3PfioacDtFmT8) ~~
+The library also now supports mechanistic interpretability with SAEs (sparse autoencoders)! With [HookedSAETransformer](https://colab.research.google.com/github/neelnanda-io/TransformerLens/blob/hooked-sae-transformer/demos/HookedSAETransformerDemo.ipynb), you can splice in SAEs during inference and cache + intervene on SAE activations. We recommend [SAELens](https://github.com/jbloomAus/SAELens) (built on top of TransformerLens) for training SAEs.
## Quick Start
@@ -51,6 +51,7 @@ logits, activations = model.run_with_cache("Hello World")
* [Introduction to the Library and Mech
Interp](https://arena-ch1-transformers.streamlit.app/[1.2]_Intro_to_Mech_Interp)
* [Demo of Main TransformerLens Features](https://neelnanda.io/transformer-lens-demo)
+* [Demo of HookedSAETransformer Features](https://colab.research.google.com/github/neelnanda-io/TransformerLens/blob/hooked-sae-transformer/demos/HookedSAETransformerDemo.ipynb)
## Gallery
diff --git a/demos/Attribution_Patching_Demo.ipynb b/demos/Attribution_Patching_Demo.ipynb
index 284b75412..8d8796629 100644
--- a/demos/Attribution_Patching_Demo.ipynb
+++ b/demos/Attribution_Patching_Demo.ipynb
@@ -1 +1 @@
-{"cells":[{"cell_type":"markdown","metadata":{},"source":["\n"," ![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
\n",""]},{"cell_type":"markdown","metadata":{},"source":[" # Attribution Patching Demo\n"," **Read [the accompanying blog post here](https://neelnanda.io/attribution-patching) for more context**\n"," This is an interim research report, giving a whirlwind tour of some unpublished work I did at Anthropic (credit to the then team - Chris Olah, Catherine Olsson, Nelson Elhage and Tristan Hume for help, support, and mentorship!)\n","\n"," The goal of this work is run activation patching at an industrial scale, by using gradient based attribution to approximate the technique - allow an arbitrary number of patches to be made on two forwards and a single backward pass\n","\n"," I have had less time than hoped to flesh out this investigation, but am writing up a rough investigation and comparison to standard activation patching on a few tasks to give a sense of the potential of this approach, and where it works vs falls down."]},{"cell_type":"markdown","metadata":{},"source":[" To use this notebook, go to Runtime > Change Runtime Type and select GPU as the hardware accelerator.\n","\n"," **Tips for reading this Colab:**\n"," * You can run all this code for yourself!\n"," * The graphs are interactive!\n"," * Use the table of contents pane in the sidebar to navigate\n"," * Collapse irrelevant sections with the dropdown arrows\n"," * Search the page using the search in the sidebar, not CTRL+F"]},{"cell_type":"markdown","metadata":{},"source":[" ## Setup (Ignore)"]},{"cell_type":"code","execution_count":2,"metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["Running as a Jupyter notebook - intended for development only!\n"]}],"source":["# Janky code to do different setup when run in a Colab notebook vs VSCode\n","DEBUG_MODE = False\n","try:\n"," import google.colab\n"," IN_COLAB = True\n"," print(\"Running as a Colab notebook\")\n"," %pip install transformer_lens\n"," %pip install torchtyping\n"," # Install my janky personal plotting utils\n"," %pip install git+https://github.com/neelnanda-io/neel-plotly.git\n"," # Install another version of node that makes PySvelte work way faster\n"," !curl -fsSL https://deb.nodesource.com/setup_16.x | sudo -E bash -; sudo apt-get install -y nodejs\n"," %pip install git+https://github.com/neelnanda-io/PySvelte.git\n"," # Needed for PySvelte to work, v3 came out and broke things...\n"," %pip install typeguard==2.13.3\n","except:\n"," IN_COLAB = False\n"," print(\"Running as a Jupyter notebook - intended for development only!\")\n"," from IPython import get_ipython\n","\n"," ipython = get_ipython()\n"," # Code to automatically update the HookedTransformer code as its edited without restarting the kernel\n"," ipython.magic(\"load_ext autoreload\")\n"," ipython.magic(\"autoreload 2\")"]},{"cell_type":"code","execution_count":3,"metadata":{},"outputs":[],"source":["# Plotly needs a different renderer for VSCode/Notebooks vs Colab argh\n","import plotly.io as pio\n","\n","if IN_COLAB or not DEBUG_MODE:\n"," # Thanks to annoying rendering issues, Plotly graphics will either show up in colab OR Vscode depending on the renderer - this is bad for developing demos! Thus creating a debug mode.\n"," pio.renderers.default = \"colab\"\n","else:\n"," pio.renderers.default = \"notebook_connected\""]},{"cell_type":"code","execution_count":4,"metadata":{},"outputs":[],"source":["# Import stuff\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import torch.optim as optim\n","import numpy as np\n","import einops\n","from fancy_einsum import einsum\n","import tqdm.notebook as tqdm\n","import random\n","from pathlib import Path\n","import plotly.express as px\n","from torch.utils.data import DataLoader\n","\n","from torchtyping import TensorType as TT\n","from typing import List, Union, Optional, Callable\n","from functools import partial\n","import copy\n","import itertools\n","import json\n","\n","from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer\n","import dataclasses\n","import datasets\n","from IPython.display import HTML, Markdown"]},{"cell_type":"code","execution_count":5,"metadata":{},"outputs":[],"source":["import pysvelte\n","\n","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 HookedTransformer, HookedTransformerConfig, FactoredMatrix, ActivationCache"]},{"cell_type":"markdown","metadata":{},"source":[" Plotting helper functions from a janky personal library of plotting utils. The library is not documented and I recommend against trying to read it, just use your preferred plotting library if you want to do anything non-obvious:"]},{"cell_type":"code","execution_count":6,"metadata":{},"outputs":[],"source":["from neel_plotly import line, imshow, scatter"]},{"cell_type":"code","execution_count":7,"metadata":{},"outputs":[],"source":["import transformer_lens.patching as patching"]},{"cell_type":"markdown","metadata":{},"source":[" ## IOI Patching Setup\n"," This just copies the relevant set up from Exploratory Analysis Demo, and isn't very important."]},{"cell_type":"code","execution_count":8,"metadata":{},"outputs":[{"name":"stderr","output_type":"stream","text":["Using pad_token, but it is not set yet.\n"]},{"name":"stdout","output_type":"stream","text":["Loaded pretrained model gpt2-small into HookedTransformer\n"]}],"source":["model = HookedTransformer.from_pretrained(\"gpt2-small\")\n","model.set_use_attn_result(True)"]},{"cell_type":"code","execution_count":9,"metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["Clean string 0 <|endoftext|>When John and Mary went to the shops, John gave the bag to\n","Corrupted string 0 <|endoftext|>When John and Mary went to the shops, Mary gave the bag to\n","Answer token indices tensor([[ 5335, 1757],\n"," [ 1757, 5335],\n"," [ 4186, 3700],\n"," [ 3700, 4186],\n"," [ 6035, 15686],\n"," [15686, 6035],\n"," [ 5780, 14235],\n"," [14235, 5780]], device='cuda:0')\n"]}],"source":["prompts = ['When John and Mary went to the shops, John gave the bag to', 'When John and Mary went to the shops, Mary gave the bag to', 'When Tom and James went to the park, James gave the ball to', 'When Tom and James went to the park, Tom gave the ball to', 'When Dan and Sid went to the shops, Sid gave an apple to', 'When Dan and Sid went to the shops, Dan gave an apple to', 'After Martin and Amy went to the park, Amy gave a drink to', 'After Martin and Amy went to the park, Martin gave a drink to']\n","answers = [(' Mary', ' John'), (' John', ' Mary'), (' Tom', ' James'), (' James', ' Tom'), (' Dan', ' Sid'), (' Sid', ' Dan'), (' Martin', ' Amy'), (' Amy', ' Martin')]\n","\n","clean_tokens = model.to_tokens(prompts)\n","# Swap each adjacent pair, with a hacky list comprehension\n","corrupted_tokens = clean_tokens[\n"," [(i+1 if i%2==0 else i-1) for i in range(len(clean_tokens)) ]\n"," ]\n","print(\"Clean string 0\", model.to_string(clean_tokens[0]))\n","print(\"Corrupted string 0\", model.to_string(corrupted_tokens[0]))\n","\n","answer_token_indices = torch.tensor([[model.to_single_token(answers[i][j]) for j in range(2)] for i in range(len(answers))], device=model.cfg.device)\n","print(\"Answer token indices\", answer_token_indices)"]},{"cell_type":"code","execution_count":10,"metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["Clean logit diff: 3.5519\n","Corrupted logit diff: -3.5519\n"]}],"source":["def get_logit_diff(logits, answer_token_indices=answer_token_indices):\n"," if len(logits.shape)==3:\n"," # Get final logits only\n"," logits = logits[:, -1, :]\n"," correct_logits = logits.gather(1, answer_token_indices[:, 0].unsqueeze(1))\n"," incorrect_logits = logits.gather(1, answer_token_indices[:, 1].unsqueeze(1))\n"," return (correct_logits - incorrect_logits).mean()\n","\n","clean_logits, clean_cache = model.run_with_cache(clean_tokens)\n","corrupted_logits, corrupted_cache = model.run_with_cache(corrupted_tokens)\n","\n","clean_logit_diff = get_logit_diff(clean_logits, answer_token_indices).item()\n","print(f\"Clean logit diff: {clean_logit_diff:.4f}\")\n","\n","corrupted_logit_diff = get_logit_diff(corrupted_logits, answer_token_indices).item()\n","print(f\"Corrupted logit diff: {corrupted_logit_diff:.4f}\")"]},{"cell_type":"code","execution_count":11,"metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["Clean Baseline is 1: 1.0000\n","Corrupted Baseline is 0: 0.0000\n"]}],"source":["CLEAN_BASELINE = clean_logit_diff\n","CORRUPTED_BASELINE = corrupted_logit_diff\n","def ioi_metric(logits, answer_token_indices=answer_token_indices):\n"," return (get_logit_diff(logits, answer_token_indices) - CORRUPTED_BASELINE) / (CLEAN_BASELINE - CORRUPTED_BASELINE)\n","\n","print(f\"Clean Baseline is 1: {ioi_metric(clean_logits).item():.4f}\")\n","print(f\"Corrupted Baseline is 0: {ioi_metric(corrupted_logits).item():.4f}\")"]},{"cell_type":"markdown","metadata":{},"source":[" ## Patching\n"," In the following cells, we define attribution patching and use it in various ways on the model."]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["Metric = Callable[[TT[\"batch_and_pos_dims\", \"d_model\"]], float]"]},{"cell_type":"code","execution_count":13,"metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["Clean Value: 1.0\n","Clean Activations Cached: 220\n","Clean Gradients Cached: 220\n","Corrupted Value: 0.0\n","Corrupted Activations Cached: 220\n","Corrupted Gradients Cached: 220\n"]}],"source":["filter_not_qkv_input = lambda name: \"_input\" not in name\n","def get_cache_fwd_and_bwd(model, tokens, metric):\n"," model.reset_hooks()\n"," cache = {}\n"," def forward_cache_hook(act, hook):\n"," cache[hook.name] = act.detach()\n"," model.add_hook(filter_not_qkv_input, forward_cache_hook, \"fwd\")\n","\n"," grad_cache = {}\n"," def backward_cache_hook(act, hook):\n"," grad_cache[hook.name] = act.detach()\n"," model.add_hook(filter_not_qkv_input, backward_cache_hook, \"bwd\")\n","\n"," value = metric(model(tokens))\n"," value.backward()\n"," model.reset_hooks()\n"," return value.item(), ActivationCache(cache, model), ActivationCache(grad_cache, model)\n","\n","clean_value, clean_cache, clean_grad_cache = get_cache_fwd_and_bwd(model, clean_tokens, ioi_metric)\n","print(\"Clean Value:\", clean_value)\n","print(\"Clean Activations Cached:\", len(clean_cache))\n","print(\"Clean Gradients Cached:\", len(clean_grad_cache))\n","corrupted_value, corrupted_cache, corrupted_grad_cache = get_cache_fwd_and_bwd(model, corrupted_tokens, ioi_metric)\n","print(\"Corrupted Value:\", corrupted_value)\n","print(\"Corrupted Activations Cached:\", len(corrupted_cache))\n","print(\"Corrupted Gradients Cached:\", len(corrupted_grad_cache))"]},{"cell_type":"markdown","metadata":{},"source":[" ### Attention Attribution\n"," The easiest thing to start with is to not even engage with the corrupted tokens/patching, but to look at the attribution of the attention patterns - that is, the linear approximation to what happens if you set each element of the attention pattern to zero. This, as it turns out, is a good proxy to what is going on with each head!\n"," Note that this is *not* the same as what we will later do with patching. In particular, this does not set up a careful counterfactual! It's a good tool for what's generally going on in this problem, but does not control for eg stuff that systematically boosts John > Mary in general, stuff that says \"I should activate the IOI circuit\", etc. Though using logit diff as our metric *does*\n"," Each element of the batch is independent and the metric is an average logit diff, so we can analyse each batch element independently here. We'll look at the first one, and then at the average across the whole batch (note - 4 prompts have indirect object before subject, 4 prompts have it the other way round, making the average pattern harder to interpret - I plot it over the first sequence of tokens as a mildly misleading reference).\n"," We can compare it to the interpretability in the wild diagram, and basically instantly recover most of the circuit!"]},{"cell_type":"code","execution_count":14,"metadata":{},"outputs":[],"source":["def create_attention_attr(clean_cache, clean_grad_cache) -> TT[\"batch\", \"layer\", \"head_index\", \"dest\", \"src\"]:\n"," attention_stack = torch.stack([clean_cache[\"pattern\", l] for l in range(model.cfg.n_layers)], dim=0)\n"," attention_grad_stack = torch.stack([clean_grad_cache[\"pattern\", l] for l in range(model.cfg.n_layers)], dim=0)\n"," attention_attr = attention_grad_stack * attention_stack\n"," attention_attr = einops.rearrange(attention_attr, \"layer batch head_index dest src -> batch layer head_index dest src\")\n"," return attention_attr\n","\n","attention_attr = create_attention_attr(clean_cache, clean_grad_cache)"]},{"cell_type":"code","execution_count":15,"metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["['L0H0', 'L0H1', 'L0H2', 'L0H3', 'L0H4']\n","['L0H0+', 'L0H0-', 'L0H1+', 'L0H1-', 'L0H2+']\n","['L0H0Q', 'L0H0K', 'L0H0V', 'L0H1Q', 'L0H1K']\n"]}],"source":["HEAD_NAMES = [f\"L{l}H{h}\" for l in range(model.cfg.n_layers) for h in range(model.cfg.n_heads)]\n","HEAD_NAMES_SIGNED = [f\"{name}{sign}\" for name in HEAD_NAMES for sign in [\"+\", \"-\"]]\n","HEAD_NAMES_QKV = [f\"{name}{act_name}\" for name in HEAD_NAMES for act_name in [\"Q\", \"K\", \"V\"]]\n","print(HEAD_NAMES[:5])\n","print(HEAD_NAMES_SIGNED[:5])\n","print(HEAD_NAMES_QKV[:5])"]},{"cell_type":"markdown","metadata":{},"source":[" An extremely janky way to plot the attention attribution patterns. We scale them to be in [-1, 1], split each head into a positive and negative part (so all of it is in [0, 1]), and then plot the top 20 head-halves (a head can appear twice!) by the max value of the attribution pattern."]},{"cell_type":"code","execution_count":16,"metadata":{},"outputs":[{"data":{"text/markdown":["### Attention Attribution for first sequence"],"text/plain":[""]},"metadata":{},"output_type":"display_data"},{"data":{"text/html":["\n"," \n","\n"," \n"," \n"," \n"," \n"," "],"text/plain":[""]},"metadata":{},"output_type":"display_data"},{"data":{"text/markdown":["### Summed Attention Attribution for all sequences"],"text/plain":[""]},"metadata":{},"output_type":"display_data"},{"data":{"text/html":["\n"," \n","\n"," \n"," \n"," \n"," \n"," "],"text/plain":[""]},"metadata":{},"output_type":"display_data"},{"name":"stdout","output_type":"stream","text":["Note: Plotted over first sequence for reference, but pairs have IO and S1 in different positions.\n"]}],"source":["def plot_attention_attr(attention_attr, tokens, top_k=20, index=0, title=\"\"):\n"," if len(tokens.shape)==2:\n"," tokens = tokens[index]\n"," if len(attention_attr.shape)==5:\n"," attention_attr = attention_attr[index]\n"," attention_attr_pos = attention_attr.clamp(min=-1e-5)\n"," attention_attr_neg = - attention_attr.clamp(max=1e-5)\n"," attention_attr_signed = torch.stack([attention_attr_pos, attention_attr_neg], dim=0)\n"," attention_attr_signed = einops.rearrange(attention_attr_signed, \"sign layer head_index dest src -> (layer head_index sign) dest src\")\n"," attention_attr_signed = attention_attr_signed / attention_attr_signed.max()\n"," attention_attr_indices = attention_attr_signed.max(-1).values.max(-1).values.argsort(descending=True)\n"," # print(attention_attr_indices.shape)\n"," # print(attention_attr_indices)\n"," attention_attr_signed = attention_attr_signed[attention_attr_indices, :, :]\n"," head_labels = [HEAD_NAMES_SIGNED[i.item()] for i in attention_attr_indices]\n","\n"," if title: display(Markdown(\"### \"+title))\n"," display(pysvelte.AttentionMulti(tokens=model.to_str_tokens(tokens), attention=attention_attr_signed.permute(1, 2, 0)[:, :, :top_k], head_labels=head_labels[:top_k]))\n","\n","plot_attention_attr(attention_attr, clean_tokens, index=0, title=\"Attention Attribution for first sequence\")\n","\n","plot_attention_attr(attention_attr.sum(0), clean_tokens[0], title=\"Summed Attention Attribution for all sequences\")\n","print(\"Note: Plotted over first sequence for reference, but pairs have IO and S1 in different positions.\")"]},{"cell_type":"markdown","metadata":{},"source":[" ## Attribution Patching\n"," In the following sections, I will implement various kinds of attribution patching, and then compare them to the activation patching patterns (activation patching code copied from [Exploratory Analysis Demo](https://neelnanda.io/exploratory-analysis-demo))\n"," ### Residual Stream Patching\n"," Note: We add up across both d_model and batch (Explanation).
\n"," We add up along d_model because we're taking the dot product - the derivative *is* the linear map that locally linearly approximates the metric, and so we take the dot product of our change vector with the derivative vector. Equivalent, we look at the effect of changing each coordinate independently, and then combine them by adding it up - it's linear, so this totally works.\n"," We add up across batch because we're taking the average of the metric, so each individual batch element provides `1/batch_size` of the overall effect. Because each batch element is independent of the others and no information moves between activations for different inputs, the batched version is equivalent to doing attribution patching separately for each input, and then averaging - in this second version the metric per input is *not* divided by batch_size because we don't average.