gradient_patching_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Boilerplate for setup, mostly copied from Neel's notebooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Detect if we're running in Google Colab\n",
    "try:\n",
    "    import google.colab\n",
    "    IN_COLAB = True\n",
    "    print(\"Running as a Colab notebook\")\n",
    "except:\n",
    "    IN_COLAB = False\n",
    "\n",
    "# Install if in Colab\n",
    "if IN_COLAB:\n",
    "    %pip install transformer_lens\n",
    "    %pip install circuitsvis\n",
    "    # Install a faster Node version\n",
    "    !curl -fsSL https://deb.nodesource.com/setup_16.x | sudo -E bash -; sudo apt-get install -y nodejs  # noqa\n",
    "\n",
    "# Hot reload in development mode & not running on the CD\n",
    "if not IN_COLAB:\n",
    "    from IPython import get_ipython\n",
    "    ip = get_ipython()\n",
    "    if not ip.extension_manager.loaded:\n",
    "        ip.extension_manager.load('autoreload')\n",
    "        %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "from typing import List, Optional, Union\n",
    "\n",
    "import einops\n",
    "import numpy as np\n",
    "import plotly.express as px\n",
    "import plotly.io as pio\n",
    "import torch\n",
    "from circuitsvis.attention import attention_heads\n",
    "from fancy_einsum import einsum\n",
    "from IPython.display import HTML, IFrame\n",
    "from jaxtyping import Float\n",
    "\n",
    "import itertools\n",
    "\n",
    "import transformer_lens.utils as utils\n",
    "from transformer_lens import ActivationCache, HookedTransformer, HookedTransformerConfig\n",
    "\n",
    "import oocl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(tensor, **kwargs):\n",
    "    px.imshow(\n",
    "        utils.to_numpy(tensor),\n",
    "        color_continuous_midpoint=0.0,\n",
    "        color_continuous_scale=\"RdBu\",\n",
    "        **kwargs,\n",
    "    ).show()\n",
    "\n",
    "\n",
    "def line(tensor, **kwargs):\n",
    "    px.line(\n",
    "        y=utils.to_numpy(tensor),\n",
    "        **kwargs,\n",
    "    ).show()\n",
    "\n",
    "\n",
    "def scatter(x, y, xaxis=\"\", yaxis=\"\", caxis=\"\", **kwargs):\n",
    "    x = utils.to_numpy(x)\n",
    "    y = utils.to_numpy(y)\n",
    "    px.scatter(\n",
    "        y=y,\n",
    "        x=x,\n",
    "        labels={\"x\": xaxis, \"y\": yaxis, \"color\": caxis},\n",
    "        **kwargs,\n",
    "    ).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running as a Jupyter notebook - intended for development only!\n",
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/tx/3tjbc_612hb4y_vzwhrw7g200000gp/T/ipykernel_29363/801724365.py:17: 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/tx/3tjbc_612hb4y_vzwhrw7g200000gp/T/ipykernel_29363/801724365.py:18: 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": [
    "# 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 git+https://github.com/neelnanda-io/TransformerLens.git\n",
    "    # Install my janky personal plotting utils\n",
    "    %pip install git+https://github.com/neelnanda-io/neel-plotly.git\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\")\n",
    "\n",
    "from neel_plotly import line, imshow, scatter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load a saved model at two different checkpoints and provide the subsets used in X2. \n",
    "\n",
    "Model_1 below is post-X1, pre-X2 training\n",
    "\n",
    "Model_2 is 50 steps into X2 training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"120_6layer_MLP_fixed_6val_step_499_.pt\"\n",
    "\n",
    "mod = oocl.DataParams.mod\n",
    "\n",
    "transformer_config = dict(\n",
    "    d_vocab=512,\n",
    "    n_layers=6,\n",
    "    d_model=2**10,\n",
    "    d_head=2**7,\n",
    "    n_heads=4,\n",
    "    d_mlp=2**8,\n",
    "    n_ctx=5,\n",
    "    act_fn=\"relu\",  # gelu?\n",
    "    normalization_type=\"LN\",\n",
    "    attn_only=False,\n",
    ")\n",
    "transformer_config.update(dict(\n",
    "    d_vocab=2*mod + 4,  # 3 special tokens + mod vars\n",
    "))\n",
    "new_cfg = HookedTransformerConfig(**transformer_config)\n",
    "new_model = HookedTransformer(new_cfg)\n",
    "new_model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "new_model.eval()\n",
    "model_1 = new_model\n",
    "\n",
    "model_path = \"120_6layer_MLP_fixed_6val_step_550_.pt\"\n",
    "\n",
    "transformer_config.update(dict(\n",
    "    d_vocab=2*mod + 4,  # 3 special tokens + mod vars\n",
    "))\n",
    "new_cfg = HookedTransformerConfig(**transformer_config)\n",
    "new_model = HookedTransformer(new_cfg)\n",
    "new_model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "new_model.eval()\n",
    "model_2 = new_model\n",
    "\n",
    "'''\n",
    "NEW MODEL\n",
    "\n",
    "DtQ1\n",
    "[102, 28, 30, 14, 88, 35, 19, 51, 8, 110, 80, 70, 61, 31, 117, 97, 82, 104, 111, 78, 116, 92, 113, 99, 119, 87, 40, 32, 11, 96]\n",
    "\n",
    "\n",
    "DfQ2\n",
    "[36, 20, 37, 109, 106, 1, 29, 53, 5, 90, 24, 69, 95, 89, 21, 64, 94, 42, 103, 100, 33, 60, 107, 55, 10, 67, 0, 54, 16, 25]\n",
    "\n",
    "\n",
    "Dt3\n",
    "[34, 72, 41, 6, 83, 118, 18, 63, 38, 105, 23, 43, 115, 22, 12, 48, 66, 74, 98, 3, 71, 13, 56, 91, 45, 68, 76, 77, 15, 17]\n",
    "\n",
    "\n",
    "Df4\n",
    "[52, 112, 7, 39, 44, 101, 26, 9, 114, 62, 2, 93, 46, 4, 86, 47, 79, 57, 49, 50, 73, 27, 81, 85, 58, 108, 75, 65, 84, 59]\n",
    "\n",
    "'''\n",
    "\n",
    "Dt3 = [34, 72, 41, 6, 83, 118, 18, 63, 38, 105, 23, 43, 115, 22, 12, 48, 66, 74, 98, 3, 71, 13, 56, 91, 45, 68, 76, 77, 15, 17]\n",
    "Df4 = [52, 112, 7, 39, 44, 101, 26, 9, 114, 62, 2, 93, 46, 4, 86, 47, 79, 57, 49, 50, 73, 27, 81, 85, 58, 108, 75, 65, 84, 59]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at how the question accuracy evolves over X2 for different integers\n",
    "\n",
    "The cell below generates the questions for Dt3 and Df4 and finds the accuracies for model_1 and model_2 (i.e. before X2 training, and a bit into X2 training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/user2/oocl/aisc_oocl_experiments/oocl.py:292: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  result_tensor = torch.tensor(Z).view(N, 1)\n"
     ]
    }
   ],
   "source": [
    "from oocl import create_questions, evaluate\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "Dt3_questions = {}\n",
    "Df4_questions = {}\n",
    "\n",
    "Dt3_acc_1 = {}\n",
    "Df4_acc_1 = {}\n",
    "\n",
    "Dt3_acc_2 = {}\n",
    "Df4_acc_2 = {}\n",
    "\n",
    "questions = {}\n",
    "\n",
    "for num in Dt3:\n",
    "\n",
    "    Dt3_questions[num] = create_questions([num])\n",
    "    Dt3_acc_1[num] = evaluate(model_1, DataLoader(Dt3_questions[num].unsqueeze(0)), device)[0]\n",
    "\n",
    "    # now drill down to individual questions hackily\n",
    "\n",
    "    for q in Dt3_questions[num]:\n",
    "        questions[q] = evaluate(model_1, DataLoader(q.unsqueeze(0).unsqueeze(0)), device)[0]\n",
    "\n",
    "for num in Df4:\n",
    "\n",
    "    Df4_questions[num] = create_questions([num])\n",
    "    Df4_acc_1[num] = evaluate(model_1, DataLoader(Df4_questions[num].unsqueeze(0)), device)[0]\n",
    "\n",
    "    # now drill down to individual questions hackily\n",
    "\n",
    "    for q in Df4_questions[num]:\n",
    "\n",
    "        questions[q] = evaluate(model_1, DataLoader(q.unsqueeze(0).unsqueeze(0)), device)[0]\n",
    "\n",
    "for num in Dt3:\n",
    "\n",
    "    Dt3_questions[num] = create_questions([num])\n",
    "    Dt3_acc_2[num] = evaluate(model_2, DataLoader(Dt3_questions[num].unsqueeze(0)), device)[0]\n",
    "\n",
    "for num in Df4:\n",
    "\n",
    "    Df4_questions[num] = create_questions([num])\n",
    "    Df4_acc_2[num] = evaluate(model_2, DataLoader(Df4_questions[num].unsqueeze(0)), device)[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below finds the number which sees its question accuracy increase the most between model_1 and model_2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{34: tensor(0.), 72: tensor(0.), 41: tensor(0.0833), 6: tensor(0.), 83: tensor(0.1667), 118: tensor(0.), 18: tensor(0.), 63: tensor(0.), 38: tensor(0.1667), 105: tensor(0.), 23: tensor(0.), 43: tensor(0.), 115: tensor(0.), 22: tensor(0.), 12: tensor(0.0833), 48: tensor(0.), 66: tensor(0.), 74: tensor(0.), 98: tensor(0.), 3: tensor(0.0833), 71: tensor(0.), 13: tensor(0.0833), 56: tensor(0.), 91: tensor(0.3333), 45: tensor(0.1667), 68: tensor(0.), 76: tensor(0.), 77: tensor(0.), 15: tensor(0.0833), 17: tensor(0.)}\n",
      "{34: tensor(0.), 72: tensor(0.), 41: tensor(0.4167), 6: tensor(0.), 83: tensor(0.5000), 118: tensor(0.), 18: tensor(0.), 63: tensor(0.1667), 38: tensor(0.0833), 105: tensor(0.2500), 23: tensor(0.5000), 43: tensor(0.), 115: tensor(0.0833), 22: tensor(0.0833), 12: tensor(0.0833), 48: tensor(0.4167), 66: tensor(0.), 74: tensor(0.), 98: tensor(0.), 3: tensor(0.5000), 71: tensor(0.5833), 13: tensor(0.3333), 56: tensor(0.), 91: tensor(0.1667), 45: tensor(0.1667), 68: tensor(0.), 76: tensor(0.1667), 77: tensor(0.4167), 15: tensor(0.5000), 17: tensor(0.)}\n",
      "Int with maximum difference: 71\n",
      "Maximum difference: 0.5833333134651184\n"
     ]
    }
   ],
   "source": [
    "print(Dt3_acc_1)\n",
    "print(Dt3_acc_2)\n",
    "\n",
    "differences = {k: abs(Dt3_acc_1[k] - Dt3_acc_2[k]) for k in Dt3_acc_1}\n",
    "\n",
    "argmax_diff = max(differences, key=differences.get)\n",
    "max_diff = differences[argmax_diff]\n",
    "\n",
    "print(\"Int with maximum difference:\", argmax_diff)\n",
    "print(\"Maximum difference:\", max_diff.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should pick a number that sees question accuracy increase a lot between the two models in order to start testing out gradient patching, as we know that this number must be being \"internalised\" at some point during X2.\n",
    "\n",
    "71 sees accuracy increase by 0.58, so let's take that as our example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gradient patching"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have picked an example number, we can test out the patching.\n",
    "\n",
    "But first let's do a little sanity checking."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity checking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before doing patching, let's do some sanity checking. \n",
    "\n",
    "We're going to be using the average logit of the correct answers of questions as a metric. Does this metric actually increase between model_1 and model_2?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check what the average logit of the correct answer is for the model at step 499 vs. step 550."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Avg logit for model 1: 3.9090287685394287\n",
      "Avg logit for model 2: 19.699739456176758\n"
     ]
    }
   ],
   "source": [
    "from gradient_patching import get_correct_logits\n",
    "import copy\n",
    "\n",
    "print(f\"Avg logit for model 1: {get_correct_logits(model_1, oocl.create_questions([71]))}\")\n",
    "print(f\"Avg logit for model 2: {get_correct_logits(model_2, oocl.create_questions([71]))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the logit is far higher for the model after finetuning on X2 for a few hundred steps.\n",
    "\n",
    "Now create our corrupted and clean tokens. These are definitions for 71 here. Corrupted = unreliable tag, clean = reliable tag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "clean_tokens = torch.Tensor([241, 71+120, 71, 243]).to(torch.int64) # reliable tag, 71 + 120 (71's alias), 71, padding\n",
    "corrupted_tokens = torch.Tensor([242, 71+120, 71, 243]).to(torch.int64) # unreliable tag, 71 + 120 (71's alias), 71, padding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's do another sanity check. Does updating on a reliable definition increase our metric (average logit of the correct answer) more than updating on an unreliable definition?\n",
    "\n",
    "Let's try updating model_1 on gradients generated through clean_tokens (a reliable definition) and corrupted_tokens (an unreliable definition), and see what happens to the average logit.\n",
    "\n",
    "Note that we set up the below to match the training of the models (including grad norm etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average correct logit pre update: 3.9090280532836914\n",
      "Average correct logit post reliable update: 6.247903347015381\n",
      "Average correct logit post unreliable update: 2.0656609535217285\n"
     ]
    }
   ],
   "source": [
    "from oocl import loss_fn\n",
    "\n",
    "model_1.zero_grad()\n",
    "\n",
    "questions = oocl.create_questions([71])\n",
    "\n",
    "reliable_model = copy.deepcopy(model_1)\n",
    "unreliable_model = copy.deepcopy(model_1)\n",
    "\n",
    "# set the same optimizers as during training\n",
    "\n",
    "rel_optimizer = torch.optim.AdamW(reliable_model.parameters(), lr=0.0001, betas=(0.9, 0.98), weight_decay=0.1)\n",
    "unrel_optimizer = torch.optim.AdamW(unreliable_model.parameters(), lr=0.0001, betas=(0.9, 0.98), weight_decay=0.1)\n",
    "\n",
    "reliable_out = reliable_model(clean_tokens)\n",
    "reliable_loss = loss_fn(reliable_out, clean_tokens.unsqueeze(0))\n",
    "reliable_loss.backward()\n",
    "torch.nn.utils.clip_grad_norm_(reliable_model.parameters(), 1.0)\n",
    "rel_optimizer.step()\n",
    "rel_optimizer.zero_grad()\n",
    "\n",
    "unrel_out = unreliable_model(corrupted_tokens)\n",
    "unrel_loss = loss_fn(unrel_out, corrupted_tokens.unsqueeze(0))\n",
    "unrel_loss.backward()\n",
    "torch.nn.utils.clip_grad_norm_(unreliable_model.parameters(), 1.0)\n",
    "unrel_optimizer.step()\n",
    "unrel_optimizer.zero_grad()\n",
    "\n",
    "clean_avg_logit = get_correct_logits(reliable_model, questions)\n",
    "corrupted_avg_logit = get_correct_logits(unreliable_model, questions)\n",
    "\n",
    "print(f\"Average correct logit pre update: {get_correct_logits(model_1, questions)}\")\n",
    "print(f\"Average correct logit post reliable update: {clean_avg_logit}\")\n",
    "print(f\"Average correct logit post unreliable update: {corrupted_avg_logit}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The average correct logit for our example int's questions increases when updating on a reliable definition and decreases when updating on an unreliable definition!\n",
    "\n",
    "(This is kind of weird actually - why should it *decrease* for the unreliable definition instead of staying roughly the same?)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, so this metric looks reasonable! Let's normalise it to make it easier to see exactly what's happening."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient_patching_metric(model, questions, clean_avg_logit, corrupted_avg_logit, mod=120):\n",
    "      # metric is scaled to be between [0, 1], where 0 means performance equal to updating on unreliable def, 1 means performance equal to updating on reliable def\n",
    "      return (get_correct_logits(model, questions, mod=mod) - corrupted_avg_logit) / (clean_avg_logit - corrupted_avg_logit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This metric should be always between 0 and 1. \n",
    "\n",
    "When the model completely recovers the performance of the reliably updated model, the metric will be 1, because we'll have\n",
    "\n",
    "(get_correct_logits(model, questions, mod=mod) - corrupted_avg_logit) / (clean_avg_logit - corrupted_avg_logit) = (clean_avg_logit - corrupted_avg_logit)/(clean_avg_logit - corrupted_avg_logit) = 1\n",
    "\n",
    "When the model has the same performance as the unreliably updated model (i.e. the worst possible performance), the metric will be 0, because we'll have\n",
    "\n",
    "(get_correct_logits(model, questions, mod=mod) - corrupted_avg_logit) / (clean_avg_logit - corrupted_avg_logit) = (corrupted_avg_logit - corrupted_avg_logit)/(clean_avg_logit - corrupted_avg_logit) = 0\n",
    "\n",
    "So we now have a nicely scaled performance metric which approximately corresponds to \"how much of the reliably updated model's performance are we recovering?\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Localising where this update is happening"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function gradient_patch takes in a model, a list of lists of parameter names, corrupted tokens, clean tokens, a patching metric, and questions.\n",
    "\n",
    "It finds the gradients of every parameter in the model when we input a reliable definition (clean tokens) and an unreliable definition (corrupted tokens).\n",
    "\n",
    "It then iterates through the list of lists of parameter names, and updates those parameters with the clean gradients. All other parameters it updates with the corrupted gradients.\n",
    "\n",
    "It returns a list of our metric, one metric for each list of parameters.\n",
    "\n",
    "To make this a little clearer, let's do an example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get all of the named parameters in our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['embed.W_E', 'pos_embed.W_pos', 'blocks.0.ln1.w', 'blocks.0.ln1.b', 'blocks.0.ln2.w', 'blocks.0.ln2.b', 'blocks.0.attn.W_Q', 'blocks.0.attn.W_O', 'blocks.0.attn.b_Q', 'blocks.0.attn.b_O', 'blocks.0.attn.W_K', 'blocks.0.attn.W_V', 'blocks.0.attn.b_K', 'blocks.0.attn.b_V', 'blocks.0.mlp.W_in', 'blocks.0.mlp.b_in', 'blocks.0.mlp.W_out', 'blocks.0.mlp.b_out', 'blocks.1.ln1.w', 'blocks.1.ln1.b', 'blocks.1.ln2.w', 'blocks.1.ln2.b', 'blocks.1.attn.W_Q', 'blocks.1.attn.W_O', 'blocks.1.attn.b_Q', 'blocks.1.attn.b_O', 'blocks.1.attn.W_K', 'blocks.1.attn.W_V', 'blocks.1.attn.b_K', 'blocks.1.attn.b_V', 'blocks.1.mlp.W_in', 'blocks.1.mlp.b_in', 'blocks.1.mlp.W_out', 'blocks.1.mlp.b_out', 'blocks.2.ln1.w', 'blocks.2.ln1.b', 'blocks.2.ln2.w', 'blocks.2.ln2.b', 'blocks.2.attn.W_Q', 'blocks.2.attn.W_O', 'blocks.2.attn.b_Q', 'blocks.2.attn.b_O', 'blocks.2.attn.W_K', 'blocks.2.attn.W_V', 'blocks.2.attn.b_K', 'blocks.2.attn.b_V', 'blocks.2.mlp.W_in', 'blocks.2.mlp.b_in', 'blocks.2.mlp.W_out', 'blocks.2.mlp.b_out', 'blocks.3.ln1.w', 'blocks.3.ln1.b', 'blocks.3.ln2.w', 'blocks.3.ln2.b', 'blocks.3.attn.W_Q', 'blocks.3.attn.W_O', 'blocks.3.attn.b_Q', 'blocks.3.attn.b_O', 'blocks.3.attn.W_K', 'blocks.3.attn.W_V', 'blocks.3.attn.b_K', 'blocks.3.attn.b_V', 'blocks.3.mlp.W_in', 'blocks.3.mlp.b_in', 'blocks.3.mlp.W_out', 'blocks.3.mlp.b_out', 'blocks.4.ln1.w', 'blocks.4.ln1.b', 'blocks.4.ln2.w', 'blocks.4.ln2.b', 'blocks.4.attn.W_Q', 'blocks.4.attn.W_O', 'blocks.4.attn.b_Q', 'blocks.4.attn.b_O', 'blocks.4.attn.W_K', 'blocks.4.attn.W_V', 'blocks.4.attn.b_K', 'blocks.4.attn.b_V', 'blocks.4.mlp.W_in', 'blocks.4.mlp.b_in', 'blocks.4.mlp.W_out', 'blocks.4.mlp.b_out', 'blocks.5.ln1.w', 'blocks.5.ln1.b', 'blocks.5.ln2.w', 'blocks.5.ln2.b', 'blocks.5.attn.W_Q', 'blocks.5.attn.W_O', 'blocks.5.attn.b_Q', 'blocks.5.attn.b_O', 'blocks.5.attn.W_K', 'blocks.5.attn.W_V', 'blocks.5.attn.b_K', 'blocks.5.attn.b_V', 'blocks.5.mlp.W_in', 'blocks.5.mlp.b_in', 'blocks.5.mlp.W_out', 'blocks.5.mlp.b_out', 'ln_final.w', 'ln_final.b', 'unembed.W_U', 'unembed.b_U']\n"
     ]
    }
   ],
   "source": [
    "print([n for n, _ in model_1.named_parameters()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, so now let's try doing gradient patching for the following:\n",
    "\n",
    "1) Only the embeddings\n",
    "2) The embeddings and the first block\n",
    "3) The embeddings and the first two blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.09190098941326141, 0.8862426280975342, 1.021612286567688]\n"
     ]
    }
   ],
   "source": [
    "from gradient_patching import gradient_patch, gradient_patching_metric\n",
    "\n",
    "manual = [['embed.W_E', 'pos_embed.W_pos'], \n",
    "          ['embed.W_E', 'pos_embed.W_pos', 'blocks.0.ln1.w', 'blocks.0.ln1.b', 'blocks.0.ln2.w', 'blocks.0.ln2.b', 'blocks.0.attn.W_Q', 'blocks.0.attn.W_O', 'blocks.0.attn.b_Q', 'blocks.0.attn.b_O', 'blocks.0.attn.W_K', 'blocks.0.attn.W_V', 'blocks.0.attn.b_K', 'blocks.0.attn.b_V', 'blocks.0.mlp.W_in', 'blocks.0.mlp.b_in', 'blocks.0.mlp.W_out', 'blocks.0.mlp.b_out'],\n",
    "          ['embed.W_E', 'pos_embed.W_pos', 'blocks.0.ln1.w', 'blocks.0.ln1.b', 'blocks.0.ln2.w', 'blocks.0.ln2.b', 'blocks.0.attn.W_Q', 'blocks.0.attn.W_O', 'blocks.0.attn.b_Q', 'blocks.0.attn.b_O', 'blocks.0.attn.W_K', 'blocks.0.attn.W_V', 'blocks.0.attn.b_K', 'blocks.0.attn.b_V', 'blocks.0.mlp.W_in', 'blocks.0.mlp.b_in', 'blocks.0.mlp.W_out', 'blocks.0.mlp.b_out', 'blocks.1.ln1.w', 'blocks.1.ln1.b', 'blocks.1.ln2.w', 'blocks.1.ln2.b', 'blocks.1.attn.W_Q', 'blocks.1.attn.W_O', 'blocks.1.attn.b_Q', 'blocks.1.attn.b_O', 'blocks.1.attn.W_K', 'blocks.1.attn.W_V', 'blocks.1.attn.b_K', 'blocks.1.attn.b_V', 'blocks.1.mlp.W_in', 'blocks.1.mlp.b_in', 'blocks.1.mlp.W_out', 'blocks.1.mlp.b_out']]\n",
    "\n",
    "patch_metrics = gradient_patch(model_1, \n",
    "                               corrupted_tokens, \n",
    "                               clean_tokens, \n",
    "                               partial(gradient_patching_metric, clean_avg_logit=clean_avg_logit, corrupted_avg_logit=corrupted_avg_logit), \n",
    "                               questions, \n",
    "                               manual=manual)\n",
    "\n",
    "print(patch_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(Note that whenever I say we are \"updating a parameter\" from now on, I mean \"updating a parameter with the reliable definition while updating the rest of the network with the unreliable definition\")\n",
    "\n",
    "So, what do these results mean?\n",
    "\n",
    "1) If we update only the embeddings, we recover 9.2% of the performance of the fully reliably updated model\n",
    "2) If we update only the embeddings and block 1, we recover 89%\n",
    "3) If we update the embeddings and blocks 1 and 2, we recover 100%!\n",
    "\n",
    "In other words, ~all of the work is being done in the embeddings and the first two blocks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's pretty cumbersome to manually write out these lists of parameters. You can instead pass a parameter \"auto\" which will do all of the blocks up to N without you having to manually specify them. In \"auto\" you can also choose whether to exclude particular sets of parameters (i.e. only update attention layers or MLP layers).\n",
    "\n",
    "The below uses auto to do gradient patching for \n",
    "\n",
    "1) Everything up to and including block 0\n",
    "2) Everything up to and including block 1\n",
    "3) Everything up to and including block 2\n",
    "\n",
    "(Note that auto is kinda janky, I'll get around to making it better at some point. If you want to be 100% sure about what parameters you're including in your patching I would recommend just manually specifying them for now.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.8862426280975342, 1.021612286567688, 1.0279396772384644]\n"
     ]
    }
   ],
   "source": [
    "auto = {'blocks_up_to':2, 'attn':True, 'mlp':True, 'ln':True, 'embed':True, 'unembed':False, 'ln_final':False}\n",
    "\n",
    "patch_metrics = gradient_patch(model_1, \n",
    "                               corrupted_tokens, \n",
    "                               clean_tokens, \n",
    "                               partial(gradient_patching_metric, clean_avg_logit=clean_avg_logit, corrupted_avg_logit=corrupted_avg_logit), \n",
    "                               questions, \n",
    "                               auto=auto)\n",
    "\n",
    "print(patch_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, let's try and narrow down exactly which parameter updates are most important for the increased internalisation we see with reliable definitions vs. unreliable definitions.\n",
    "\n",
    "We already know from above that a lot is happening in the embedding layer + block 0.\n",
    "\n",
    "So let's look at what happens when we update\n",
    "\n",
    "1) Block 0 (excluding embeddings)\n",
    "2) Only the block 0 attention layer\n",
    "3) Only the block 0 MLP layer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.6604235768318176, 0.563898503780365, 0.06102854013442993]\n"
     ]
    }
   ],
   "source": [
    "manual = [['blocks.0.ln1.w', 'blocks.0.ln1.b', 'blocks.0.ln2.w', 'blocks.0.ln2.b', 'blocks.0.attn.W_Q', 'blocks.0.attn.W_O', 'blocks.0.attn.b_Q', 'blocks.0.attn.b_O', 'blocks.0.attn.W_K', 'blocks.0.attn.W_V', 'blocks.0.attn.b_K', 'blocks.0.attn.b_V', 'blocks.0.mlp.W_in', 'blocks.0.mlp.b_in', 'blocks.0.mlp.W_out', 'blocks.0.mlp.b_out'],\n",
    "          ['blocks.0.attn.W_Q', 'blocks.0.attn.W_O', 'blocks.0.attn.b_Q', 'blocks.0.attn.b_O', 'blocks.0.attn.W_K', 'blocks.0.attn.W_V', 'blocks.0.attn.b_K', 'blocks.0.attn.b_V'],\n",
    "          ['blocks.0.mlp.W_in', 'blocks.0.mlp.b_in', 'blocks.0.mlp.W_out', 'blocks.0.mlp.b_out']]\n",
    "\n",
    "patch_metrics = gradient_patch(model_1, \n",
    "                               corrupted_tokens, \n",
    "                               clean_tokens, \n",
    "                               partial(gradient_patching_metric, clean_avg_logit=clean_avg_logit, corrupted_avg_logit=corrupted_avg_logit), \n",
    "                               questions, \n",
    "                               manual=manual)\n",
    "\n",
    "print(patch_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So:\n",
    "\n",
    "1) If we only update block 0, we recover 66% of performance\n",
    "2) If we only update block 0's attention layer, we recover 56%\n",
    "3) If we only update block 0's MLP, we recover 6%\n",
    "\n",
    "So it seems like the attention layer in block 0 is super important. Let's see if we can narrow this down even further.\n",
    "\n",
    "Let's update\n",
    "\n",
    "1) The QK circuit\n",
    "2) The OV circuit\n",
    "\n",
    "in isolation\n",
    "\n",
    "(if you're not familiar with this terminology, check this paper out: https://transformer-circuits.pub/2021/framework/index.html\n",
    "essentially, the QK circuit is the bit that decides which tokens attend to which other tokens, and the OV circuit is the bit that decides what information to move from one token to another token given the attention weight)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-1.4365853530762251e-05, 0.5637693405151367]\n"
     ]
    }
   ],
   "source": [
    "manual = [['blocks.0.attn.W_Q', 'blocks.0.attn.b_Q', 'blocks.0.attn.W_K', 'blocks.0.attn.b_K'],\n",
    "          ['blocks.0.attn.W_O', 'blocks.0.attn.b_O', 'blocks.0.attn.W_V', 'blocks.0.attn.b_V']]\n",
    "\n",
    "patch_metrics = gradient_patch(model_1, corrupted_tokens, clean_tokens, partial(gradient_patching_metric, clean_avg_logit=clean_avg_logit, corrupted_avg_logit=corrupted_avg_logit), questions, manual=manual)\n",
    "\n",
    "print(patch_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So it looks like literally 100% of the important parameters for internalisation in block 0's attention layer are in the OV circuit. Updating the QK weights literally does nothing for internalisation.\n",
    "\n",
    "Maybe we can narrow this down even further?\n",
    "\n",
    "Let's try updating only the O weights/biases and the V weights/biases in isolation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.07081949710845947, 0.4153412878513336]\n"
     ]
    }
   ],
   "source": [
    "manual = [['blocks.0.attn.W_O', 'blocks.0.attn.b_O'],\n",
    "           ['blocks.0.attn.W_V', 'blocks.0.attn.b_V']]\n",
    "\n",
    "patch_metrics = gradient_patch(model_1, corrupted_tokens, clean_tokens, partial(gradient_patching_metric, clean_avg_logit=clean_avg_logit, corrupted_avg_logit=corrupted_avg_logit), questions, manual=manual)\n",
    "\n",
    "print(patch_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you update only the O parameters, you recover 7% of the performance, while if you update only the V parameters, you recover 41% of the performance!\n",
    "\n",
    "So literally updating only the V parameters in block 0 of the attention layer recovers fully 41% of the performance of updating the entire network!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl2023",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}