tutorials/Hooked_SAE_Transformer_Demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Hooked_SAE_Transformer_Demo.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HookedSAETransformer Demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HookedSAETransformer is a lightweight extension of HookedTransformer that allows you to \"splice in\" Sparse Autoencoders. This makes it easy to do exploratory analysis such as: running inference with SAEs attached, caching SAE feature activations, and intervening on SAE activations with hooks.\n",
    "\n",
    "I (Connor Kissane) implemented this to accelerate research on [Attention SAEs](https://www.lesswrong.com/posts/DtdzGwFh9dCfsekZZ/sparse-autoencoders-work-on-attention-layer-outputs) based on suggestions from Arthur Conmy and Neel Nanda, and found that it was well worth the time and effort. I hope other researchers will also find the library useful! This notebook demonstrates how it works and how to use it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Janky code to do different setup when run in a Colab notebook vs VSCode\n",
    "DEVELOPMENT_MODE = False\n",
    "try:\n",
    "    import google.colab\n",
    "    IN_COLAB = True\n",
    "    print(\"Running as a Colab notebook\")\n",
    "    %pip install git+https://github.com/jbloomAus/SAELens\n",
    "  \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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformer_lens.utils as utils\n",
    "\n",
    "import plotly.express as px\n",
    "import tqdm\n",
    "from functools import partial\n",
    "import einops\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "update_layout_set = {\n",
    "    \"xaxis_range\", \"yaxis_range\", \"hovermode\", \"xaxis_title\", \"yaxis_title\", \"colorbar\", \"colorscale\", \"coloraxis\",\n",
    "     \"title_x\", \"bargap\", \"bargroupgap\", \"xaxis_tickformat\", \"yaxis_tickformat\", \"title_y\", \"legend_title_text\", \"xaxis_showgrid\",\n",
    "     \"xaxis_gridwidth\", \"xaxis_gridcolor\", \"yaxis_showgrid\", \"yaxis_gridwidth\"\n",
    "}\n",
    "\n",
    "def imshow(tensor, renderer=None, xaxis=\"\", yaxis=\"\", **kwargs):\n",
    "    if isinstance(tensor, list):\n",
    "        tensor = torch.stack(tensor)\n",
    "    kwargs_post = {k: v for k, v in kwargs.items() if k in update_layout_set}\n",
    "    kwargs_pre = {k: v for k, v in kwargs.items() if k not in update_layout_set}\n",
    "    if \"facet_labels\" in kwargs_pre:\n",
    "        facet_labels = kwargs_pre.pop(\"facet_labels\")\n",
    "    else:\n",
    "        facet_labels = None\n",
    "    if \"color_continuous_scale\" not in kwargs_pre:\n",
    "        kwargs_pre[\"color_continuous_scale\"] = \"RdBu\"\n",
    "    fig = px.imshow(utils.to_numpy(tensor), color_continuous_midpoint=0.0,labels={\"x\":xaxis, \"y\":yaxis}, **kwargs_pre).update_layout(**kwargs_post)\n",
    "    if facet_labels:\n",
    "        for i, label in enumerate(facet_labels):\n",
    "            fig.layout.annotations[i]['text'] = label\n",
    "\n",
    "    fig.show(renderer)\n",
    "\n",
    "def scatter(x, y, xaxis=\"\", yaxis=\"\", caxis=\"\", renderer=None, return_fig=False, **kwargs):\n",
    "    x = utils.to_numpy(x)\n",
    "    y = utils.to_numpy(y)\n",
    "    fig = px.scatter(y=y, x=x, labels={\"x\":xaxis, \"y\":yaxis, \"color\":caxis}, **kwargs)\n",
    "    if return_fig:\n",
    "        return fig\n",
    "    fig.show(renderer)\n",
    "\n",
    "from typing import List\n",
    "def show_avg_logit_diffs(x_axis: List[str], per_prompt_logit_diffs: List[torch.tensor]):\n",
    "\n",
    "\n",
    "    y_data = [per_prompt_logit_diff.mean().item() for per_prompt_logit_diff in per_prompt_logit_diffs]\n",
    "    error_y_data = [per_prompt_logit_diff.std().item() for per_prompt_logit_diff in per_prompt_logit_diffs] \n",
    "\n",
    "    fig = go.Figure(data=[go.Bar(\n",
    "        x=x_axis,\n",
    "        y=y_data,\n",
    "        error_y=dict(\n",
    "            type='data',  # specifies that the actual values are given\n",
    "            array=error_y_data,  # the magnitudes of the errors\n",
    "            visible=True  # make error bars visible\n",
    "        ),\n",
    "    )])\n",
    "\n",
    "    # Customize layout\n",
    "    fig.update_layout(title_text=f'Logit Diff after Interventions',\n",
    "                    xaxis_title_text='Intervention',\n",
    "                    yaxis_title_text='Logit diff',\n",
    "                    plot_bgcolor='white')\n",
    "\n",
    "    # Show the figure\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "# elif torch.backends.mps.is_available():\n",
    "#     device = \"mps\"\n",
    "else: \n",
    "    device = \"cpu\"\n",
    "torch.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading and Running Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like a [HookedTransformer](https://TransformerLensOrg.github.io/TransformerLens/generated/demos/Main_Demo.html#Loading-and-Running-Models), we can load in any model that's supported in TransformerLens with the `HookedSAETransformer.from_pretrained(MODEL_NAME)`. In this demo we'll use GPT-2 small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sae_lens import HookedSAETransformer\n",
    "model: HookedSAETransformer = HookedSAETransformer.from_pretrained(\"gpt2-small\").to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default HookedSAETransformer will behave exactly like a HookedTransformer. We'll explore the main features of HookedSAETransformer on the classic IOI task, so let's first sanity check that GPT2-small can do the IOI task without any SAEs attached:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_format = [\n",
    "    \"When John and Mary went to the shops,{} gave the bag to\",\n",
    "    \"When Tom and James went to the park,{} gave the ball to\",\n",
    "    \"When Dan and Sid went to the shops,{} gave an apple to\",\n",
    "    \"After Martin and Amy went to the park,{} gave a drink to\",\n",
    "]\n",
    "names = [\n",
    "    (\" John\", \" Mary\",),\n",
    "    (\" Tom\", \" James\"),\n",
    "    (\" Dan\", \" Sid\"),\n",
    "    (\" Martin\", \" Amy\"),\n",
    "]\n",
    "# List of prompts\n",
    "prompts = []\n",
    "# List of answers, in the format (correct, incorrect)\n",
    "answers = []\n",
    "# List of the token (ie an integer) corresponding to each answer, in the format (correct_token, incorrect_token)\n",
    "answer_tokens = []\n",
    "for i in range(len(prompt_format)):\n",
    "    for j in range(2):\n",
    "        answers.append((names[i][j], names[i][1 - j]))\n",
    "        answer_tokens.append(\n",
    "            (\n",
    "                model.to_single_token(answers[-1][0]),\n",
    "                model.to_single_token(answers[-1][1]),\n",
    "            )\n",
    "        )\n",
    "        # Insert the *incorrect* answer to the prompt, making the correct answer the indirect object.\n",
    "        prompts.append(prompt_format[i].format(answers[-1][1]))\n",
    "answer_tokens = torch.tensor(answer_tokens).to(device)\n",
    "print(prompts)\n",
    "print(answers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def logits_to_ave_logit_diff(logits, answer_tokens, per_prompt=False):\n",
    "    # Only the final logits are relevant for the answer\n",
    "    final_logits = logits[:, -1, :]\n",
    "    answer_logits = final_logits.gather(dim=-1, index=answer_tokens)\n",
    "    answer_logit_diff = answer_logits[:, 0] - answer_logits[:, 1]\n",
    "    if per_prompt:\n",
    "        return answer_logit_diff\n",
    "    else:\n",
    "        return answer_logit_diff.mean()\n",
    "    \n",
    "tokens = model.to_tokens(prompts, prepend_bos=True)\n",
    "original_logits, cache = model.run_with_cache(tokens)\n",
    "original_average_logit_diff = logits_to_ave_logit_diff(original_logits, answer_tokens)\n",
    "print(f\"Original average logit diff: {original_average_logit_diff}\")\n",
    "original_per_prompt_logit_diff = logits_to_ave_logit_diff(original_logits, answer_tokens, per_prompt=True)\n",
    "print(f\"Original per prompt logit diff: {original_per_prompt_logit_diff}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HookedSAEs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to use the key features of HookedSAETransformer, we first need to load in SAEs.\n",
    "\n",
    "SAE is a class we've implemented to have TransformerLens hooks around the SAE activations. While we will use it out of the box, it is designed to be hackable: you can copy and paste the SAE class into a notebook and completely change the architecture / hook names, and as long as it reconstructs the activations, it should still work.\n",
    "\n",
    "You can initialize a SAE with an SAEConfig, but note you'll likely have to write some basic conversion code to match configs / state dicts to the SAE class when loading in an open sourced SAE (eg from HuggingFace). For convenience, we've developed a `SAE.from_pretrained` to automatically load certain open sourced SAEs. We'll use our GPT-2 Small [Attention SAEs](https://www.alignmentforum.org/posts/FSTRedtjuHa4Gfdbr/attention-saes-scale-to-gpt-2-small) to demonstrate. We'll load in all of our attention SAEs with `SAE.from_pretrained`, and store them in a dictionary that maps each hook_name (str) to the corresponding HookedSAE.\n",
    "\n",
    "<details>\n",
    "\n",
    "Later we'll show how to add SAEs to the HookedSAETransformer (replacing model activations with their SAE reconstructions). When you add an SAE, HookedSAETransformer just treats this a black box that takes some activation as an input, and outputs a tensor of the same shape. \n",
    "\n",
    "With this in mind, the SAE is designed to be simple and hackable. Think of it as a convenient default class that you can copy and edit. As long as it takes a TransformerLens activation as input, and outputs a tensor of the same shape, you should be able to add it to your HookedSAETransformer.\n",
    "\n",
    "You probably don't even need to use the SAE class, although it's recommended. The sae can be any pytorch module that takes in some activation at hook_name and outputs a tensor of the same shape. The two assumptions that HookedSAETransformer makes when adding SAEs are:\n",
    "1. The SAE class has a cfg attribute, sae.cfg.hook_name (str), for the activation that the SAE was trained to reconstruct (in TransformerLens notation e.g. 'blocks.0.attn.hook_z')\n",
    "2. The SAE takes that activation as input, and outputs a tensor of the same shape.\n",
    "\n",
    "The main benefit of HookedSAE is that it's a subclass of HookedRootModule, so we can add hooks to SAE activations. This makes it easy to leverage existing TransformerLens functionality like run_with_cache and run_with_hooks with SAEs.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sae_lens import SAE\n",
    "\n",
    "hook_name_to_sae = {}\n",
    "for layer in tqdm.tqdm(range(12)):\n",
    "    sae, cfg_dict, _ = SAE.from_pretrained(\n",
    "        \"gpt2-small-hook-z-kk\",\n",
    "        f\"blocks.{layer}.hook_z\",\n",
    "        device=device,\n",
    "    )\n",
    "    hook_name_to_sae[sae.cfg.hook_name] = sae\n",
    "    \n",
    "\n",
    "print(hook_name_to_sae.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run with SAEs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The key feature of HookedSAETransformer is being able to \"splice in\" SAEs, replacing model activations with their SAE reconstructions. \n",
    "\n",
    "To run a forward pass with SAEs attached use `model.run_with_saes(tokens, saes=saes)`, where saes is a list of SAEs that you want to add for just this forward pass. These will be reset immediately after the forward pass, returning the model to its original state.\n",
    "\n",
    "I expect this to be particularly useful for evaluating SAEs (eg [Gurnee](https://www.alignmentforum.org/posts/rZPiuFxESMxCDHe4B/sae-reconstruction-errors-are-empirically-pathological)), including evaluating how SAE reconstructions affect the models ability to perform certain tasks (eg [Makelov et al.](https://openreview.net/forum?id=MHIX9H8aYF&referrer=%5Bthe%20profile%20of%20Neel%20Nanda%5D(%2Fprofile%3Fid%3D~Neel_Nanda1)))\n",
    "\n",
    "To demonstrate, let's use `run_with_saes` to evaluate many combinations of SAEs on different cross sections of the IOI circuit.\n",
    "\n",
    "<details>\n",
    "\n",
    "Under the hood, TransformerLens already wraps activations with a HookPoint object. HookPoint is a dummy pytorch module that acts as an identity function by default, and is only used to access the activation with PyTorch hooks. When you run_with_saes, HookedSAETransformer temporarily replaces these HookPoints with the given SAEs, which take the activation as input and replace it with the SAE output (the reconstructed activation) during the forward pass. \n",
    "\n",
    "Since SAE is a subclass of HookedRootModule, we also are able to add PyTorch hooks to the corresponding SAE activations, as we'll use later.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_layers = [[0, 3], [2, 4], [5,6], [7, 8], [9, 10, 11]]\n",
    "x_axis = ['Clean Baseline']\n",
    "per_prompt_logit_diffs = [\n",
    "    original_per_prompt_logit_diff, \n",
    "]\n",
    "\n",
    "for layers in all_layers:\n",
    "    hooked_saes = [hook_name_to_sae[utils.get_act_name('z', layer)] for layer in layers]\n",
    "    logits_with_saes = model.run_with_saes(tokens, saes=hooked_saes)\n",
    "    average_logit_diff_with_saes = logits_to_ave_logit_diff(logits_with_saes, answer_tokens)\n",
    "    per_prompt_diff_with_saes = logits_to_ave_logit_diff(logits_with_saes, answer_tokens, per_prompt=True)\n",
    "    \n",
    "    x_axis.append(f\"With SAEs L{layers}\")\n",
    "    per_prompt_logit_diffs.append(per_prompt_diff_with_saes)\n",
    "\n",
    "show_avg_logit_diffs(x_axis, per_prompt_logit_diffs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run with cache (with SAEs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We often want to see what SAE features are active on a given prompt. With HookedSAETransformer, you can cache SAE activations (and all the other standard activations) with `logits, cache = model.run_with_cache_with_saes(tokens, saes=saes)`. Just as `run_with_saes` is a wapper around the standard forward pass, `run_with_cache_with_saes` is a wrapper around `run_with_cache`, and will also only add these saes for one forward pass before returning the model to its original state. \n",
    "\n",
    "To access SAE activations from the cache, the corresponding hook names will generally be the HookedTransformer hook_name (eg blocks.5.attn.hook_z) + the SAE hooked name preceeded by a period (eg .hook_sae_acts_post).\n",
    "\n",
    "`run_with_cache_with_saes` makes it easy to explore which SAE features are active across any input. Let's explore the active features at the S2 position for our L5 Attention SAE across all of our IOI prompts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer, s2_pos = 5, 10\n",
    "saes = [hook_name_to_sae[utils.get_act_name('z', layer)]]\n",
    "_, cache = model.run_with_cache_with_saes(tokens, saes=saes)\n",
    "sae_acts = cache[utils.get_act_name('z', layer) + \".hook_sae_acts_post\"][:, s2_pos, :]\n",
    "live_feature_mask = sae_acts > 0\n",
    "live_feature_union = live_feature_mask.any(dim=0)\n",
    "\n",
    "imshow(\n",
    "    sae_acts[:, live_feature_union],\n",
    "    title = \"Activations of Live SAE features at L5 S2 position per prompt\",\n",
    "    xaxis=\"Feature Id\", yaxis=\"Prompt\",\n",
    "    x=list(map(str, live_feature_union.nonzero().flatten().tolist())),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could then interpret some of the commonly activating features, like 7515, using [neuronpedia](https://www.neuronpedia.org/gpt2-small/5-att-kk/7515)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run with Hooks (with SAEs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HookedSAETransformer also allows you to intervene on SAE activations with `model.run_with_hooks_with_saes(tokens, saes=saes, fwd_hooks=fwd_hooks)`. This works exactly like the standard TransformerLens `run_with_hooks`, with the added benefit that we can now intervene on SAE activations from the SAEs that we splice in. Along the same lines as `run_with_saes` and `run_with_cache_with_saes`, this will only temporarily add SAEs before returning the model to it's original state. \n",
    "\n",
    "I expect this to be useful when doing circuit analysis with SAEs. To demonstrate, let's zero ablate individual layer 5 attention SAE features to localize causally important features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ablate_sae_feature(sae_acts, hook, pos, feature_id):\n",
    "    if pos is None:\n",
    "        sae_acts[:, :, feature_id] = 0.\n",
    "    else:\n",
    "        sae_acts[:, pos, feature_id] = 0.\n",
    "    return sae_acts\n",
    "\n",
    "layer = 5\n",
    "sae = hook_name_to_sae[utils.get_act_name('z', layer)]\n",
    "\n",
    "logits_with_saes = model.run_with_saes(tokens, saes=sae)\n",
    "clean_sae_baseline_per_prompt = logits_to_ave_logit_diff(logits_with_saes, answer_tokens, per_prompt=True)\n",
    "\n",
    "all_live_features = torch.arange(sae.cfg.d_sae)[live_feature_union.cpu()]\n",
    "\n",
    "causal_effects = torch.zeros((len(prompts), all_live_features.shape[0]))\n",
    "fid_to_idx = {fid.item(): idx for idx, fid in enumerate(all_live_features)}\n",
    "\n",
    "\n",
    "abl_layer, abl_pos  = 5, 10\n",
    "for feature_id in tqdm.tqdm(all_live_features):\n",
    "    feature_id = feature_id.item()\n",
    "    abl_feature_logits = model.run_with_hooks_with_saes(\n",
    "        tokens,\n",
    "        saes=sae,\n",
    "        fwd_hooks=[(utils.get_act_name('z', abl_layer) + \".hook_sae_acts_post\", partial(ablate_sae_feature, pos=abl_pos, feature_id=feature_id))]\n",
    "    ) # [batch, seq, vocab]\n",
    "    \n",
    "    abl_feature_logit_diff = logits_to_ave_logit_diff(abl_feature_logits, answer_tokens, per_prompt=True) # [batch]\n",
    "    causal_effects[:, fid_to_idx[feature_id]] = abl_feature_logit_diff - clean_sae_baseline_per_prompt\n",
    "\n",
    "\n",
    "imshow(causal_effects, title=f\"Change in logit diff when ablating L{abl_layer} SAE features for all prompts at pos {abl_pos}\", xaxis=\"Feature Idx\", yaxis=\"Prompt Idx\", x=list(map(str, all_live_features.tolist())))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although it's not super clean, we see a few features stand out, where ablating them causes a nontrivial drop in logit diff on multiple prompts: 7515 and 27535 for BABA prompts, with 44256 for ABBA prompts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Add SAEs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the `run_with_saes` family of methods are great for evaluating SAEs and exploratory analysis, you may want to permanently attach SAEs to your model. You can attach SAEs to any activation with `model.add_sae(sae)`, where sae is an SAE. \n",
    "\n",
    "When you add an SAE, it gets stored in `model.acts_to_saes`, a dictionary that maps the activation name to the SAE that is attached. The main benefit of permanently adding SAEs is that we can now just run the model like a normal HookedTransformer (with `forward`, `run_with_cache`, `run_with_hooks`), but some activations will be replaced with the reconstructed activations from the corresponding SAEs.\n",
    "\n",
    "I expect this to be most useful when you've already identified a good set of SAEs that you want to use for interpretability, and don't feel like passing in a massive list of saes for every forward pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Attached SAEs before add_sae\", model.acts_to_saes)\n",
    "layer = 5\n",
    "sae = hook_name_to_sae[utils.get_act_name('z', layer)]\n",
    "model.add_sae(sae)\n",
    "print(\"Attached SAEs after add_sae\", model.acts_to_saes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can just call the standard HookedTransformer forward, and the sae that we added will automatically be spliced in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits_with_saes = model(tokens)\n",
    "assert not torch.allclose(original_logits, logits_with_saes, atol=1e-4)\n",
    "\n",
    "average_logit_diff_with_saes = logits_to_ave_logit_diff(logits_with_saes, answer_tokens)\n",
    "print(f\"Average logit diff with SAEs: {average_logit_diff_with_saes}\")\n",
    "per_prompt_diff_with_saes = logits_to_ave_logit_diff(logits_with_saes, answer_tokens, per_prompt=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run with cache"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, we can also use `logits, cache = model.run_with_cache(tokens)` directly to cache SAE activations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = 5\n",
    "_, cache = model.run_with_cache(tokens)\n",
    "s2_pos = 10\n",
    "sae_acts = cache[utils.get_act_name('z', layer) + \".hook_sae_acts_post\"][:, s2_pos, :]\n",
    "\n",
    "live_feature_mask = sae_acts > 0\n",
    "live_feature_union = live_feature_mask.any(dim=0)\n",
    "\n",
    "imshow(\n",
    "    sae_acts[:, live_feature_union],\n",
    "    title = \"Activations of Live SAE features at L5 S2 position per prompt\",\n",
    "    xaxis=\"Feature Id\", yaxis=\"Prompt\",\n",
    "    x=list(map(str, live_feature_union.nonzero().flatten().tolist())),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run with hooks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we can also use `run_with_hooks` and intervene on the added SAE's activations. To show a more complicated intervention, we'll try path patching the feature from the S-inhibition head's value vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_use_split_qkv_input(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def path_patch_v_input(v_input, hook, feature_dirs, pos, head_index):\n",
    "    v_input[:, pos, head_index, :] = v_input[:, pos, head_index, :] - feature_dirs\n",
    "    return v_input\n",
    "\n",
    "\n",
    "s_inhib_heads = [(7, 3), (7, 9), (8,6), (8,10)]\n",
    "\n",
    "results = torch.zeros(tokens.shape[0], all_live_features.shape[0])\n",
    "\n",
    "W_O_cat = einops.rearrange(\n",
    "    model.W_O,\n",
    "    \"n_layers n_heads d_head d_model -> n_layers (n_heads d_head) d_model\"\n",
    ")\n",
    "\n",
    "for feature_id in tqdm.tqdm(all_live_features):\n",
    "    feature_id = feature_id.item()\n",
    "    feature_acts = cache[utils.get_act_name('z', abl_layer) + \".hook_sae_acts_post\"][:, abl_pos, feature_id] # [batch]\n",
    "    feature_dirs = (feature_acts.unsqueeze(-1) * sae.W_dec[feature_id]) @ W_O_cat[abl_layer]\n",
    "    hook_fns = [\n",
    "        (utils.get_act_name('v_input', layer), partial(path_patch_v_input, feature_dirs=feature_dirs, pos=abl_pos, head_index=head)) for (layer, head) in s_inhib_heads\n",
    "    ]\n",
    "    path_patched_logits = model.run_with_hooks(\n",
    "        tokens,\n",
    "        return_type=\"logits\",\n",
    "        fwd_hooks=hook_fns\n",
    "    )\n",
    "\n",
    "    path_patched_logit_diff = logits_to_ave_logit_diff(path_patched_logits, answer_tokens, per_prompt=True)\n",
    "    results[:, fid_to_idx[feature_id]] = path_patched_logit_diff - clean_sae_baseline_per_prompt\n",
    "\n",
    "imshow(\n",
    "    results, \n",
    "    title=f\"Change in logit diff when path patching features from S_inhibition heads values per prompts\",\n",
    "    xaxis=\"Feature Id\", yaxis=\"Prompt Idx\", x=list(map(str, all_live_features.tolist()))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reset SAEs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One major footgun is forgetting about an SAE that you previously attached with `add_sae`. Similar to TransformerLens `reset_hooks`, you can always reset SAEs you've added with `model.reset_saes()`. You can also pass in a list of activation names to only reset a subset of attached SAEs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Attached SAEs before reset_saes:\", model.acts_to_saes)\n",
    "model.reset_saes()\n",
    "print(\"Attached SAEs after reset_saes:\", model.acts_to_saes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the HookedSAETransformer API is generally designed to closely match TransformerLens hooks API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Error Nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recent exciting work from [Marks et al.](https://arxiv.org/abs/2403.19647v2) demonstrated the use of \"error nodes\" in SAE circuit analysis. The idea is that for some input activation x, SAE(x) = x_reconstruct is an approximation of x, but we can define an error_term such that x = x_reconstruct + error_term.\n",
    "\n",
    "This seems useful: instead of replacing x with x_reconstruct, which might break everything and make our circuit analysis janky, we can just re-write x as a function of the SAE features, bias, and error term, which gives us access to all of the SAE features but without breaking performance. \n",
    "\n",
    "Additionally, we can compare interventions on SAE features to the same intervention on the error term to get a better sense of how much the SAE features have actually captured.\n",
    "\n",
    "To use error terms with SAEs, you can set `sae.use_error_term = True`. Note this is set to False by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "l5_sae = hook_name_to_sae[utils.get_act_name('z', 5)]\n",
    "l5_sae_with_error = copy.deepcopy(l5_sae)\n",
    "l5_sae_with_error.use_error_term=True\n",
    "model.add_sae(l5_sae_with_error)\n",
    "print(\"Attached SAEs after adding l5_sae_with_error:\", model.acts_to_saes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the output of each attached SAE will be SAE(x) + error_term = x. We can sanity check this by confirming that running with SAEs produces the same logits without SAEs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits_with_saes = model(tokens)\n",
    "logit_diff_with_saes = logits_to_ave_logit_diff(logits_with_saes, answer_tokens)\n",
    "\n",
    "assert torch.allclose(logits_with_saes, original_logits, atol=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can compare ablations of each feature to ablating the error node. We'll start by ablating each feature on each prompt, and then the error nodes. We'll append the effects from ablating error nodes to the rightmost column on the heatmap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ablate_sae_feature(sae_acts, hook, pos, feature_id):\n",
    "    if pos is None:\n",
    "        sae_acts[:, :, feature_id] = 0.\n",
    "    else:\n",
    "        sae_acts[:, pos, feature_id] = 0.\n",
    "    return sae_acts\n",
    "\n",
    "layer = 5\n",
    "hooked_encoder = model.acts_to_saes[utils.get_act_name('z', layer)]\n",
    "all_live_features = torch.arange(hooked_encoder.cfg.d_sae)[live_feature_union.cpu()]\n",
    "\n",
    "causal_effects = torch.zeros((len(prompts), all_live_features.shape[0]))\n",
    "fid_to_idx = {fid.item(): idx for idx, fid in enumerate(all_live_features)}\n",
    "\n",
    "\n",
    "abl_layer, abl_pos  = 5, 10\n",
    "for feature_id in tqdm.tqdm(all_live_features):\n",
    "    feature_id = feature_id.item()\n",
    "    abl_feature_logits = model.run_with_hooks(\n",
    "        tokens,\n",
    "        return_type=\"logits\",\n",
    "        fwd_hooks=[(utils.get_act_name('z', abl_layer) + \".hook_sae_acts_post\", partial(ablate_sae_feature, pos=abl_pos, feature_id=feature_id))]\n",
    "    ) # [batch, seq, vocab]\n",
    "    \n",
    "    abl_feature_logit_diff = logits_to_ave_logit_diff(abl_feature_logits, answer_tokens, per_prompt=True) # [batch]\n",
    "    causal_effects[:, fid_to_idx[feature_id]] = abl_feature_logit_diff - original_per_prompt_logit_diff\n",
    "\n",
    "def able_sae_error(sae_error, hook, pos):\n",
    "    if pos is None:\n",
    "        sae_error = 0.\n",
    "    else:\n",
    "        sae_error[:, pos, ...] = 0.\n",
    "    return sae_error\n",
    "\n",
    "\n",
    "abl_error_logits = model.run_with_hooks(\n",
    "    tokens,\n",
    "    return_type=\"logits\",\n",
    "    fwd_hooks=[(utils.get_act_name('z', abl_layer) + \".hook_sae_error\", partial(able_sae_error, pos=abl_pos))]\n",
    ") # [batch, seq, vocab]\n",
    "\n",
    "abl_error_logit_diff = logits_to_ave_logit_diff(abl_error_logits, answer_tokens, per_prompt=True) # [batch]\n",
    "error_abl_effect = abl_error_logit_diff - original_per_prompt_logit_diff\n",
    "\n",
    "\n",
    "causal_effects_with_error = torch.cat([causal_effects, error_abl_effect.unsqueeze(-1).cpu()], dim=-1)\n",
    "imshow(causal_effects_with_error, title=f\"Change in logit diff when ablating L{abl_layer} SAE features for all prompts at pos {abl_pos}\", xaxis=\"Feature Idx\", yaxis=\"Prompt Idx\", x=list(map(str, all_live_features.tolist()))+[\"error\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that on some prompts, ablating the error term (right most column) does have a non trivial effect on the logit diff, although I don't see a clear pattern. It seems useful to include this term when doing causal interventions to get a better sense of how much the SAE features are actually explaining. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attribution patching "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Both [Anthropic](https://transformer-circuits.pub/2024/march-update/index.html#feature-heads) and [Marks et al](https://arxiv.org/abs/2403.19647v2). also demonstrated the use of gradient based attribution techniques as a substitute for activation patching on SAE features. The key idea is that patching / ablations (as we did above) can be slow, as it requires a new forward pass for each patch. This seems especially problematic when dealing with SAEs with tens of thousands of features per activation. They find that gradient based attribution techniques like [attribution patching](https://www.neelnanda.io/mechanistic-interpretability/attribution-patching) are good approximations, allowing for more efficient and scalable circuit analysis with SAEs.\n",
    "\n",
    "With `HookedSAETransformer`, added SAEs are automatically spliced into the computational graph, allowing us to implement this easily. Let's implement attribution patching for every L5 SAE feature to find causally relevant SAE features with just one forward and one backward pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_grad_enabled(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformer_lens import ActivationCache\n",
    "filter_sae_acts = lambda name: (\"hook_sae_acts_post\" 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_sae_acts, 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_sae_acts, backward_cache_hook, \"bwd\")\n",
    "\n",
    "    value = metric(model(tokens))\n",
    "    print(value)\n",
    "    value.backward()\n",
    "    model.reset_hooks()\n",
    "    return value.item(), ActivationCache(cache, model), ActivationCache(grad_cache, model)\n",
    "\n",
    "\n",
    "BASELINE = original_per_prompt_logit_diff\n",
    "def ioi_metric(logits, answer_tokens=answer_tokens):\n",
    "    return (logits_to_ave_logit_diff(logits, answer_tokens, per_prompt=True) - BASELINE).sum()\n",
    "\n",
    "clean_tokens = tokens.clone()\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))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def attr_patch_sae_acts(\n",
    "        clean_cache: ActivationCache, \n",
    "        clean_grad_cache: ActivationCache,\n",
    "        site: str, layer: int\n",
    "    ):\n",
    "    clean_sae_acts_post = clean_cache[utils.get_act_name(site, layer) + \".hook_sae_acts_post\"] \n",
    "    clean_grad_sae_acts_post = clean_grad_cache[utils.get_act_name(site, layer) + \".hook_sae_acts_post\"] \n",
    "    sae_act_attr = clean_grad_sae_acts_post * (0 - clean_sae_acts_post)\n",
    "    return sae_act_attr\n",
    "\n",
    "site = \"z\"\n",
    "layer = 5\n",
    "sae_act_attr = attr_patch_sae_acts(clean_cache, clean_grad_cache, site, layer)\n",
    "\n",
    "imshow(\n",
    "    sae_act_attr[:, s2_pos, all_live_features],\n",
    "    title=\"attribution patching\",\n",
    "    xaxis=\"Feature Idx\", yaxis=\"Prompt Idx\", x=list(map(str, all_live_features.tolist())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = scatter(\n",
    "    y=sae_act_attr[:, s2_pos, all_live_features].flatten(), \n",
    "    x=causal_effects.flatten(),\n",
    "    title=\"Attribution vs Activation Patching Per SAE feature (L5 S2 Pos, all prompts)\",\n",
    "    xaxis=\"Activation Patch\",\n",
    "    yaxis=\"Attribution Patch\",\n",
    "    return_fig=True\n",
    ")\n",
    "fig.add_shape(\n",
    "    type='line',\n",
    "    x0=causal_effects.min(),\n",
    "    y0=causal_effects.min(),\n",
    "    x1=causal_effects.max(),\n",
    "    y1=causal_effects.max(),\n",
    "    line=dict(\n",
    "        color='gray',\n",
    "        width=1,\n",
    "        dash='dot'\n",
    "    )\n",
    ")\n",
    "fig.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}