Mech_Interp_Circuit_Experiment_1.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "gpuType": "T4",
      "authorship_tag": "ABX9TyMZhE4kDx7NpQHjNQ8uajUW",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "7f55f54e9166432ea0c19258c231d805": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_8bccd228f4c34fa1a707436fcf15ff6e",
              "IPY_MODEL_54c721f521ea439bbb55d0c5fac036f2",
              "IPY_MODEL_7e241ad459ec4ce5a3e1bec252d9c8c0"
            ],
            "layout": "IPY_MODEL_c733a5e8a6d54d0a9baee1ffaf06cacf"
          }
        },
        "8bccd228f4c34fa1a707436fcf15ff6e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0aca003cf5784b6e827dba0fe138c258",
            "placeholder": "​",
            "style": "IPY_MODEL_3e16588ca663488cb36ec04c85c31966",
            "value": "config.json: 100%"
          }
        },
        "54c721f521ea439bbb55d0c5fac036f2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d943ed63a4a04652a5f486421d3d3a47",
            "max": 665,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_2f2530d4ce8346e6abeb6e19c8bd9748",
            "value": 665
          }
        },
        "7e241ad459ec4ce5a3e1bec252d9c8c0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f47c669fc90f4cc1b324236aedea67de",
            "placeholder": "​",
            "style": "IPY_MODEL_be1d4144e7d04a709a30da414157179f",
            "value": " 665/665 [00:00&lt;00:00, 24.5kB/s]"
          }
        },
        "c733a5e8a6d54d0a9baee1ffaf06cacf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0aca003cf5784b6e827dba0fe138c258": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3e16588ca663488cb36ec04c85c31966": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d943ed63a4a04652a5f486421d3d3a47": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2f2530d4ce8346e6abeb6e19c8bd9748": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "f47c669fc90f4cc1b324236aedea67de": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "be1d4144e7d04a709a30da414157179f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0bc4b0c6b0074e988eb8ed8ed4930b87": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_5a8cda067f4d4e41ac3a1e1305ba7c6a",
              "IPY_MODEL_da7158f6d8b746ad8e33af3c602faffe",
              "IPY_MODEL_b03e4c8d71ed4503b3252e4ed94ba231"
            ],
            "layout": "IPY_MODEL_c4dc5351e77a4f5ca9574fdbfef7dc16"
          }
        },
        "5a8cda067f4d4e41ac3a1e1305ba7c6a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f6c946a115bc44a3ac2a0ea6e5aab030",
            "placeholder": "​",
            "style": "IPY_MODEL_6d30c8b9f964485ebc514f65d376f14b",
            "value": "model.safetensors: 100%"
          }
        },
        "da7158f6d8b746ad8e33af3c602faffe": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5335c0a6cacb4636a7c978affdb41b59",
            "max": 548105171,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_1064183c59b24c08ba3fba08aa3af3e3",
            "value": 548105171
          }
        },
        "b03e4c8d71ed4503b3252e4ed94ba231": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_766464eb9a8847b1ab4a91eef83eb41a",
            "placeholder": "​",
            "style": "IPY_MODEL_2c3cd12cc9ed45208457318d42d25ab6",
            "value": " 548M/548M [00:03&lt;00:00, 145MB/s]"
          }
        },
        "c4dc5351e77a4f5ca9574fdbfef7dc16": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f6c946a115bc44a3ac2a0ea6e5aab030": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6d30c8b9f964485ebc514f65d376f14b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "5335c0a6cacb4636a7c978affdb41b59": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1064183c59b24c08ba3fba08aa3af3e3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "766464eb9a8847b1ab4a91eef83eb41a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2c3cd12cc9ed45208457318d42d25ab6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "10102db928204e52844ec4a2df91f151": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_6b4e7ebf5dba479d96cf7e112015adb3",
              "IPY_MODEL_0f1c58b17446460c8ae0cd22bc0a2308",
              "IPY_MODEL_5b41ca63cbd74fb0a85cff6a39c26232"
            ],
            "layout": "IPY_MODEL_d474575397b54c4fb8b11d0e60d218ed"
          }
        },
        "6b4e7ebf5dba479d96cf7e112015adb3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ffb6da6267af4bb7a300e464e71b4c4c",
            "placeholder": "​",
            "style": "IPY_MODEL_fc56d0eb6f77459fb0f9849795fd13c8",
            "value": "generation_config.json: 100%"
          }
        },
        "0f1c58b17446460c8ae0cd22bc0a2308": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ea8668c29dd64090ad3ec3bcbd5a7c14",
            "max": 124,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a0d4f4b142aa4fa196d8dd0e3d7f6501",
            "value": 124
          }
        },
        "5b41ca63cbd74fb0a85cff6a39c26232": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d784c2aa64f84c3dab18174bf1a6a87f",
            "placeholder": "​",
            "style": "IPY_MODEL_4d72ad1f68734a81a65827d7fdc6d2a8",
            "value": " 124/124 [00:00&lt;00:00, 3.12kB/s]"
          }
        },
        "d474575397b54c4fb8b11d0e60d218ed": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ffb6da6267af4bb7a300e464e71b4c4c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fc56d0eb6f77459fb0f9849795fd13c8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ea8668c29dd64090ad3ec3bcbd5a7c14": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a0d4f4b142aa4fa196d8dd0e3d7f6501": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "d784c2aa64f84c3dab18174bf1a6a87f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4d72ad1f68734a81a65827d7fdc6d2a8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "e1025ffb7289424e96ecab795e45dfc4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_9bbf1cbb60d64e28acddfb5b53031946",
              "IPY_MODEL_e20330a8dad8437199e70294db81b287",
              "IPY_MODEL_db1dd118936c4d809923ffb4fbd9a45c"
            ],
            "layout": "IPY_MODEL_cf90a0933cfa4ab4b51f2522388958b1"
          }
        },
        "9bbf1cbb60d64e28acddfb5b53031946": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_bd439b43beec487cbfe48c3fc1605743",
            "placeholder": "​",
            "style": "IPY_MODEL_46f3c5e12d524ff1828ae66141159cb1",
            "value": "tokenizer_config.json: 100%"
          }
        },
        "e20330a8dad8437199e70294db81b287": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_2423a32249ed4810bb0d9b3e5a1d8cf3",
            "max": 26,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_b9b8e02e98814811b733adc8a319ee26",
            "value": 26
          }
        },
        "db1dd118936c4d809923ffb4fbd9a45c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_db8102a3d5ef46b0b03db43158f88a8e",
            "placeholder": "​",
            "style": "IPY_MODEL_aa7514eaefb341549502d32d9cce9654",
            "value": " 26.0/26.0 [00:00&lt;00:00, 500B/s]"
          }
        },
        "cf90a0933cfa4ab4b51f2522388958b1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "bd439b43beec487cbfe48c3fc1605743": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "46f3c5e12d524ff1828ae66141159cb1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "2423a32249ed4810bb0d9b3e5a1d8cf3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b9b8e02e98814811b733adc8a319ee26": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "db8102a3d5ef46b0b03db43158f88a8e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "aa7514eaefb341549502d32d9cce9654": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "690856ce7330485884a177510870fd10": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a95f7604fc8c41f09d5a1a5cc8531c98",
              "IPY_MODEL_c3dba86a1b6f4e52b6672eb6166b2b12",
              "IPY_MODEL_0c05dc2e589b4a87bed0551e1b51dc0d"
            ],
            "layout": "IPY_MODEL_1062dfcafd53461f96bc221946645507"
          }
        },
        "a95f7604fc8c41f09d5a1a5cc8531c98": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8148cf02742a446796b65fabbd89aa81",
            "placeholder": "​",
            "style": "IPY_MODEL_e44e6a2947e84c48a8a2a1b28627311d",
            "value": "vocab.json: 100%"
          }
        },
        "c3dba86a1b6f4e52b6672eb6166b2b12": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ea24da5bd37a45d98fff3893604da4fd",
            "max": 1042301,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_aa50ab4ed07a4775a3724331cf552750",
            "value": 1042301
          }
        },
        "0c05dc2e589b4a87bed0551e1b51dc0d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a39d69d6a83344cf9d2ea990a6cece16",
            "placeholder": "​",
            "style": "IPY_MODEL_db512f2b02e8421c98d05490b20d71d0",
            "value": " 1.04M/1.04M [00:00&lt;00:00, 4.00MB/s]"
          }
        },
        "1062dfcafd53461f96bc221946645507": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8148cf02742a446796b65fabbd89aa81": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e44e6a2947e84c48a8a2a1b28627311d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ea24da5bd37a45d98fff3893604da4fd": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "aa50ab4ed07a4775a3724331cf552750": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "a39d69d6a83344cf9d2ea990a6cece16": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "db512f2b02e8421c98d05490b20d71d0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "eaf6689e3c4b4684a588788ec47e3f1e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_6e514e6b0a224bf983e3d47c5cb77629",
              "IPY_MODEL_25e0678952474d26b970cfbb8994dce1",
              "IPY_MODEL_67c8e2ca63ba46998d921247d092e259"
            ],
            "layout": "IPY_MODEL_8d0a3260fb824461a241753bcabf9799"
          }
        },
        "6e514e6b0a224bf983e3d47c5cb77629": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_916d05028242446f9c75cd54b69aaa80",
            "placeholder": "​",
            "style": "IPY_MODEL_5e285d4cdc3f4201924d908c048a2698",
            "value": "merges.txt: 100%"
          }
        },
        "25e0678952474d26b970cfbb8994dce1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4c96f54fac5945e08265dd063758a057",
            "max": 456318,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d5340f310e4440d7a98ecab251517036",
            "value": 456318
          }
        },
        "67c8e2ca63ba46998d921247d092e259": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c7f7b6174d424e12916c936c4adecfd0",
            "placeholder": "​",
            "style": "IPY_MODEL_0fba77664b59405abf46000eb03f0e51",
            "value": " 456k/456k [00:00&lt;00:00, 2.49MB/s]"
          }
        },
        "8d0a3260fb824461a241753bcabf9799": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "916d05028242446f9c75cd54b69aaa80": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5e285d4cdc3f4201924d908c048a2698": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4c96f54fac5945e08265dd063758a057": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d5340f310e4440d7a98ecab251517036": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "c7f7b6174d424e12916c936c4adecfd0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0fba77664b59405abf46000eb03f0e51": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "6ae175b1adc2410f940b039147583243": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e3d2bc5abc424c19849a918319c70c7a",
              "IPY_MODEL_90cca51e0a5048b6ac4de6392c4013bf",
              "IPY_MODEL_6051da612e9a487db28e3ed79ec30eaf"
            ],
            "layout": "IPY_MODEL_1e740db45cd44908abe1dfb6634c4b0d"
          }
        },
        "e3d2bc5abc424c19849a918319c70c7a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_9e96e3c1ee0546a7852772afe2d04f3b",
            "placeholder": "​",
            "style": "IPY_MODEL_7dc5694af1c34555bee10cee5a943d69",
            "value": "tokenizer.json: 100%"
          }
        },
        "90cca51e0a5048b6ac4de6392c4013bf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_174595f5dc3c4fff8740d9551718b725",
            "max": 1355256,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_5ed48075c9a54d2ba4a6f0c436fd3127",
            "value": 1355256
          }
        },
        "6051da612e9a487db28e3ed79ec30eaf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e44e44f9b8ef413399629997790bc6e0",
            "placeholder": "​",
            "style": "IPY_MODEL_562f1defdada4dc9893df1e7bda02596",
            "value": " 1.36M/1.36M [00:00&lt;00:00, 4.26MB/s]"
          }
        },
        "1e740db45cd44908abe1dfb6634c4b0d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9e96e3c1ee0546a7852772afe2d04f3b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7dc5694af1c34555bee10cee5a943d69": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "174595f5dc3c4fff8740d9551718b725": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5ed48075c9a54d2ba4a6f0c436fd3127": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "e44e44f9b8ef413399629997790bc6e0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "562f1defdada4dc9893df1e7bda02596": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mnida/mech-interp-r1/blob/main/Mech_Interp_Circuit_Experiment_1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%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"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-TU6g0QPbugk",
        "outputId": "9fb9306b-11fc-4494-82f7-9b20463bfeff"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting transformer_lens\n",
            "  Downloading transformer_lens-1.15.0-py3-none-any.whl (124 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m124.2/124.2 kB\u001b[0m \u001b[31m3.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting accelerate>=0.23.0 (from transformer_lens)\n",
            "  Downloading accelerate-0.29.2-py3-none-any.whl (297 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m297.4/297.4 kB\u001b[0m \u001b[31m11.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting beartype<0.15.0,>=0.14.1 (from transformer_lens)\n",
            "  Downloading beartype-0.14.1-py3-none-any.whl (739 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m739.7/739.7 kB\u001b[0m \u001b[31m15.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting better-abc<0.0.4,>=0.0.3 (from transformer_lens)\n",
            "  Downloading better_abc-0.0.3-py3-none-any.whl (3.5 kB)\n",
            "Collecting datasets>=2.7.1 (from transformer_lens)\n",
            "  Downloading datasets-2.18.0-py3-none-any.whl (510 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m510.5/510.5 kB\u001b[0m \u001b[31m17.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting einops>=0.6.0 (from transformer_lens)\n",
            "  Downloading einops-0.7.0-py3-none-any.whl (44 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.6/44.6 kB\u001b[0m \u001b[31m6.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting fancy-einsum>=0.0.3 (from transformer_lens)\n",
            "  Downloading fancy_einsum-0.0.3-py3-none-any.whl (6.2 kB)\n",
            "Collecting jaxtyping>=0.2.11 (from transformer_lens)\n",
            "  Downloading jaxtyping-0.2.28-py3-none-any.whl (40 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m40.7/40.7 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.24 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (1.25.2)\n",
            "Requirement already satisfied: pandas>=1.1.5 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (2.0.3)\n",
            "Requirement already satisfied: rich>=12.6.0 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (13.7.1)\n",
            "Requirement already satisfied: torch!=2.0,!=2.1.0,>=1.10 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (2.2.1+cu121)\n",
            "Requirement already satisfied: tqdm>=4.64.1 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (4.66.2)\n",
            "Requirement already satisfied: transformers>=4.37.2 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (4.38.2)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (4.10.0)\n",
            "Collecting wandb>=0.13.5 (from transformer_lens)\n",
            "  Downloading wandb-0.16.6-py3-none-any.whl (2.2 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m25.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from accelerate>=0.23.0->transformer_lens) (24.0)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate>=0.23.0->transformer_lens) (5.9.5)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from accelerate>=0.23.0->transformer_lens) (6.0.1)\n",
            "Requirement already satisfied: huggingface-hub in /usr/local/lib/python3.10/dist-packages (from accelerate>=0.23.0->transformer_lens) (0.20.3)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from accelerate>=0.23.0->transformer_lens) (0.4.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets>=2.7.1->transformer_lens) (3.13.3)\n",
            "Requirement already satisfied: pyarrow>=12.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets>=2.7.1->transformer_lens) (14.0.2)\n",
            "Requirement already satisfied: pyarrow-hotfix in /usr/local/lib/python3.10/dist-packages (from datasets>=2.7.1->transformer_lens) (0.6)\n",
            "Collecting dill<0.3.9,>=0.3.0 (from datasets>=2.7.1->transformer_lens)\n",
            "  Downloading dill-0.3.8-py3-none-any.whl (116 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m16.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.10/dist-packages (from datasets>=2.7.1->transformer_lens) (2.31.0)\n",
            "Collecting xxhash (from datasets>=2.7.1->transformer_lens)\n",
            "  Downloading xxhash-3.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (194 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.1/194.1 kB\u001b[0m \u001b[31m25.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting multiprocess (from datasets>=2.7.1->transformer_lens)\n",
            "  Downloading multiprocess-0.70.16-py310-none-any.whl (134 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m19.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: fsspec[http]<=2024.2.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from datasets>=2.7.1->transformer_lens) (2023.6.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets>=2.7.1->transformer_lens) (3.9.3)\n",
            "Collecting typeguard==2.13.3 (from jaxtyping>=0.2.11->transformer_lens)\n",
            "  Downloading typeguard-2.13.3-py3-none-any.whl (17 kB)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.5->transformer_lens) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.5->transformer_lens) (2023.4)\n",
            "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.5->transformer_lens) (2024.1)\n",
            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich>=12.6.0->transformer_lens) (3.0.0)\n",
            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich>=12.6.0->transformer_lens) (2.16.1)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens) (3.1.3)\n",
            "Collecting nvidia-cuda-nvrtc-cu12==12.1.105 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (23.7 MB)\n",
            "Collecting nvidia-cuda-runtime-cu12==12.1.105 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (823 kB)\n",
            "Collecting nvidia-cuda-cupti-cu12==12.1.105 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (14.1 MB)\n",
            "Collecting nvidia-cudnn-cu12==8.9.2.26 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl (731.7 MB)\n",
            "Collecting nvidia-cublas-cu12==12.1.3.1 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl (410.6 MB)\n",
            "Collecting nvidia-cufft-cu12==11.0.2.54 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl (121.6 MB)\n",
            "Collecting nvidia-curand-cu12==10.3.2.106 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl (56.5 MB)\n",
            "Collecting nvidia-cusolver-cu12==11.4.5.107 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl (124.2 MB)\n",
            "Collecting nvidia-cusparse-cu12==12.1.0.106 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl (196.0 MB)\n",
            "Collecting nvidia-nccl-cu12==2.19.3 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl (166.0 MB)\n",
            "Collecting nvidia-nvtx-cu12==12.1.105 (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (99 kB)\n",
            "Requirement already satisfied: triton==2.2.0 in /usr/local/lib/python3.10/dist-packages (from torch!=2.0,!=2.1.0,>=1.10->transformer_lens) (2.2.0)\n",
            "Collecting nvidia-nvjitlink-cu12 (from nvidia-cusolver-cu12==11.4.5.107->torch!=2.0,!=2.1.0,>=1.10->transformer_lens)\n",
            "  Using cached nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (21.1 MB)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.37.2->transformer_lens) (2023.12.25)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers>=4.37.2->transformer_lens) (0.15.2)\n",
            "Requirement already satisfied: Click!=8.0.0,>=7.1 in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (8.1.7)\n",
            "Collecting GitPython!=3.1.29,>=1.0.0 (from wandb>=0.13.5->transformer_lens)\n",
            "  Downloading GitPython-3.1.43-py3-none-any.whl (207 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.3/207.3 kB\u001b[0m \u001b[31m29.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting sentry-sdk>=1.0.0 (from wandb>=0.13.5->transformer_lens)\n",
            "  Downloading sentry_sdk-1.45.0-py2.py3-none-any.whl (267 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m267.1/267.1 kB\u001b[0m \u001b[31m31.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting docker-pycreds>=0.4.0 (from wandb>=0.13.5->transformer_lens)\n",
            "  Downloading docker_pycreds-0.4.0-py2.py3-none-any.whl (9.0 kB)\n",
            "Collecting setproctitle (from wandb>=0.13.5->transformer_lens)\n",
            "  Downloading setproctitle-1.3.3-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (30 kB)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (67.7.2)\n",
            "Requirement already satisfied: appdirs>=1.4.3 in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (1.4.4)\n",
            "Requirement already satisfied: protobuf!=4.21.0,<5,>=3.19.0 in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (3.20.3)\n",
            "Requirement already satisfied: six>=1.4.0 in /usr/local/lib/python3.10/dist-packages (from docker-pycreds>=0.4.0->wandb>=0.13.5->transformer_lens) (1.16.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (23.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (1.4.1)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (6.0.5)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (1.9.4)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (4.0.3)\n",
            "Collecting gitdb<5,>=4.0.1 (from GitPython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer_lens)\n",
            "  Downloading gitdb-4.0.11-py3-none-any.whl (62 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 kB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich>=12.6.0->transformer_lens) (0.1.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets>=2.7.1->transformer_lens) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets>=2.7.1->transformer_lens) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets>=2.7.1->transformer_lens) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets>=2.7.1->transformer_lens) (2024.2.2)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch!=2.0,!=2.1.0,>=1.10->transformer_lens) (2.1.5)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch!=2.0,!=2.1.0,>=1.10->transformer_lens) (1.3.0)\n",
            "Collecting smmap<6,>=3.0.1 (from gitdb<5,>=4.0.1->GitPython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer_lens)\n",
            "  Downloading smmap-5.0.1-py3-none-any.whl (24 kB)\n",
            "Installing collected packages: better-abc, xxhash, typeguard, smmap, setproctitle, sentry-sdk, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, fancy-einsum, einops, docker-pycreds, dill, beartype, nvidia-cusparse-cu12, nvidia-cudnn-cu12, multiprocess, jaxtyping, gitdb, nvidia-cusolver-cu12, GitPython, wandb, datasets, accelerate, transformer_lens\n",
            "Successfully installed GitPython-3.1.43 accelerate-0.29.2 beartype-0.14.1 better-abc-0.0.3 datasets-2.18.0 dill-0.3.8 docker-pycreds-0.4.0 einops-0.7.0 fancy-einsum-0.0.3 gitdb-4.0.11 jaxtyping-0.2.28 multiprocess-0.70.16 nvidia-cublas-cu12-12.1.3.1 nvidia-cuda-cupti-cu12-12.1.105 nvidia-cuda-nvrtc-cu12-12.1.105 nvidia-cuda-runtime-cu12-12.1.105 nvidia-cudnn-cu12-8.9.2.26 nvidia-cufft-cu12-11.0.2.54 nvidia-curand-cu12-10.3.2.106 nvidia-cusolver-cu12-11.4.5.107 nvidia-cusparse-cu12-12.1.0.106 nvidia-nccl-cu12-2.19.3 nvidia-nvjitlink-cu12-12.4.127 nvidia-nvtx-cu12-12.1.105 sentry-sdk-1.45.0 setproctitle-1.3.3 smmap-5.0.1 transformer_lens-1.15.0 typeguard-2.13.3 wandb-0.16.6 xxhash-3.4.1\n",
            "Collecting circuitsvis\n",
            "  Downloading circuitsvis-1.43.2-py3-none-any.whl (1.8 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 MB\u001b[0m \u001b[31m13.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: importlib-metadata>=5.1.0 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (7.1.0)\n",
            "Requirement already satisfied: numpy>=1.24 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (1.25.2)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (8.9.2.26)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (10.3.2.106)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.0.106)\n",
            "Collecting nvidia-nccl-cu12==2.18.1 (from circuitsvis)\n",
            "  Downloading nvidia_nccl_cu12-2.18.1-py3-none-manylinux1_x86_64.whl (209.8 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m209.8/209.8 MB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n",
            "Requirement already satisfied: torch>=1.10 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (2.2.1+cu121)\n",
            "Collecting triton==2.1.0 (from circuitsvis)\n",
            "  Downloading triton-2.1.0-0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (89.2 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m89.2/89.2 MB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->circuitsvis) (12.4.127)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from triton==2.1.0->circuitsvis) (3.13.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.10/dist-packages (from importlib-metadata>=5.1.0->circuitsvis) (3.18.1)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (4.10.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (2023.6.0)\n",
            "INFO: pip is looking at multiple versions of torch to determine which version is compatible with other requirements. This could take a while.\n",
            "Collecting torch>=1.10 (from circuitsvis)\n",
            "  Downloading torch-2.2.2-cp310-cp310-manylinux1_x86_64.whl (755.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m755.5/755.5 MB\u001b[0m \u001b[31m2.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Downloading torch-2.2.1-cp310-cp310-manylinux1_x86_64.whl (755.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m755.5/755.5 MB\u001b[0m \u001b[31m2.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Downloading torch-2.2.0-cp310-cp310-manylinux1_x86_64.whl (755.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m755.5/755.5 MB\u001b[0m \u001b[31m1.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Downloading torch-2.1.2-cp310-cp310-manylinux1_x86_64.whl (670.2 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m670.2/670.2 MB\u001b[0m \u001b[31m1.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.10->circuitsvis) (2.1.5)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.10->circuitsvis) (1.3.0)\n",
            "Installing collected packages: triton, nvidia-nccl-cu12, torch, circuitsvis\n",
            "  Attempting uninstall: triton\n",
            "    Found existing installation: triton 2.2.0\n",
            "    Uninstalling triton-2.2.0:\n",
            "      Successfully uninstalled triton-2.2.0\n",
            "  Attempting uninstall: nvidia-nccl-cu12\n",
            "    Found existing installation: nvidia-nccl-cu12 2.19.3\n",
            "    Uninstalling nvidia-nccl-cu12-2.19.3:\n",
            "      Successfully uninstalled nvidia-nccl-cu12-2.19.3\n",
            "  Attempting uninstall: torch\n",
            "    Found existing installation: torch 2.2.1+cu121\n",
            "    Uninstalling torch-2.2.1+cu121:\n",
            "      Successfully uninstalled torch-2.2.1+cu121\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "torchaudio 2.2.1+cu121 requires torch==2.2.1, but you have torch 2.1.2 which is incompatible.\n",
            "torchtext 0.17.1 requires torch==2.2.1, but you have torch 2.1.2 which is incompatible.\n",
            "torchvision 0.17.1+cu121 requires torch==2.2.1, but you have torch 2.1.2 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed circuitsvis-1.43.2 nvidia-nccl-cu12-2.18.1 torch-2.1.2 triton-2.1.0\n",
            "\n",
            "\u001b[31m================================================================================\u001b[m\n",
            "\u001b[31m================================================================================\u001b[m\n",
            "\n",
            "  \u001b[1m\u001b[33m                            DEPRECATION WARNING                            \u001b[m\n",
            "\n",
            "    \u001b[1m\u001b[4m Node.js 16.x is no longer actively supported!\u001b[m\n",
            "\n",
            "  \u001b[1mYou will not receive security or critical stability updates\u001b[m for this version.\n",
            "\n",
            "  You should migrate to a supported version of Node.js as soon as possible.\n",
            "  Use the installation script that corresponds to the version of Node.js you\n",
            "  wish to install. e.g.\n",
            "  \n",
            "   * \u001b[31mhttps://deb.nodesource.com/setup_16.x — Node.js 16 \"Gallium\" \u001b[1m(deprecated)\u001b[m\n",
            "   * \u001b[32mhttps://deb.nodesource.com/setup_18.x — Node.js 18 \"Hydrogen\" (Maintenance)\u001b[m\n",
            "   * \u001b[31mhttps://deb.nodesource.com/setup_19.x — Node.js 19 \"Nineteen\" \u001b[1m(deprecated)\u001b[m\n",
            "   * \u001b[1m\u001b[32mhttps://deb.nodesource.com/setup_20.x — Node.js 20 LTS \"Iron\" (recommended)\u001b[m\n",
            "   * \u001b[32mhttps://deb.nodesource.com/setup_21.x — Node.js 21 \"Iron\" (current)\u001b[m\n",
            "   \n",
            "\n",
            "\n",
            "  Please see \u001b[1mhttps://github.com/nodejs/Release\u001b[m for details about which\n",
            "  version may be appropriate for you.\n",
            "\n",
            "  The \u001b[32m\u001b[1mNodeSource\u001b[m Node.js distributions repository contains\n",
            "  information both about supported versions of Node.js and supported Linux\n",
            "  distributions. To learn more about usage, see the repository:\n",
            "   \u001b[4m\u001b[1mhttps://github.com/nodesource/distributions\u001b[m\n",
            "\n",
            "\u001b[31m================================================================================\u001b[m\n",
            "\u001b[31m================================================================================\u001b[m\n",
            "\n",
            "Continuing in 10 seconds ...\n",
            "\n",
            "\u001b[38;5;79m2024-04-10 21:50:21 - Installing pre-requisites\u001b[0m\n",
            "Get:1 https://cloud.r-project.org/bin/linux/ubuntu jammy-cran40/ InRelease [3,626 B]\n",
            "Hit:2 https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64  InRelease\n",
            "Get:3 http://security.ubuntu.com/ubuntu jammy-security InRelease [110 kB]\n",
            "Hit:4 http://archive.ubuntu.com/ubuntu jammy InRelease\n",
            "Get:5 http://archive.ubuntu.com/ubuntu jammy-updates InRelease [119 kB]\n",
            "Get:6 http://security.ubuntu.com/ubuntu jammy-security/main amd64 Packages [1,691 kB]\n",
            "Hit:7 https://ppa.launchpadcontent.net/c2d4u.team/c2d4u4.0+/ubuntu jammy InRelease\n",
            "Hit:8 https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu jammy InRelease\n",
            "Get:9 http://security.ubuntu.com/ubuntu jammy-security/universe amd64 Packages [1,081 kB]\n",
            "Hit:10 http://archive.ubuntu.com/ubuntu jammy-backports InRelease\n",
            "Hit:11 https://ppa.launchpadcontent.net/graphics-drivers/ppa/ubuntu jammy InRelease\n",
            "Get:12 http://archive.ubuntu.com/ubuntu jammy-updates/universe amd64 Packages [1,357 kB]\n",
            "Hit:13 https://ppa.launchpadcontent.net/ubuntugis/ppa/ubuntu jammy InRelease\n",
            "Get:14 http://archive.ubuntu.com/ubuntu jammy-updates/restricted amd64 Packages [2,173 kB]\n",
            "Get:15 http://archive.ubuntu.com/ubuntu jammy-updates/main amd64 Packages [1,970 kB]\n",
            "Fetched 8,505 kB in 2s (3,966 kB/s)\n",
            "Reading package lists... Done\n",
            "Reading package lists... Done\n",
            "Building dependency tree... Done\n",
            "Reading state information... Done\n",
            "ca-certificates is already the newest version (20230311ubuntu0.22.04.1).\n",
            "curl is already the newest version (7.81.0-1ubuntu1.16).\n",
            "gnupg is already the newest version (2.2.27-3ubuntu2.1).\n",
            "gnupg set to manually installed.\n",
            "The following NEW packages will be installed:\n",
            "  apt-transport-https\n",
            "0 upgraded, 1 newly installed, 0 to remove and 50 not upgraded.\n",
            "Need to get 1,510 B of archives.\n",
            "After this operation, 170 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu jammy-updates/universe amd64 apt-transport-https all 2.4.12 [1,510 B]\n",
            "Fetched 1,510 B in 0s (4,244 B/s)\n",
            "Selecting previously unselected package apt-transport-https.\n",
            "(Reading database ... 121752 files and directories currently installed.)\n",
            "Preparing to unpack .../apt-transport-https_2.4.12_all.deb ...\n",
            "Unpacking apt-transport-https (2.4.12) ...\n",
            "Setting up apt-transport-https (2.4.12) ...\n",
            "Hit:1 https://cloud.r-project.org/bin/linux/ubuntu jammy-cran40/ InRelease\n",
            "Hit:2 https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64  InRelease\n",
            "Get:3 https://deb.nodesource.com/node_16.x nodistro InRelease [12.1 kB]\n",
            "Hit:4 http://security.ubuntu.com/ubuntu jammy-security InRelease\n",
            "Get:5 https://deb.nodesource.com/node_16.x nodistro/main amd64 Packages [7,253 B]\n",
            "Hit:6 http://archive.ubuntu.com/ubuntu jammy InRelease\n",
            "Hit:7 http://archive.ubuntu.com/ubuntu jammy-updates InRelease\n",
            "Hit:8 http://archive.ubuntu.com/ubuntu jammy-backports InRelease\n",
            "Hit:9 https://ppa.launchpadcontent.net/c2d4u.team/c2d4u4.0+/ubuntu jammy InRelease\n",
            "Hit:10 https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu jammy InRelease\n",
            "Hit:11 https://ppa.launchpadcontent.net/graphics-drivers/ppa/ubuntu jammy InRelease\n",
            "Hit:12 https://ppa.launchpadcontent.net/ubuntugis/ppa/ubuntu jammy InRelease\n",
            "Fetched 19.4 kB in 1s (17.0 kB/s)\n",
            "Reading package lists... Done\n",
            "\u001b[1;32m2024-04-10 21:50:31 - Repository configured successfully. To install Node.js, run: apt-get install nodejs -y\u001b[0m\n",
            "Reading package lists... Done\n",
            "Building dependency tree... Done\n",
            "Reading state information... Done\n",
            "The following NEW packages will be installed:\n",
            "  nodejs\n",
            "0 upgraded, 1 newly installed, 0 to remove and 50 not upgraded.\n",
            "Need to get 27.5 MB of archives.\n",
            "After this operation, 128 MB of additional disk space will be used.\n",
            "Get:1 https://deb.nodesource.com/node_16.x nodistro/main amd64 nodejs amd64 16.20.2-1nodesource1 [27.5 MB]\n",
            "Fetched 27.5 MB in 0s (65.5 MB/s)\n",
            "debconf: unable to initialize frontend: Dialog\n",
            "debconf: (No usable dialog-like program is installed, so the dialog based frontend cannot be used. at /usr/share/perl5/Debconf/FrontEnd/Dialog.pm line 78, <> line 1.)\n",
            "debconf: falling back to frontend: Readline\n",
            "debconf: unable to initialize frontend: Readline\n",
            "debconf: (This frontend requires a controlling tty.)\n",
            "debconf: falling back to frontend: Teletype\n",
            "dpkg-preconfigure: unable to re-open stdin: \n",
            "Selecting previously unselected package nodejs.\n",
            "(Reading database ... 121756 files and directories currently installed.)\n",
            "Preparing to unpack .../nodejs_16.20.2-1nodesource1_amd64.deb ...\n",
            "Unpacking nodejs (16.20.2-1nodesource1) ...\n",
            "Setting up nodejs (16.20.2-1nodesource1) ...\n",
            "Processing triggers for man-db (2.10.2-1) ...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Imuyk0gsW7Ta"
      },
      "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 plotly.graph_objects as go\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 transformer_lens.utils as utils\n",
        "from transformer_lens import ActivationCache, HookedTransformer"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.set_grad_enabled(False)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zcQRqZRbgDop",
        "outputId": "a42b0359-7311-498c-d282-35db7767b31e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch.autograd.grad_mode.set_grad_enabled at 0x7f3c39777a00>"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "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()"
      ],
      "metadata": {
        "id": "Q8p4l-vVPOxv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Idea 1: Exploring Linearity of Time Within GPT 2 Small (short-lived)\n",
        "\n",
        "The motivation of this work is to get my feet wet in the world of mechanistic interpretability research by learning more about how models represent or learn about time.\n",
        "\n",
        "I choose this idea because I am imagining a future in which GPT-8 (or\n",
        " superintellegence equivalent), might represent time in a non-intuitive manner. Maybe something like the aliens in the movie Arrival. This new representation for time could allow the model to make tons of progress on unsolved physics questions.\n"
      ],
      "metadata": {
        "id": "_jabSNSocBMh"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "Y2twJjzwb0FZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# NBVAL_IGNORE_OUTPUT\n",
        "model = HookedTransformer.from_pretrained(\n",
        "    \"gpt2-small\",\n",
        "    center_unembed=True,\n",
        "    center_writing_weights=True,\n",
        "    fold_ln=True,\n",
        "    refactor_factored_attn_matrices=True,\n",
        ")\n",
        "\n",
        "# Get the default device used\n",
        "device: torch.device = utils.get_device()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 386,
          "referenced_widgets": [
            "7f55f54e9166432ea0c19258c231d805",
            "8bccd228f4c34fa1a707436fcf15ff6e",
            "54c721f521ea439bbb55d0c5fac036f2",
            "7e241ad459ec4ce5a3e1bec252d9c8c0",
            "c733a5e8a6d54d0a9baee1ffaf06cacf",
            "0aca003cf5784b6e827dba0fe138c258",
            "3e16588ca663488cb36ec04c85c31966",
            "d943ed63a4a04652a5f486421d3d3a47",
            "2f2530d4ce8346e6abeb6e19c8bd9748",
            "f47c669fc90f4cc1b324236aedea67de",
            "be1d4144e7d04a709a30da414157179f",
            "0bc4b0c6b0074e988eb8ed8ed4930b87",
            "5a8cda067f4d4e41ac3a1e1305ba7c6a",
            "da7158f6d8b746ad8e33af3c602faffe",
            "b03e4c8d71ed4503b3252e4ed94ba231",
            "c4dc5351e77a4f5ca9574fdbfef7dc16",
            "f6c946a115bc44a3ac2a0ea6e5aab030",
            "6d30c8b9f964485ebc514f65d376f14b",
            "5335c0a6cacb4636a7c978affdb41b59",
            "1064183c59b24c08ba3fba08aa3af3e3",
            "766464eb9a8847b1ab4a91eef83eb41a",
            "2c3cd12cc9ed45208457318d42d25ab6",
            "10102db928204e52844ec4a2df91f151",
            "6b4e7ebf5dba479d96cf7e112015adb3",
            "0f1c58b17446460c8ae0cd22bc0a2308",
            "5b41ca63cbd74fb0a85cff6a39c26232",
            "d474575397b54c4fb8b11d0e60d218ed",
            "ffb6da6267af4bb7a300e464e71b4c4c",
            "fc56d0eb6f77459fb0f9849795fd13c8",
            "ea8668c29dd64090ad3ec3bcbd5a7c14",
            "a0d4f4b142aa4fa196d8dd0e3d7f6501",
            "d784c2aa64f84c3dab18174bf1a6a87f",
            "4d72ad1f68734a81a65827d7fdc6d2a8",
            "e1025ffb7289424e96ecab795e45dfc4",
            "9bbf1cbb60d64e28acddfb5b53031946",
            "e20330a8dad8437199e70294db81b287",
            "db1dd118936c4d809923ffb4fbd9a45c",
            "cf90a0933cfa4ab4b51f2522388958b1",
            "bd439b43beec487cbfe48c3fc1605743",
            "46f3c5e12d524ff1828ae66141159cb1",
            "2423a32249ed4810bb0d9b3e5a1d8cf3",
            "b9b8e02e98814811b733adc8a319ee26",
            "db8102a3d5ef46b0b03db43158f88a8e",
            "aa7514eaefb341549502d32d9cce9654",
            "690856ce7330485884a177510870fd10",
            "a95f7604fc8c41f09d5a1a5cc8531c98",
            "c3dba86a1b6f4e52b6672eb6166b2b12",
            "0c05dc2e589b4a87bed0551e1b51dc0d",
            "1062dfcafd53461f96bc221946645507",
            "8148cf02742a446796b65fabbd89aa81",
            "e44e6a2947e84c48a8a2a1b28627311d",
            "ea24da5bd37a45d98fff3893604da4fd",
            "aa50ab4ed07a4775a3724331cf552750",
            "a39d69d6a83344cf9d2ea990a6cece16",
            "db512f2b02e8421c98d05490b20d71d0",
            "eaf6689e3c4b4684a588788ec47e3f1e",
            "6e514e6b0a224bf983e3d47c5cb77629",
            "25e0678952474d26b970cfbb8994dce1",
            "67c8e2ca63ba46998d921247d092e259",
            "8d0a3260fb824461a241753bcabf9799",
            "916d05028242446f9c75cd54b69aaa80",
            "5e285d4cdc3f4201924d908c048a2698",
            "4c96f54fac5945e08265dd063758a057",
            "d5340f310e4440d7a98ecab251517036",
            "c7f7b6174d424e12916c936c4adecfd0",
            "0fba77664b59405abf46000eb03f0e51",
            "6ae175b1adc2410f940b039147583243",
            "e3d2bc5abc424c19849a918319c70c7a",
            "90cca51e0a5048b6ac4de6392c4013bf",
            "6051da612e9a487db28e3ed79ec30eaf",
            "1e740db45cd44908abe1dfb6634c4b0d",
            "9e96e3c1ee0546a7852772afe2d04f3b",
            "7dc5694af1c34555bee10cee5a943d69",
            "174595f5dc3c4fff8740d9551718b725",
            "5ed48075c9a54d2ba4a6f0c436fd3127",
            "e44e44f9b8ef413399629997790bc6e0",
            "562f1defdada4dc9893df1e7bda02596"
          ]
        },
        "id": "LXl6qcyobqMb",
        "outputId": "a3c3139f-9597-47a1-9d86-b1701d9f8939"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:88: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "7f55f54e9166432ea0c19258c231d805"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "model.safetensors:   0%|          | 0.00/548M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0bc4b0c6b0074e988eb8ed8ed4930b87"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "10102db928204e52844ec4a2df91f151"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "e1025ffb7289424e96ecab795e45dfc4"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "690856ce7330485884a177510870fd10"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "eaf6689e3c4b4684a588788ec47e3f1e"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "6ae175b1adc2410f940b039147583243"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loaded pretrained model gpt2-small into HookedTransformer\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "I first attempted the type of prompts below, where I was testing if the model could correctly understand time and what 2 hours later meant. Unfortunately GPT-2 small was terrible at this, so I don't think its worth digging deeper."
      ],
      "metadata": {
        "id": "dvhbbmNJ4Mvq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "example_prompt = \"In New York the time is 8:00 PM. In Los Angeles the time is\"\n",
        "example_answer = \"5\"\n",
        "example_prompt_written = \"In New York the time is eight PM. In Los Angeles the time is\"\n",
        "example_answer_written = \"five\"\n",
        "utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)\n",
        "utils.test_prompt(example_prompt_written, example_answer_written, model, prepend_bos=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 565
        },
        "id": "v7HuQzY5XMuT",
        "outputId": "99e7c83c-c3c2-4c5e-920d-b151bd2a3db1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tokenized prompt: ['<|endoftext|>', 'In', ' New', ' York', ' the', ' time', ' is', ' 8', ':', '00', ' PM', '.', ' In', ' Los', ' Angeles', ' the', ' time', ' is']\n",
            "Tokenized answer: [' 5']\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Performance on answer token:\n",
              "\u001b[1mRank: \u001b[0m\u001b[1;36m6\u001b[0m\u001b[1m        Logit: \u001b[0m\u001b[1;36m16.48\u001b[0m\u001b[1m Prob:  \u001b[0m\u001b[1;36m6.49\u001b[0m\u001b[1m% Token: | \u001b[0m\u001b[1;36m5\u001b[0m\u001b[1m|\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Performance on answer token:\n",
              "<span style=\"font-weight: bold\">Rank: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span><span style=\"font-weight: bold\">        Logit: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16.48</span><span style=\"font-weight: bold\"> Prob:  </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.49</span><span style=\"font-weight: bold\">% Token: | </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span><span style=\"font-weight: bold\">|</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Top 0th token. Logit: 17.43 Prob: 16.80% Token: | 9|\n",
            "Top 1th token. Logit: 17.38 Prob: 16.00% Token: | 8|\n",
            "Top 2th token. Logit: 16.89 Prob:  9.80% Token: | 7|\n",
            "Top 3th token. Logit: 16.75 Prob:  8.46% Token: | 6|\n",
            "Top 4th token. Logit: 16.56 Prob:  7.01% Token: | 10|\n",
            "Top 5th token. Logit: 16.54 Prob:  6.89% Token: | 2|\n",
            "Top 6th token. Logit: 16.48 Prob:  6.49% Token: | 5|\n",
            "Top 7th token. Logit: 16.29 Prob:  5.36% Token: | 11|\n",
            "Top 8th token. Logit: 16.26 Prob:  5.19% Token: | 4|\n",
            "Top 9th token. Logit: 16.20 Prob:  4.92% Token: | 1|\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' 5'\u001b[0m, \u001b[1;36m6\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Ranks of the answer tokens:</span> <span style=\"font-weight: bold\">[(</span><span style=\"color: #008000; text-decoration-color: #008000\">' 5'</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span><span style=\"font-weight: bold\">)]</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tokenized prompt: ['<|endoftext|>', 'In', ' New', ' York', ' the', ' time', ' is', ' eight', ' PM', '.', ' In', ' Los', ' Angeles', ' the', ' time', ' is']\n",
            "Tokenized answer: [' five']\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Performance on answer token:\n",
              "\u001b[1mRank: \u001b[0m\u001b[1;36m9\u001b[0m\u001b[1m        Logit: \u001b[0m\u001b[1;36m14.70\u001b[0m\u001b[1m Prob:  \u001b[0m\u001b[1;36m2.95\u001b[0m\u001b[1m% Token: | five|\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Performance on answer token:\n",
              "<span style=\"font-weight: bold\">Rank: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9</span><span style=\"font-weight: bold\">        Logit: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14.70</span><span style=\"font-weight: bold\"> Prob:  </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.95</span><span style=\"font-weight: bold\">% Token: | five|</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Top 0th token. Logit: 15.66 Prob:  7.70% Token: | 8|\n",
            "Top 1th token. Logit: 15.61 Prob:  7.31% Token: | 10|\n",
            "Top 2th token. Logit: 15.57 Prob:  7.07% Token: | 9|\n",
            "Top 3th token. Logit: 15.37 Prob:  5.75% Token: | 11|\n",
            "Top 4th token. Logit: 15.11 Prob:  4.44% Token: | midnight|\n",
            "Top 5th token. Logit: 15.08 Prob:  4.32% Token: | 7|\n",
            "Top 6th token. Logit: 14.99 Prob:  3.95% Token: | ten|\n",
            "Top 7th token. Logit: 14.95 Prob:  3.78% Token: | 12|\n",
            "Top 8th token. Logit: 14.87 Prob:  3.50% Token: | seven|\n",
            "Top 9th token. Logit: 14.70 Prob:  2.95% Token: | five|\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' five'\u001b[0m, \u001b[1;36m9\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Ranks of the answer tokens:</span> <span style=\"font-weight: bold\">[(</span><span style=\"color: #008000; text-decoration-color: #008000\">' five'</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9</span><span style=\"font-weight: bold\">)]</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "example_prompt = \"John gets one new coin every day. Yesterday he had two coins. After getting a new coin, today he has\"\n",
        "example_answer = \" three\"\n",
        "utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "rmCp3iIt5le4",
        "outputId": "14a5ff35-c541-4480-f198-cab7990b3f44"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tokenized prompt: ['<|endoftext|>', 'John', ' gets', ' one', ' new', ' coin', ' every', ' day', '.', ' Yesterday', ' he', ' had', ' two', ' coins', '.', ' After', ' getting', ' a', ' new', ' coin', ',', ' today', ' he', ' has']\n",
            "Tokenized answer: [' three']\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Performance on answer token:\n",
              "\u001b[1mRank: \u001b[0m\u001b[1;36m1\u001b[0m\u001b[1m        Logit: \u001b[0m\u001b[1;36m16.76\u001b[0m\u001b[1m Prob: \u001b[0m\u001b[1;36m17.13\u001b[0m\u001b[1m% Token: | three|\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Performance on answer token:\n",
              "<span style=\"font-weight: bold\">Rank: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">        Logit: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16.76</span><span style=\"font-weight: bold\"> Prob: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">17.13</span><span style=\"font-weight: bold\">% Token: | three|</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Top 0th token. Logit: 17.26 Prob: 28.07% Token: | two|\n",
            "Top 1th token. Logit: 16.76 Prob: 17.13% Token: | three|\n",
            "Top 2th token. Logit: 16.69 Prob: 15.92% Token: | one|\n",
            "Top 3th token. Logit: 15.73 Prob:  6.08% Token: | four|\n",
            "Top 4th token. Logit: 15.39 Prob:  4.33% Token: | a|\n",
            "Top 5th token. Logit: 14.86 Prob:  2.55% Token: | five|\n",
            "Top 6th token. Logit: 14.71 Prob:  2.20% Token: | just|\n",
            "Top 7th token. Logit: 14.57 Prob:  1.92% Token: | only|\n",
            "Top 8th token. Logit: 14.28 Prob:  1.43% Token: | six|\n",
            "Top 9th token. Logit: 14.07 Prob:  1.17% Token: | 2|\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' three'\u001b[0m, \u001b[1;36m1\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Ranks of the answer tokens:</span> <span style=\"font-weight: bold\">[(</span><span style=\"color: #008000; text-decoration-color: #008000\">' three'</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">)]</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "After trying a lot of math type of tasks, I realized gpt2-small is just not good at numbers. I'm going to try to find a different task related to language for this specific project, just so that I can get started.\n",
        "\n",
        "However, I think there is an interesting question that I found which is how the model interprets tokens of numbers in numerical form (ex: 1,2,3) vs tokens of numbers in english (ex: one, two, three). I think there could be an interesting analysis done of how the model acts differently when the only difference in the prompt is how the number is represented. After experimenting for a few minutes, I see slight differences in the outputs."
      ],
      "metadata": {
        "id": "wwvM2g4RIAd1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "example_prompt = \"A mouse is small, but an elephant is\"\n",
        "example_answer = \"big\"\n",
        "g= utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "UX8VqaSP-Jbq",
        "outputId": "f839ca7c-0331-462d-b607-099f0ef99c2e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tokenized prompt: ['<|endoftext|>', 'A', ' mouse', ' is', ' small', ',', ' but', ' an', ' elephant', ' is']\n",
            "Tokenized answer: [' big']\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Performance on answer token:\n",
              "\u001b[1mRank: \u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m        Logit: \u001b[0m\u001b[1;36m18.63\u001b[0m\u001b[1m Prob: \u001b[0m\u001b[1;36m26.37\u001b[0m\u001b[1m% Token: | big|\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Performance on answer token:\n",
              "<span style=\"font-weight: bold\">Rank: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">        Logit: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">18.63</span><span style=\"font-weight: bold\"> Prob: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">26.37</span><span style=\"font-weight: bold\">% Token: | big|</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Top 0th token. Logit: 18.63 Prob: 26.37% Token: | big|\n",
            "Top 1th token. Logit: 18.14 Prob: 16.10% Token: | large|\n",
            "Top 2th token. Logit: 17.25 Prob:  6.64% Token: | huge|\n",
            "Top 3th token. Logit: 17.23 Prob:  6.45% Token: | bigger|\n",
            "Top 4th token. Logit: 17.00 Prob:  5.18% Token: | larger|\n",
            "Top 5th token. Logit: 16.91 Prob:  4.69% Token: | a|\n",
            "Top 6th token. Logit: 16.41 Prob:  2.87% Token: | small|\n",
            "Top 7th token. Logit: 16.17 Prob:  2.25% Token: | much|\n",
            "Top 8th token. Logit: 15.98 Prob:  1.85% Token: | very|\n",
            "Top 9th token. Logit: 15.68 Prob:  1.38% Token: | tiny|\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' big'\u001b[0m, \u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Ranks of the answer tokens:</span> <span style=\"font-weight: bold\">[(</span><span style=\"color: #008000; text-decoration-color: #008000\">' big'</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">)]</span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "def logit_answers(logits, answer_tokens, per_prompt=False):\n",
        "    print(logits.shape)\n",
        "    final_logits = logits[:, -1,:]\n",
        "    probs = torch.softmax(final_logits, dim=-1)\n",
        "\n",
        "    answer_probs = probs.index_select(dim=-1, index=answer_tokens)\n",
        "\n",
        "    return answer_probs[:, 0].item(), answer_probs[:, 1].item()\n",
        "\n",
        "example_prompt = \"A mouse is small but an elephant is\"\n",
        "example_prompt_comma = \"A mouse is small, but an elephant is\"\n",
        "tokens_test = model.to_tokens(example_prompt, prepend_bos=True)\n",
        "tokens_test_comma = model.to_tokens(example_prompt_comma, prepend_bos=True)\n",
        "answer_tokens = [ model.to_single_token(\"big\"),\n",
        "                model.to_single_token(\"small\")]\n",
        "answer_tokens = torch.tensor(answer_tokens).to(device)\n",
        "print(answer_tokens.shape)\n",
        "\n",
        "example_answer = \"big\"\n",
        "logits, cache_1 = model.run_with_cache(\n",
        "    tokens_test\n",
        ")\n",
        "\n",
        "logits_comma, cache_2 = model.run_with_cache(tokens_test_comma)\n",
        "logits_small, logits_big = logit_answers(logits,answer_tokens)\n",
        "logits_small_comma, logits_big_comma = logit_answers(logits_comma,answer_tokens)\n",
        "print(logits_small_comma,logits_big_comma)\n",
        "trace_small = go.Bar(\n",
        "    x=['No Comma','Comma'],\n",
        "    y=[1.23,],\n",
        "    name='small',\n",
        "    marker=dict(color='blue')\n",
        ")\n",
        "\n",
        "trace_big = go.Bar(\n",
        "    x=['No Comma'],\n",
        "    y=[28.53],\n",
        "    name='big',\n",
        "    marker=dict(color='red')\n",
        ")\n",
        "\n",
        "trace_small_comma = go.Bar(\n",
        "    x=['Comma'],\n",
        "    y=[2.87],\n",
        "    name='small',\n",
        "    marker=dict(color='blue')\n",
        ")\n",
        "\n",
        "trace_big_comma = go.Bar(\n",
        "    x=['Comma'],\n",
        "    y=[26.37],\n",
        "    name='big',\n",
        "    marker=dict(color='red')\n",
        ")\n",
        "\n",
        "# Create the layout for the graph\n",
        "layout = go.Layout(\n",
        "    title='Token Probabilities for Prompts With and Without Commas',\n",
        "    xaxis=dict(title='Prompt'),\n",
        "    yaxis=dict(title='Probabilities (%)'),\n",
        "    barmode='group'\n",
        ")\n",
        "\n",
        "# Create the figure and add the traces\n",
        "fig = go.Figure(data=[trace_small, trace_small_comma, trace_big, trace_big_comma], layout=layout)\n",
        "\n",
        "# Display the graph\n",
        "fig.show()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 613
        },
        "id": "g7jm49Ki5CSm",
        "outputId": "99f0a819-4335-45ae-8ad3-6f353c862532"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([2])\n",
            "torch.Size([1, 9, 50257])\n",
            "torch.Size([1, 10, 50257])\n",
            "1.937619117597933e-06 2.1398871297151345e-07\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"5353e2ce-47e9-4a43-8ded-9444b1aca4a0\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"5353e2ce-47e9-4a43-8ded-9444b1aca4a0\")) {                    Plotly.newPlot(                        \"5353e2ce-47e9-4a43-8ded-9444b1aca4a0\",                        [{\"marker\":{\"color\":\"blue\"},\"name\":\"small\",\"x\":[\"No Comma\",\"Comma\"],\"y\":[1.23],\"type\":\"bar\"},{\"marker\":{\"color\":\"blue\"},\"name\":\"small\",\"x\":[\"Comma\"],\"y\":[2.87],\"type\":\"bar\"},{\"marker\":{\"color\":\"red\"},\"name\":\"big\",\"x\":[\"No Comma\"],\"y\":[28.53],\"type\":\"bar\"},{\"marker\":{\"color\":\"red\"},\"name\":\"big\",\"x\":[\"Comma\"],\"y\":[26.37],\"type\":\"bar\"}],                        {\"barmode\":\"group\",\"title\":{\"text\":\"Token Probabilities for Prompts With and Without Commas\"},\"xaxis\":{\"title\":{\"text\":\"Prompt\"}},\"yaxis\":{\"title\":{\"text\":\"Probabilities (%)\"}},\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('5353e2ce-47e9-4a43-8ded-9444b1aca4a0');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "example_prompt = \"A pillow is soft, but a brick is\"\n",
        "example_answer = \"hard\"\n",
        "utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 180
        },
        "id": "CsUr4WQz4eIu",
        "outputId": "f3e64cc4-ed13-4998-fc03-62c341d05a23"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "name 'utils' is not defined",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-1-43d78f7b26a1>\u001b[0m in \u001b[0;36m<cell line: 3>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mexample_prompt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"A pillow is soft, but a brick is\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mexample_answer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"hard\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest_prompt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexample_prompt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexample_answer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprepend_bos\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m: name 'utils' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Idea 2: Opposite Adjective\n",
        "I am going to study the opposite adjective task where the model has to output the opposite of an adjactive given context. Some example adjactives are old vs young and big vs small.\n",
        "\n",
        "Since we are going to fix the output to one correct answer, I am going to use the logit difference metric and ignore synonyms of the answer that might work as well (for instance \"new\" instead of \"young\" might be outputed but we won't consider it a correct or incorrect answer to simplify the problem)."
      ],
      "metadata": {
        "id": "2W-IKpIqJTmx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt_format = [\n",
        "    \"A mouse is{} but an elephant is\",\n",
        "    \"A turtle is{} but a kitten is\",\n",
        "    \"A desert is{} but an ocean is\",\n",
        "    \"A pillow is{} but a brick is\",\n",
        "    \"A rock is{} but a sponge is\",\n",
        "    \"A rose is{} but a weed is\"\n",
        "]\n",
        "prompt_format_commas = [\n",
        "    \"A mouse is{}, but an elephant is\",\n",
        "    \"A turtle is{}, but a kitten is\",\n",
        "    \"A desert is{}, but an ocean is\",\n",
        "    \"A pillow is{}, but a brick is\",\n",
        "    \"A rock is{}, but a sponge is\",\n",
        "    \"A rose is{}, but a weed is\"\n",
        "]\n",
        "names = [\n",
        "    (\" small\", \" big\"),\n",
        "    (\" old\", \" young\"),\n",
        "    (\" dry\", \" wet\"),\n",
        "    (\" soft\", \" hard\"),\n",
        "    (\" hard\", \" soft\"),\n",
        "    (\" beautiful\", \" ugly\")\n",
        "\n",
        "]\n",
        "# List of prompts\n",
        "prompts = []\n",
        "prompts_commas = []\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",
        "\n",
        "        answers.append((names[i][1], names[i][0]))\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",
        "        prompts_commas.append(prompt_format_commas[i].format(answers[-1][1]))\n",
        "answer_tokens = torch.tensor(answer_tokens).to(device)\n",
        "print(prompts)\n",
        "print(prompts_commas)\n",
        "print(answers)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FPbVMNTXI8_q",
        "outputId": "b2d3cc79-8d5a-4a8a-98ad-dec7d66e73d2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['A mouse is small but an elephant is', 'A turtle is old but a kitten is', 'A desert is dry but an ocean is', 'A pillow is soft but a brick is', 'A rock is hard but a sponge is', 'A rose is beautiful but a weed is']\n",
            "['A mouse is small, but an elephant is', 'A turtle is old, but a kitten is', 'A desert is dry, but an ocean is', 'A pillow is soft, but a brick is', 'A rock is hard, but a sponge is', 'A rose is beautiful, but a weed is']\n",
            "[(' big', ' small'), (' young', ' old'), (' wet', ' dry'), (' hard', ' soft'), (' soft', ' hard'), (' ugly', ' beautiful')]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Code to make sure positions are aligned across prompts and lengths are the same\n",
        "for prompt in prompts:\n",
        "    str_tokens = model.to_str_tokens(prompt)\n",
        "    print(\"Prompt length:\", len(str_tokens))\n",
        "    print(\"Prompt as tokens:\", str_tokens)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UavTmxJJKsl_",
        "outputId": "4b5fdeba-cb8b-475a-961e-0030b9675251"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt length: 9\n",
            "Prompt as tokens: ['<|endoftext|>', 'A', ' mouse', ' is', ' small', ' but', ' an', ' elephant', ' is']\n",
            "Prompt length: 9\n",
            "Prompt as tokens: ['<|endoftext|>', 'A', ' turtle', ' is', ' old', ' but', ' a', ' kitten', ' is']\n",
            "Prompt length: 9\n",
            "Prompt as tokens: ['<|endoftext|>', 'A', ' desert', ' is', ' dry', ' but', ' an', ' ocean', ' is']\n",
            "Prompt length: 9\n",
            "Prompt as tokens: ['<|endoftext|>', 'A', ' pillow', ' is', ' soft', ' but', ' a', ' brick', ' is']\n",
            "Prompt length: 9\n",
            "Prompt as tokens: ['<|endoftext|>', 'A', ' rock', ' is', ' hard', ' but', ' a', ' sponge', ' is']\n",
            "Prompt length: 9\n",
            "Prompt as tokens: ['<|endoftext|>', 'A', ' rose', ' is', ' beautiful', ' but', ' a', ' weed', ' is']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tokens = model.to_tokens(prompts, prepend_bos=True)\n",
        "\n",
        "# Run the model and cache all activations\n",
        "original_logits, cache = model.run_with_cache(tokens)\n",
        "\n",
        "#Getting the logits for the prompts with commas\n",
        "\n",
        "tokens_commas = model.to_tokens(prompts_commas, prepend_bos=True)\n",
        "\n",
        "original_logits_commas, cache_commas = model.run_with_cache(tokens_commas)\n"
      ],
      "metadata": {
        "id": "YQJGtYqHLBCy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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",
        "\n",
        "print(\n",
        "    \"Per prompt logit difference:\",\n",
        "    logits_to_ave_logit_diff(original_logits, answer_tokens, per_prompt=True)\n",
        "    .detach()\n",
        "    .cpu()\n",
        "    .round(decimals=3),\n",
        ")\n",
        "print(\n",
        "    \"Average logit difference:\",\n",
        "    round(logits_to_ave_logit_diff(original_logits, answer_tokens).item(), 3),\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1WJD11YZXFMw",
        "outputId": "f54f9ab0-52c8-49ad-a049-bf73e50c97ce"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Per prompt logit difference: tensor([3.1440, 2.0830, 0.3140, 0.5160, 1.7990, 2.6010])\n",
            "Average logit difference: 1.743\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\n",
        "    \"Per prompt logit difference with commas:\",\n",
        "    logits_to_ave_logit_diff(original_logits_commas, answer_tokens, per_prompt=True)\n",
        "    .detach()\n",
        "    .cpu()\n",
        "    .round(decimals=3),\n",
        ")\n",
        "original_average_logit_diff_commas = logits_to_ave_logit_diff(original_logits_commas, answer_tokens)\n",
        "print(\n",
        "    \"Average logit difference with commas:\",\n",
        "    round(logits_to_ave_logit_diff(original_logits_commas, answer_tokens).item(), 3),\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WTcrGVv2WtPD",
        "outputId": "2caccbe0-f191-4749-f1ca-5caa39a861a0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Per prompt logit difference with commas: tensor([2.2180, 1.2900, 0.3890, 0.1310, 0.8260, 1.9750])\n",
            "Average logit difference with commas: 1.138\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Hmm, this logit difference is not that large compared to the IOI example. I think part of the reason is because of the ambuguity of the answer/opposite, for instance the model has some probability that the token is \"large\" as well as \"big\".\n",
        "\n",
        "## Some general predictions as to what is happening.\n",
        "\n",
        "I assume the most relevant tokens in the prompt for the first data point are, old, but, tortoise, and kitten. It seems the \"but\" is very helpful for signifying the opposite of old, since I tried replacing \"but\" with \"then\" and the logit difference was significantly smaller.\n",
        "\n",
        "The model needs to figure out what the previous adjactive is and then predict the opposite. Therefore, maybe there is a head that attends the noun to adjective relationship and then some other component (MLP?) that is able to find the opposite of the adjective. This is likely combined with a similar layer that processes both the noun we want to describe (eg: \"kitten\") and also the opposite (eg: \"old\")."
      ],
      "metadata": {
        "id": "NKKWmarZLdvb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Logit Difference"
      ],
      "metadata": {
        "id": "R8wZGkTiV-CO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "answer_residual_directions = model.tokens_to_residual_directions(answer_tokens)\n",
        "print(\"Answer residual directions shape:\", answer_residual_directions.shape)\n",
        "logit_diff_directions = (\n",
        "    answer_residual_directions[:, 0] - answer_residual_directions[:, 1]\n",
        ")\n",
        "print(\"Logit difference directions shape:\", logit_diff_directions.shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eviJZXU2LdSg",
        "outputId": "a77c06b5-988b-41aa-a4d0-8c1b1f93bb7f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Answer residual directions shape: torch.Size([6, 2, 768])\n",
            "Logit difference directions shape: torch.Size([6, 768])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# cache syntax - resid_post is the residual stream at the end of the layer, -1 gets the final layer. The general syntax is [activation_name, layer_index, sub_layer_type].\n",
        "final_residual_stream = cache[\"resid_post\", -1]\n",
        "print(\"Final residual stream shape:\", final_residual_stream.shape)\n",
        "final_token_residual_stream = final_residual_stream[:, -1, :]\n",
        "# Apply LayerNorm scaling\n",
        "# pos_slice is the subset of the positions we take - here the final token of each prompt\n",
        "scaled_final_token_residual_stream = cache.apply_ln_to_stack(\n",
        "    final_token_residual_stream, layer=-1, pos_slice=-1\n",
        ")\n",
        "\n",
        "average_logit_diff = einsum(\n",
        "    \"batch d_model, batch d_model -> \",\n",
        "    scaled_final_token_residual_stream,\n",
        "    logit_diff_directions,\n",
        ") / len(prompts)\n",
        "original_average_logit_diff = logits_to_ave_logit_diff(original_logits, answer_tokens)\n",
        "\n",
        "print(\"Calculated average logit diff:\", round(average_logit_diff.item(), 3))\n",
        "print(\"Original logit difference:\", round(original_average_logit_diff.item(), 3))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WAhzwsCuO2Ad",
        "outputId": "7270b347-4982-4c55-fd27-9dd6d32447c2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Final residual stream shape: torch.Size([6, 9, 768])\n",
            "Calculated average logit diff: 2.043\n",
            "Original logit difference: 1.743\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, we'll take the logit lens approach to decompose where the logit differences are coming from."
      ],
      "metadata": {
        "id": "yG4EBoXKO_6Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def residual_stack_to_logit_diff(\n",
        "    residual_stack: Float[torch.Tensor, \"components batch d_model\"],\n",
        "    cache: ActivationCache,\n",
        ") -> float:\n",
        "    scaled_residual_stack = cache.apply_ln_to_stack(\n",
        "        residual_stack, layer=-1, pos_slice=-1\n",
        "    )\n",
        "    return einsum(\n",
        "        \"... batch d_model, batch d_model -> ...\",\n",
        "        scaled_residual_stack,\n",
        "        logit_diff_directions,\n",
        "    ) / len(prompts)"
      ],
      "metadata": {
        "id": "q2zsnDoOPHiA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "accumulated_residual, labels = cache.accumulated_resid(\n",
        "    layer=-1, incl_mid=True, pos_slice=-1, return_labels=True\n",
        ")\n",
        "logit_lens_logit_diffs = residual_stack_to_logit_diff(accumulated_residual, cache)\n",
        "line(\n",
        "    logit_lens_logit_diffs,\n",
        "    x=np.arange(model.cfg.n_layers * 2 + 1) / 2,\n",
        "    hover_name=labels,\n",
        "    title=\"Logit Difference From Accumulate Residual Stream\",\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "P89NpQGBPIys",
        "outputId": "55a2e7ff-0722-4455-eee5-df450a0ee210"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"8066f5d4-de6f-4dc5-ae74-b6aa6191bfff\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"8066f5d4-de6f-4dc5-ae74-b6aa6191bfff\")) {                    Plotly.newPlot(                        \"8066f5d4-de6f-4dc5-ae74-b6aa6191bfff\",                        [{\"hovertemplate\":\"\\u003cb\\u003e%{hovertext}\\u003c\\u002fb\\u003e\\u003cbr\\u003e\\u003cbr\\u003ex=%{x}\\u003cbr\\u003ey=%{y}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"hovertext\":[\"0_pre\",\"0_mid\",\"1_pre\",\"1_mid\",\"2_pre\",\"2_mid\",\"3_pre\",\"3_mid\",\"4_pre\",\"4_mid\",\"5_pre\",\"5_mid\",\"6_pre\",\"6_mid\",\"7_pre\",\"7_mid\",\"8_pre\",\"8_mid\",\"9_pre\",\"9_mid\",\"10_pre\",\"10_mid\",\"11_pre\",\"11_mid\",\"final_post\"],\"legendgroup\":\"\",\"line\":{\"color\":\"#636efa\",\"dash\":\"solid\"},\"marker\":{\"symbol\":\"circle\"},\"mode\":\"lines\",\"name\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"x\":[0.0,0.5,1.0,1.5,2.0,2.5,3.0,3.5,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.0,10.5,11.0,11.5,12.0],\"xaxis\":\"x\",\"y\":[0.005890760570764542,-0.06337859481573105,-0.10097262263298035,-0.11769779026508331,-0.12878018617630005,-0.16191592812538147,-0.15366889536380768,-0.1416461020708084,-0.11321134865283966,-0.06703229248523712,-0.08043431490659714,-0.12074237316846848,-0.12190812081098557,-0.23424939811229706,-0.14806652069091797,-0.17988072335720062,-0.10927490890026093,-0.3009437918663025,-0.11537047475576401,0.9335346221923828,1.5943492650985718,1.4560816287994385,1.6399873495101929,1.516530990600586,2.043048620223999],\"yaxis\":\"y\",\"type\":\"scatter\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"x\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"y\"}},\"legend\":{\"tracegroupgap\":0},\"title\":{\"text\":\"Logit Difference From Accumulate Residual Stream\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('8066f5d4-de6f-4dc5-ae74-b6aa6191bfff');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This stream shows that most of the performance comes from attention layer 9 as well as the last layer which looks like layer 11.\n",
        "\n",
        "More interestingly, the performance is negative up until layer 9, which means the model is more likely to predict the incorrect answer, the same adjictive in the prompt. I haven't seen this graph that stays in the negative so long and I think it might have to do with my hypothesis that the model has a lot of attention on the first noun/adjective pairing to learn that relationship.\n",
        "\n",
        "Time to dig deeper."
      ],
      "metadata": {
        "id": "aItMOX8xQyCf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "per_layer_residual, labels = cache.decompose_resid(\n",
        "    layer=-1, pos_slice=-1, return_labels=True\n",
        ")\n",
        "per_layer_logit_diffs = residual_stack_to_logit_diff(per_layer_residual, cache)\n",
        "line(per_layer_logit_diffs, hover_name=labels, title=\"Logit Difference From Each Layer\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "w7xQ9OxcQxlz",
        "outputId": "300fab5f-b0d0-4adc-9639-4cae7fa49dd1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"75ef575d-ec04-4367-b8c9-9fa79893842d\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"75ef575d-ec04-4367-b8c9-9fa79893842d\")) {                    Plotly.newPlot(                        \"75ef575d-ec04-4367-b8c9-9fa79893842d\",                        [{\"hovertemplate\":\"\\u003cb\\u003e%{hovertext}\\u003c\\u002fb\\u003e\\u003cbr\\u003e\\u003cbr\\u003ex=%{x}\\u003cbr\\u003ey=%{y}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"hovertext\":[\"embed\",\"pos_embed\",\"0_attn_out\",\"0_mlp_out\",\"1_attn_out\",\"1_mlp_out\",\"2_attn_out\",\"2_mlp_out\",\"3_attn_out\",\"3_mlp_out\",\"4_attn_out\",\"4_mlp_out\",\"5_attn_out\",\"5_mlp_out\",\"6_attn_out\",\"6_mlp_out\",\"7_attn_out\",\"7_mlp_out\",\"8_attn_out\",\"8_mlp_out\",\"9_attn_out\",\"9_mlp_out\",\"10_attn_out\",\"10_mlp_out\",\"11_attn_out\",\"11_mlp_out\"],\"legendgroup\":\"\",\"line\":{\"color\":\"#636efa\",\"dash\":\"solid\"},\"marker\":{\"symbol\":\"circle\"},\"mode\":\"lines\",\"name\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"x\":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25],\"xaxis\":\"x\",\"y\":[0.004913042299449444,0.000977715477347374,-0.06926935166120529,-0.0375940203666687,-0.01672518625855446,-0.011082402430474758,-0.03313572332262993,0.008247030898928642,0.012022806331515312,0.028434736654162407,0.046179067343473434,-0.013402005657553673,-0.04030805081129074,-0.0011657103896141052,-0.11234135925769806,0.08618294447660446,-0.03181425482034683,0.07060584425926208,-0.19166891276836395,0.18557336926460266,1.0489051342010498,0.6608145236968994,-0.1382674276828766,0.18390560150146484,-0.12345622479915619,0.526517391204834],\"yaxis\":\"y\",\"type\":\"scatter\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"x\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"y\"}},\"legend\":{\"tracegroupgap\":0},\"title\":{\"text\":\"Logit Difference From Each Layer\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('75ef575d-ec04-4367-b8c9-9fa79893842d');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This is similar to the IOI task where the attention makes the most difference in each layer since we also care a lot about moving the adjactive information around. One hypothesis for the decrease in performance in the earlier layers is that maybe the attention is simply moving the incorrect adjictive to the \"is\" token, but the MLP isn't able to process this that well. However in layer 9, maybe the attention is moving different information, like both the noun of \"kitten\" and the other adjictive \"old\" and the MLP has better information to figure out an opposite.\n",
        "\n",
        "From this graph, I can better understand how the attention sets up the MLP to process information effectively."
      ],
      "metadata": {
        "id": "MCg8SpoKSQ-d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "per_head_residual, labels = cache.stack_head_results(\n",
        "    layer=-1, pos_slice=-1, return_labels=True\n",
        ")\n",
        "per_head_logit_diffs = residual_stack_to_logit_diff(per_head_residual, cache)\n",
        "per_head_logit_diffs = einops.rearrange(\n",
        "    per_head_logit_diffs,\n",
        "    \"(layer head_index) -> layer head_index\",\n",
        "    layer=model.cfg.n_layers,\n",
        "    head_index=model.cfg.n_heads,\n",
        ")\n",
        "imshow(\n",
        "    per_head_logit_diffs,\n",
        "    labels={\"x\": \"Head\", \"y\": \"Layer\"},\n",
        "    title=\"Logit Difference From Each Head\",\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 560
        },
        "id": "_pKEeamcSQiW",
        "outputId": "dd2a2f25-9ab0-4b63-d674-e90e5fa93d78"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tried to stack head results when they weren't cached. Computing head results now\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"dd3dfe81-cca7-498c-bb09-8edc29e4e19e\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"dd3dfe81-cca7-498c-bb09-8edc29e4e19e\")) {                    Plotly.newPlot(                        \"dd3dfe81-cca7-498c-bb09-8edc29e4e19e\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-0.02118229679763317,0.011538292281329632,-0.00631052628159523,0.0013563558459281921,-0.013413703069090843,-0.02422286570072174,-0.041599102318286896,0.022703733295202255,0.008520704694092274,-0.02432115748524666,0.017299935221672058,0.004447922110557556],[0.00843064859509468,0.011206060647964478,0.004790980368852615,0.011728346347808838,0.005249427165836096,-0.006287422031164169,-0.009159574285149574,-0.018804987892508507,-0.008379396051168442,0.003910670522600412,0.008993467316031456,0.009357551112771034],[0.013590577058494091,-0.024640943855047226,-0.006494876928627491,0.009674293920397758,-0.02383767068386078,-0.008573547005653381,-0.0020972019992768764,-0.012090058997273445,0.009536845609545708,0.009338783100247383,-0.022623300552368164,0.013022209517657757],[0.004250833299010992,0.004118512850254774,-0.0021320213563740253,-0.0035438011400401592,0.004313119221478701,-0.002566268667578697,-0.00048522837460041046,0.0032204193994402885,0.0013790230732411146,-0.0005869592423550785,0.005350479390472174,0.00048591228551231325],[0.006913047283887863,0.01521211862564087,-0.003993764985352755,0.0039401850663125515,0.004815371707081795,-0.004304363392293453,0.028187748044729233,-0.021280229091644287,-0.005198651924729347,-0.00893851276487112,0.01325245015323162,0.0041809603571891785],[-0.047785013914108276,-0.003622382879257202,-0.00048214755952358246,-0.017687618732452393,0.0034946228843182325,0.014263618737459183,0.0010156277567148209,-0.006058576516807079,0.040593311190605164,-0.044284723699092865,0.0007197686354629695,0.022407621145248413],[0.03094424679875374,-0.031172428280115128,0.027751250192523003,-0.004795196931809187,-0.0067387777380645275,-0.008900457993149757,-0.019927840679883957,-0.010361410677433014,-0.0046320101246237755,-0.035804808139801025,-0.00021352464682422578,-0.03531346470117569],[-0.003438861109316349,-0.0024548920337110758,-0.009958849288523197,-0.029429331421852112,-0.0016672108322381973,0.007846878841519356,-0.002754464279860258,0.00573880597949028,0.050658732652664185,0.056175462901592255,-0.07410895079374313,-0.03415384888648987],[0.017020106315612793,-0.060182712972164154,-0.016311658546328545,-0.001255591749213636,0.023416858166456223,0.008560478687286377,-0.08667761087417603,0.0003191853757016361,-0.08750084042549133,-0.018895991146564484,-0.001131999073550105,0.023169536143541336],[-0.046013083308935165,-0.06752003729343414,0.03386766090989113,0.023042671382427216,-0.00670058187097311,-0.0009625336970202625,-0.08242174237966537,1.2604846954345703,-0.003682257141917944,-0.09276461601257324,0.017905166372656822,-0.0025231686886399984],[-0.04627007246017456,-0.362529993057251,-0.007182325702160597,-0.1551101952791214,-0.03349853307008743,-0.003574095666408539,-0.004776606336236,0.48373138904571533,0.00482338247820735,0.010850511491298676,-0.004263957031071186,-0.035184863954782486],[-0.0021513327956199646,-0.01998107135295868,-0.020695125684142113,-0.03833695501089096,0.011631982401013374,0.009166432544589043,-0.013119669631123543,-0.07825392484664917,0.011651827022433281,-0.04351112246513367,0.11233197152614594,-0.04317968338727951]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003ecolor: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Each Head\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('dd3dfe81-cca7-498c-bb09-8edc29e4e19e');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here we can see that the only heads that matter in a large way are: L9H7 and L10H7 as well as L10H1 in the negative direction. This shows that the drop in the earlier layers are due to a combination of multiple heads which have realitively small logit difference effects of about -0.1 , especially in layer 8.\n",
        "\n",
        "Let's dig deeper and do some attention head analysis."
      ],
      "metadata": {
        "id": "NfEvVE7bUB7e"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def visualize_attention_patterns(\n",
        "    heads: Union[List[int], int, Float[torch.Tensor, \"heads\"]],\n",
        "    local_cache: ActivationCache,\n",
        "    local_tokens: torch.Tensor,\n",
        "    title: Optional[str] = \"\",\n",
        "    max_width: Optional[int] = 700,\n",
        ") -> str:\n",
        "    # If a single head is given, convert to a list\n",
        "    if isinstance(heads, int):\n",
        "        heads = [heads]\n",
        "\n",
        "    # Create the plotting data\n",
        "    labels: List[str] = []\n",
        "    patterns: List[Float[torch.Tensor, \"dest_pos src_pos\"]] = []\n",
        "\n",
        "    # Assume we have a single batch item\n",
        "    batch_index = 0\n",
        "\n",
        "    for head in heads:\n",
        "        # Set the label\n",
        "        layer = head // model.cfg.n_heads\n",
        "        head_index = head % model.cfg.n_heads\n",
        "        labels.append(f\"L{layer}H{head_index}\")\n",
        "\n",
        "        # Get the attention patterns for the head\n",
        "        # Attention patterns have shape [batch, head_index, query_pos, key_pos]\n",
        "        patterns.append(local_cache[\"attn\", layer][batch_index, head_index])\n",
        "\n",
        "    # Convert the tokens to strings (for the axis labels)\n",
        "    str_tokens = model.to_str_tokens(local_tokens)\n",
        "\n",
        "    # Combine the patterns into a single tensor\n",
        "    patterns: Float[torch.Tensor, \"head_index dest_pos src_pos\"] = torch.stack(\n",
        "        patterns, dim=0\n",
        "    )\n",
        "\n",
        "    # Circuitsvis Plot (note we get the code version so we can concatenate with the title)\n",
        "    plot = attention_heads(\n",
        "        attention=patterns, tokens=str_tokens, attention_head_names=labels\n",
        "    ).show_code()\n",
        "\n",
        "    # Display the title\n",
        "    title_html = f\"<h2>{title}</h2><br/>\"\n",
        "\n",
        "    # Return the visualisation as raw code\n",
        "    return f\"<div style='max-width: {str(max_width)}px;'>{title_html + plot}</div>\""
      ],
      "metadata": {
        "id": "IEo3iGGBVjJ2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "top_k = 3\n",
        "\n",
        "top_positive_logit_attr_heads = torch.topk(\n",
        "    per_head_logit_diffs.flatten(), k=top_k\n",
        ").indices\n",
        "\n",
        "positive_html = visualize_attention_patterns(\n",
        "    top_positive_logit_attr_heads,\n",
        "    cache,\n",
        "    tokens[0],\n",
        "    f\"Top {top_k} Positive Logit Attribution Heads\",\n",
        ")\n",
        "\n",
        "top_negative_logit_attr_heads = torch.topk(\n",
        "    -per_head_logit_diffs.flatten(), k=top_k\n",
        ").indices\n",
        "\n",
        "negative_html = visualize_attention_patterns(\n",
        "    top_negative_logit_attr_heads,\n",
        "    cache,\n",
        "    tokens[0],\n",
        "    title=f\"Top {top_k} Negative Logit Attribution Heads\",\n",
        ")\n",
        "\n",
        "HTML(positive_html + negative_html)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "xx8v9DKgZ6uE",
        "outputId": "58c97085-35ae-4d28-9380-5518f4de9753"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<div style='max-width: 700px;'><h2>Top 3 Positive Logit Attribution Heads</h2><br/><div id=\"circuits-vis-c8656de9-c8f6\" style=\"margin: 15px 0;\"/>\n",
              "    <script crossorigin type=\"module\">\n",
              "    import { render, AttentionHeads } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
              "    render(\n",
              "      \"circuits-vis-c8656de9-c8f6\",\n",
              "      AttentionHeads,\n",
              "      {\"attention\": [[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9860349893569946, 0.013964998535811901, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9551160931587219, 0.027135001495480537, 0.017748890444636345, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9366254806518555, 0.015255346894264221, 0.03175199776887894, 0.01636725477874279, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9476580023765564, 0.00735104363411665, 0.015432869084179401, 0.015446932055056095, 0.014111137017607689, 0.0, 0.0, 0.0, 0.0], [0.6686481833457947, 0.010562391020357609, 0.07823613286018372, 0.054229848086833954, 0.1783769279718399, 0.009946494363248348, 0.0, 0.0, 0.0], [0.37391355633735657, 0.003834002884104848, 0.30471310019493103, 0.009269432164728642, 0.30005908012390137, 0.00566705921664834, 0.002543787006288767, 0.0, 0.0], [0.797453761100769, 0.02370990812778473, 0.026920510455965996, 0.02839079685509205, 0.0921684056520462, 0.006652399897575378, 0.006916511803865433, 0.01778770610690117, 0.0], [0.07823030650615692, 0.0009519269806332886, 0.010218545794487, 0.0013520902721211314, 0.8983247876167297, 0.0010833048727363348, 0.0005594769027084112, 0.009186544455587864, 9.30116802919656e-05]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9711289405822754, 0.028871025890111923, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.962450385093689, 0.0038900121580809355, 0.03365957736968994, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.85878986120224, 0.015433978289365768, 0.07241909205913544, 0.05335712060332298, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7033705115318298, 0.01538859587162733, 0.11321327835321426, 0.059626255184412, 0.10840138047933578, 0.0, 0.0, 0.0, 0.0], [0.5705567598342896, 0.012362794950604439, 0.012375605292618275, 0.03007964789867401, 0.32326540350914, 0.05135980620980263, 0.0, 0.0, 0.0], [0.3255925178527832, 0.018225088715553284, 0.44416120648384094, 0.03159342333674431, 0.09710187464952469, 0.018276162445545197, 0.06504977494478226, 0.0, 0.0], [0.7573192119598389, 0.006619773339480162, 0.024497462436556816, 0.10662192106246948, 0.016511093825101852, 0.026536984369158745, 0.016377350315451622, 0.04551612585783005, 0.0], [0.09326817095279694, 0.0063356757164001465, 0.0017902236431837082, 0.005481854546815157, 0.7756001949310303, 0.014800457283854485, 0.030199307948350906, 0.049663979560136795, 0.022860107943415642]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9506229758262634, 0.04937697574496269, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8636487126350403, 0.0646706372499466, 0.07168063521385193, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7691375613212585, 0.029964299872517586, 0.12727312743663788, 0.07362502813339233, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6092661619186401, 0.02264225110411644, 0.1738659292459488, 0.11226367950439453, 0.08196202665567398, 0.0, 0.0, 0.0, 0.0], [0.3979811370372772, 0.007602523546665907, 0.07248664647340775, 0.26127785444259644, 0.15910464525222778, 0.10154716670513153, 0.0, 0.0, 0.0], [0.33622974157333374, 0.060418751090765, 0.2987986207008362, 0.16040712594985962, 0.05644509196281433, 0.039555490016937256, 0.04814521223306656, 0.0, 0.0], [0.535023033618927, 0.02326866425573826, 0.04021435230970383, 0.20968644320964813, 0.04370691999793053, 0.04488035663962364, 0.06491724401712418, 0.03830297663807869, 0.0], [0.48421016335487366, 0.01896764524281025, 0.030104398727416992, 0.09323664754629135, 0.16630412638187408, 0.06683766096830368, 0.05963825434446335, 0.03186190500855446, 0.048839252442121506]]], \"attentionHeadNames\": [\"L9H7\", \"L10H7\", \"L11H10\"], \"tokens\": [\"<|endoftext|>\", \"A\", \" mouse\", \" is\", \" small\", \" but\", \" an\", \" elephant\", \" is\"]}\n",
              "    )\n",
              "    </script></div><div style='max-width: 700px;'><h2>Top 3 Negative Logit Attribution Heads</h2><br/><div id=\"circuits-vis-c31d36c0-74e5\" style=\"margin: 15px 0;\"/>\n",
              "    <script crossorigin type=\"module\">\n",
              "    import { render, AttentionHeads } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
              "    render(\n",
              "      \"circuits-vis-c31d36c0-74e5\",\n",
              "      AttentionHeads,\n",
              "      {\"attention\": [[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9824504256248474, 0.01754951849579811, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9570852518081665, 0.0024241425562649965, 0.04049060493707657, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8858200907707214, 0.009392774663865566, 0.08570577204227448, 0.019081400707364082, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8826439380645752, 0.0028362898156046867, 0.0866345539689064, 0.002867253962904215, 0.025018051266670227, 0.0, 0.0, 0.0, 0.0], [0.8823736310005188, 0.005670490674674511, 0.009323567152023315, 0.008928066119551659, 0.059599556028842926, 0.03410470485687256, 0.0, 0.0, 0.0], [0.8217462301254272, 0.03471330553293228, 0.05128922685980797, 0.0036314306780695915, 0.053689006716012955, 0.012227464467287064, 0.02270328253507614, 0.0, 0.0], [0.5749940872192383, 0.009073768742382526, 0.11085585504770279, 0.04800565913319588, 0.10685914009809494, 0.041319381445646286, 0.02710012160241604, 0.08179203420877457, 0.0], [0.5314610004425049, 0.010446934029459953, 0.006295223254710436, 0.009326216764748096, 0.39668846130371094, 0.013130798004567623, 0.015078941360116005, 0.009814027696847916, 0.0077583217062056065]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9294889569282532, 0.07051104307174683, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.961098849773407, 0.01934536173939705, 0.01955586113035679, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9262334108352661, 0.015146452002227306, 0.017152654007077217, 0.04146750271320343, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9188517928123474, 0.013149654492735863, 0.019416728988289833, 0.0368507020175457, 0.01173100620508194, 0.0, 0.0, 0.0, 0.0], [0.9183754920959473, 0.018714701756834984, 0.009916894137859344, 0.01910332776606083, 0.01682744361460209, 0.017062097787857056, 0.0, 0.0, 0.0], [0.8261955380439758, 0.034665804356336594, 0.09462044388055801, 0.008749963715672493, 0.023674068972468376, 0.00441260589286685, 0.007681544404476881, 0.0, 0.0], [0.7842328548431396, 0.05153097212314606, 0.01611478067934513, 0.08800896257162094, 0.010277516208589077, 0.02909010648727417, 0.009754310362040997, 0.010990451090037823, 0.0], [0.8479006290435791, 0.012382165528833866, 0.00697946734726429, 0.017735179513692856, 0.027184953913092613, 0.01199188269674778, 0.022834056988358498, 0.03204350546002388, 0.020948143675923347]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.997093677520752, 0.002906296867877245, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.978542685508728, 0.0007322311867028475, 0.02072509378194809, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9366994500160217, 0.008979960344731808, 0.043810389935970306, 0.010510188527405262, 0.0, 0.0, 0.0, 0.0, 0.0], [0.931016206741333, 0.000538992986548692, 0.059891778975725174, 0.0012615852756425738, 0.0072913020849227905, 0.0, 0.0, 0.0, 0.0], [0.9744870662689209, 0.0013381453463807702, 0.007696496322751045, 0.0005283160135149956, 0.012848284095525742, 0.0031016452703624964, 0.0, 0.0, 0.0], [0.9173592329025269, 0.010737542062997818, 0.05546996742486954, 0.001540816156193614, 0.0035333833657205105, 0.0005904482677578926, 0.010768645443022251, 0.0, 0.0], [0.8578725457191467, 0.0020349363330751657, 0.06400742381811142, 0.015871569514274597, 0.012477627024054527, 0.00260402774438262, 0.005950705148279667, 0.03918111324310303, 0.0], [0.9584300518035889, 0.0010149084264412522, 0.0075726197101175785, 0.00024245721579063684, 0.014895761385560036, 0.0007754968828521669, 0.0027447787579149008, 0.013773089274764061, 0.0005509259062819183]]], \"attentionHeadNames\": [\"L10H1\", \"L10H3\", \"L9H9\"], \"tokens\": [\"<|endoftext|>\", \"A\", \" mouse\", \" is\", \" small\", \" but\", \" an\", \" elephant\", \" is\"]}\n",
              "    )\n",
              "    </script></div>"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Note\n",
        "\n",
        "It's at this point where I'm wondering how solid of a task this was to predict because it has to introduce new information that isn't in the prompt. I've heard that most stored knowledge/facts exist in the earlier layers of the transformer. My intuition was that there could be an algorithm that combines some form of lookup information with an algorithm that finds the opposite of the adjictive in the prompt(\"incorrect answer\") and something that is similar to the second noun (\"elephant\").\n",
        "\n",
        "Taking a look at these heads, I see that the most impactful head L9H7 attends \"small\" to \"but\", \"an\", and \"is\". This seems to act as the \"opposite\" head, given that small is the wrong adjective but it is attended to \"but\", giving the next circuit information about which opposite to look at.\n",
        "\n",
        "It's curious that the second noun, \"elephant\" is hardly attended from in the postitive difference heads.\n",
        "\n",
        "The negative heads, seem to focus mostly on attending the first adjective \"small\", with the last token \"is\", which makes sense that moving the opposite adjective to the last token would result in a higher likelihood of predicting that \"wrong\" adjective, assuming there was little processing done after moving it.\n",
        "\n",
        "There is also some movement for the first noun to the second noun, which I expected more of in the postive heads because the relationship between the nouns should be helpful in understanding what adjective to characterize the second noun."
      ],
      "metadata": {
        "id": "yNu9md-8eqC2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Activation Patching Approach\n",
        "\n",
        "I will first try changing the prompt by swapping both the noun and the adjective. It is probably better to do this because swapping just the adjective will confuse the model that is also probably doing some work to describe the relevant noun instead of just output the opposite adjective.\n",
        "\n",
        "Therefore I will try these prompts:\n",
        "\n",
        "Correct:   \"A mouse is small but an elephant is\",\n",
        "Incorrect: \"An elephant is big but a mouse is\",\n",
        "\n",
        "I hope this can control only for the circuits that are meant to understand what proper opposite adjective should be since the structure of both prompts are the same."
      ],
      "metadata": {
        "id": "4TaZ0aS0u9D3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "corrupted_prompts = [\n",
        "\"An elephant is big but a mouse is\",\n",
        "\"A kitten is young but a turtle is\",\n",
        "\"An ocean is wet but a desert is\",\n",
        "\"A brick is hard but a pillow is\",\n",
        "\"A sponge is soft but a rock is\",\n",
        "\"A weed is ugly but a rose is\"\n",
        "]\n",
        "# for prompt in corrupted_prompts:\n",
        "#     str_tokens = model.to_str_tokens(prompt)\n",
        "#     print(\"Prompt length:\", len(str_tokens))\n",
        "#     print(\"Prompt as tokens:\", str_tokens)\n",
        "\n",
        "corrupted_tokens = model.to_tokens(corrupted_prompts, prepend_bos=True)\n",
        "corrupted_logits, corrupted_cache = model.run_with_cache(\n",
        "    corrupted_tokens, return_type=\"logits\"\n",
        ")\n",
        "\n",
        "corrupted_average_logit_diff = logits_to_ave_logit_diff(corrupted_logits, answer_tokens)\n",
        "print(\"Corrupted Average Logit Diff\", round(corrupted_average_logit_diff.item(), 2))\n",
        "print(\"Clean Average Logit Diff\", round(original_average_logit_diff.item(), 2))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RIqX5mfjgu4V",
        "outputId": "da87aaa4-c2ad-4bd4-a38b-0498ded0ecf1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Corrupted Average Logit Diff -1.63\n",
            "Clean Average Logit Diff 1.74\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def patch_residual_component(\n",
        "    corrupted_residual_component: Float[torch.Tensor, \"batch pos d_model\"],\n",
        "    hook,\n",
        "    pos,\n",
        "    clean_cache,\n",
        "):\n",
        "    corrupted_residual_component[:, pos, :] = clean_cache[hook.name][:, pos, :]\n",
        "    return corrupted_residual_component\n",
        "\n",
        "\n",
        "def normalize_patched_logit_diff(patched_logit_diff):\n",
        "    # Subtract corrupted logit diff to measure the improvement, divide by the total improvement from clean to corrupted to normalise\n",
        "    # 0 means zero change, negative means actively made worse, 1 means totally recovered clean performance, >1 means actively *improved* on clean performance\n",
        "    return (patched_logit_diff - corrupted_average_logit_diff) / (\n",
        "        original_average_logit_diff - corrupted_average_logit_diff\n",
        "    )\n",
        "\n",
        "\n",
        "patched_residual_stream_diff = torch.zeros(\n",
        "    model.cfg.n_layers, tokens.shape[1], device=device, dtype=torch.float32\n",
        ")\n",
        "for layer in range(model.cfg.n_layers):\n",
        "    for position in range(tokens.shape[1]):\n",
        "        hook_fn = partial(patch_residual_component, pos=position, clean_cache=cache)\n",
        "        patched_logits = model.run_with_hooks(\n",
        "            corrupted_tokens,\n",
        "            fwd_hooks=[(utils.get_act_name(\"resid_pre\", layer), hook_fn)],\n",
        "            return_type=\"logits\",\n",
        "        )\n",
        "        patched_logit_diff = logits_to_ave_logit_diff(patched_logits, answer_tokens)\n",
        "\n",
        "        patched_residual_stream_diff[layer, position] = normalize_patched_logit_diff(\n",
        "            patched_logit_diff\n",
        "        )"
      ],
      "metadata": {
        "id": "mSLsJz-T3Ox-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompt_position_labels = [\n",
        "    f\"{tok}_{i}\" for i, tok in enumerate(model.to_str_tokens(tokens[0]))\n",
        "]\n",
        "imshow(\n",
        "    patched_residual_stream_diff,\n",
        "    x=prompt_position_labels,\n",
        "    title=\"Logit Difference From Patched Residual Stream\",\n",
        "    labels={\"x\": \"Position\", \"y\": \"Layer\"},\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "mnM_O4Or_qnh",
        "outputId": "97138ce6-b921-4c9f-ccb3-c2948d46c08b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"6283c885-7f4e-4a08-a292-d77998deed02\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"6283c885-7f4e-4a08-a292-d77998deed02\")) {                    Plotly.newPlot(                        \"6283c885-7f4e-4a08-a292-d77998deed02\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"\\u003c|endoftext|\\u003e_0\",\"A_1\",\" mouse_2\",\" is_3\",\" small_4\",\" but_5\",\" an_6\",\" elephant_7\",\" is_8\"],\"z\":[[0.0,0.006206355523318052,-0.07651412487030029,0.0,0.9391323328018188,0.0,-0.0006512489053420722,0.1392621248960495,0.0],[0.0,0.005642754957079887,-0.09889412671327591,-0.0022722496651113033,0.9411510229110718,-0.00274258223362267,-0.0014991766074672341,0.12723499536514282,0.000988977961242199],[0.0,0.0041963797993958,-0.0938941016793251,-0.00248783640563488,0.9449524879455566,-0.0072442470118403435,0.0008845067350193858,0.1382954716682434,-0.0026957541704177856],[0.0,0.004080174956470728,-0.06540977954864502,-0.027685021981596947,0.9401445388793945,-0.00584286218509078,0.006234912201762199,0.1350691318511963,-0.009853271767497063],[0.0,0.006485982332378626,-0.05723949894309044,-0.030049549415707588,0.936122715473175,0.002015100559219718,0.005136903841048479,0.17174172401428223,-0.01923048496246338],[0.0,0.005493964999914169,-0.036623887717723846,0.02556481398642063,0.9176777601242065,0.05862339213490486,0.00692196236923337,0.2067686766386032,-0.036024484783411026],[0.0,0.005258162505924702,-0.03966313228011131,0.010650413110852242,0.9165608286857605,0.06386002898216248,0.009220294654369354,0.2485589236021042,-0.006420811638236046],[0.0,0.0011290741385892034,-0.0615423321723938,5.145809700479731e-05,0.7767050266265869,0.17294614017009735,0.026023555546998978,0.22001716494560242,0.023008272051811218],[0.0,0.0007463191286660731,-0.05236441269516945,-0.0091181555762887,0.8165079951286316,0.11064837872982025,0.025930818170309067,0.21341021358966827,0.09788917005062103],[0.0,-7.255733362399042e-05,-0.0456385612487793,-0.010310527868568897,0.8444496393203735,0.030314192175865173,0.003985140006989241,0.16322170197963715,0.23948104679584503],[0.0,-0.00015419757983181626,-0.02274833247065544,-0.011981572955846786,0.008634216152131557,0.004557399079203606,0.003112331498414278,0.04617505520582199,1.1030603647232056],[0.0,0.0007535996264778078,-0.0077230972237885,-0.010864479467272758,-0.030716808512806892,-0.0037821349687874317,3.180789178713894e-07,0.008826371282339096,1.075160264968872]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003ecolor: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Patched Residual Stream\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('6283c885-7f4e-4a08-a292-d77998deed02');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "del patched_residual_stream_diff # delete to free up memory"
      ],
      "metadata": {
        "id": "5wY65BUJPOBW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ok this graph from the patched residual stream helps a lot and shows that the most important token that the circuit is paying attention to is the first adjective \"small\" in this case. Elephant or the second noun is also being looked at closely, which we didn't understand from our initial logit difference experiments that only loooked at the en of the circuits.\n",
        "\n",
        "In addition, the last two layers move all information to the last token and have both actively improved on clean performance."
      ],
      "metadata": {
        "id": "DDz1eBSBF2iX"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we try patching in Layers to see the effect from attention vs MLP. We already know that both play a role from our earlier graphs but attention seems more important especially in layer 9. But let's see.\n",
        "\n",
        "##### Issue: Running into a ton of OOM issue here on this layer patching script. The solution was to include the disable automatic differentiation at the start"
      ],
      "metadata": {
        "id": "LsyxCPxiNFe-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "patched_attn_diff = torch.zeros(\n",
        "    model.cfg.n_layers, tokens.shape[1], device=device, dtype=torch.float32\n",
        ")\n",
        "patched_mlp_diff = torch.zeros(\n",
        "    model.cfg.n_layers, tokens.shape[1], device=device, dtype=torch.float32\n",
        ")\n",
        "for layer in range(model.cfg.n_layers):\n",
        "    for position in range(tokens.shape[1]):\n",
        "        hook_fn = partial(patch_residual_component, pos=position, clean_cache=cache)\n",
        "        patched_attn_logits = model.run_with_hooks(\n",
        "            corrupted_tokens,\n",
        "            fwd_hooks=[(utils.get_act_name(\"attn_out\", layer), hook_fn)],\n",
        "            return_type=\"logits\",\n",
        "        )\n",
        "        patched_attn_logit_diff = logits_to_ave_logit_diff(\n",
        "            patched_attn_logits, answer_tokens\n",
        "        )\n",
        "        patched_mlp_logits = model.run_with_hooks(\n",
        "            corrupted_tokens,\n",
        "            fwd_hooks=[(utils.get_act_name(\"mlp_out\", layer), hook_fn)],\n",
        "            return_type=\"logits\",\n",
        "        )\n",
        "        patched_mlp_logit_diff = logits_to_ave_logit_diff(\n",
        "            patched_mlp_logits, answer_tokens\n",
        "        )\n",
        "\n",
        "        patched_attn_diff[layer, position] = normalize_patched_logit_diff(\n",
        "            patched_attn_logit_diff\n",
        "        )\n",
        "        patched_mlp_diff[layer, position] = normalize_patched_logit_diff(\n",
        "            patched_mlp_logit_diff\n",
        "        )"
      ],
      "metadata": {
        "id": "W2fsmwlHCx8U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "imshow(\n",
        "    patched_attn_diff,\n",
        "    x=prompt_position_labels,\n",
        "    title=\"Logit Difference From Patched Attention Layer\",\n",
        "    labels={\"x\": \"Position\", \"y\": \"Layer\"},\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "lF30qB8RJjX0",
        "outputId": "e5dfb2d2-f7e9-4992-cd66-8847ba631770"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"60a07172-08fe-455e-8a96-e9d589f16b5f\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"60a07172-08fe-455e-8a96-e9d589f16b5f\")) {                    Plotly.newPlot(                        \"60a07172-08fe-455e-8a96-e9d589f16b5f\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"\\u003c|endoftext|\\u003e_0\",\"A_1\",\" mouse_2\",\" is_3\",\" small_4\",\" but_5\",\" an_6\",\" elephant_7\",\" is_8\"],\"z\":[[0.0,0.0010505439713597298,-0.03872193768620491,-0.0022722496651113033,0.04121754691004753,-0.00274258223362267,-0.0015293234027922153,0.016955867409706116,0.000988977961242199],[0.0,0.0005171255907043815,0.01046875398606062,-0.002787537407130003,-0.002682854188606143,-0.0028958963230252266,0.000481041322927922,0.016538018360733986,-0.003334774635732174],[0.0,0.000970458728261292,-0.00535132410004735,-0.02537718415260315,-0.012367756105959415,0.0013339875731617212,0.0022443647030740976,0.004454801324754953,-0.008503839373588562],[0.0,-0.0006621695938520133,-0.0020133687648922205,-0.01889851689338684,0.007719032932072878,-0.008077330887317657,-0.002810156438499689,-0.005080073606222868,-0.015335255302488804],[0.0,0.00013589038280770183,0.0005907785380259156,-0.01956920325756073,-0.030304331332445145,0.03025057539343834,0.003199273021891713,0.007126204669475555,-0.02330034039914608],[0.0,0.00012769100430887192,1.378341949020978e-06,-0.0030319634824991226,-0.004451867658644915,-0.005541111342608929,0.006245196796953678,0.016499247401952744,0.019847240298986435],[0.0,-0.00030139743466861546,0.0008334020967595279,-0.0010799132287502289,-0.0017296071164309978,0.10823225229978561,0.019274592399597168,0.0144593371078372,0.030725715681910515],[0.0,-4.65455450466834e-05,4.6368837502086535e-05,-0.00033136754063889384,-0.0008256268338300288,0.0023080159444361925,0.002406903076916933,0.021001514047384262,0.05719747766852379],[0.0,-7.842412014724687e-05,7.421841132781992e-07,-0.0010534420143812895,-0.009273095987737179,-0.004555278457701206,-0.0002185202029068023,0.012211402878165245,0.1265639215707779],[0.0,-2.2618944058194757e-05,-0.0012902694288641214,-0.002236730884760618,-0.0002881088003050536,0.00811055302619934,0.00995222944766283,0.003726647701114416,0.9950038194656372],[0.0,-6.714998903589731e-07,-0.0003107630764134228,-0.0011888021836057305,-0.001612306572496891,-0.001392160658724606,-0.002604748122394085,-0.005803349427878857,0.04940956458449364],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.007823292165994644]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003ecolor: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Patched Attention Layer\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('60a07172-08fe-455e-8a96-e9d589f16b5f');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "imshow(\n",
        "    patched_mlp_diff,\n",
        "    x=prompt_position_labels,\n",
        "    title=\"Logit Difference From Patched MLP Layer\",\n",
        "    labels={\"x\": \"Position\", \"y\": \"Layer\"},\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "JGqA9XXrJkQ1",
        "outputId": "2d51b77d-1926-4ad0-e42d-e554621e314e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"913bb7fa-0c81-4b8c-8bb2-1127b640f77c\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"913bb7fa-0c81-4b8c-8bb2-1127b640f77c\")) {                    Plotly.newPlot(                        \"913bb7fa-0c81-4b8c-8bb2-1127b640f77c\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"\\u003c|endoftext|\\u003e_0\",\"A_1\",\" mouse_2\",\" is_3\",\" small_4\",\" but_5\",\" an_6\",\" elephant_7\",\" is_8\"],\"z\":[[0.0,0.0005480499821715057,-0.06859368085861206,-0.0013572426978498697,0.8559735417366028,0.0008630894590169191,0.003080028807744384,0.11578100174665451,0.004149021115154028],[0.0,0.0001361024333164096,-0.014474958181381226,0.0027424409054219723,0.19180530309677124,-0.006829472258687019,3.8840968045406044e-05,-0.011012704111635685,0.00013963664241600782],[0.0,-0.0008426617132499814,-0.012874102219939232,0.0028651487082242966,0.07233944535255432,-0.0059237247332930565,0.0028274033684283495,0.024114657193422318,0.005708279553800821],[0.0,-9.075851266970858e-05,0.007182186469435692,-0.016663270071148872,0.09622608125209808,-0.018814755603671074,0.009262917563319206,0.028367266058921814,-0.008827467449009418],[0.0,0.001519533572718501,-0.0011842431267723441,-0.021481528878211975,0.07809731364250183,0.0008776150643825531,0.0015773533377796412,0.015830539166927338,-0.01931983046233654],[0.0,0.00021940376609563828,0.0019926230888813734,0.009009267203509808,0.024953607469797134,0.003469215938821435,-0.0015732889296486974,0.02760695107281208,0.02625183016061783],[0.0,-0.0020235474221408367,-0.00956714153289795,0.0012210342101752758,0.026344532147049904,0.14567558467388153,0.01662587746977806,0.03479698300361633,0.06248762458562851],[0.0,-0.0021439932752400637,-0.006802965421229601,-0.002706745406612754,0.014103795401751995,0.0243988074362278,-0.007982366718351841,0.028493402525782585,0.10373447090387344],[0.0,-0.0008373603923246264,-0.0052611310966312885,0.0011237020371481776,0.0018992491532117128,0.015268494375050068,0.0006516022840514779,0.0032887591514736414,0.2245168834924698],[0.0,-0.0009460373548790812,-0.002662638435140252,-0.00041615322697907686,0.013555992394685745,0.005435897968709469,0.003437620121985674,0.003770012641325593,0.39332956075668335],[0.0,5.728954420192167e-05,-0.001732081058435142,-0.0022704119328409433,0.001972654601559043,0.0011451194295659661,-2.7142732506035827e-05,-0.0005541287828236818,0.1854439228773117],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.22812417149543762]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003ecolor: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Patched MLP Layer\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('913bb7fa-0c81-4b8c-8bb2-1127b640f77c');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Wow! This is awesome, it seems that attention only plays a big part at layer 9 where all the information is moved to the last token.\n",
        "\n",
        "However, the MLPs are doing a lot of the work, unlike the IOI task. This makes sense when we consider that this task involves more processing to figure out what the opposite adjective is since it isn't in the target.\n",
        "\n",
        "What's cool is that we can see in the earlier layers, even up to layer 0, the MLP starts processing the first adjective. The note in Neel's notebook states a hypothesis that MLP layer 0 in gpt-2 small has high activity across the board and can act as compying over the token embeddings, so that when later layers access the input tokens they do so at MLP0. This is still interesting because both the first adjective and the second noun are highly active here but it seems like the first noun is not being used/pulled much from the MLP0.\n",
        "\n",
        "However we still see a lot of activity at layer 1 on \"small\".\n",
        "\n",
        "Could it be that the general algorithm is simply to do processing and compare the first adjective with the last noun and then figure out what the correct adjective to output is in the context of the first adjective?\n",
        "\n",
        "We also see that once all the information is moved to the last token, a lot of processing happens in layer 8,9,10 and 11. Therefore, I'm not sure what kind of processing happens when because know there has to be an MLP that takes in \"small\" and outputs the opposite."
      ],
      "metadata": {
        "id": "H0nrZG70kIRr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Individual Heads\n"
      ],
      "metadata": {
        "id": "ziyt2kB3nDt2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def patch_head_vector(\n",
        "    corrupted_head_vector: Float[torch.Tensor, \"batch pos head_index d_head\"],\n",
        "    hook,\n",
        "    head_index,\n",
        "    clean_cache,\n",
        "):\n",
        "    corrupted_head_vector[:, :, head_index, :] = clean_cache[hook.name][\n",
        "        :, :, head_index, :\n",
        "    ]\n",
        "    return corrupted_head_vector\n",
        "\n",
        "\n",
        "patched_head_z_diff = torch.zeros(\n",
        "    model.cfg.n_layers, model.cfg.n_heads, device=device, dtype=torch.float32\n",
        ")\n",
        "for layer in range(model.cfg.n_layers):\n",
        "    for head_index in range(model.cfg.n_heads):\n",
        "        hook_fn = partial(patch_head_vector, head_index=head_index, clean_cache=cache)\n",
        "        patched_logits = model.run_with_hooks(\n",
        "            corrupted_tokens,\n",
        "            fwd_hooks=[(utils.get_act_name(\"z\", layer, \"attn\"), hook_fn)],\n",
        "            return_type=\"logits\",\n",
        "        )\n",
        "        patched_logit_diff = logits_to_ave_logit_diff(patched_logits, answer_tokens)\n",
        "\n",
        "        patched_head_z_diff[layer, head_index] = normalize_patched_logit_diff(\n",
        "            patched_logit_diff\n",
        "        )"
      ],
      "metadata": {
        "id": "_jgdXs6mJu8n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "imshow(\n",
        "    patched_head_z_diff,\n",
        "    title=\"Logit Difference From Patched Head Output Attention\",\n",
        "    labels={\"x\": \"Head\", \"y\": \"Layer\"},\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "usueqbioKBSt",
        "outputId": "52dadb29-4188-465d-e113-e7f5b14f5da5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"b3b3d391-a07b-47cc-9701-49228601b56f\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"b3b3d391-a07b-47cc-9701-49228601b56f\")) {                    Plotly.newPlot(                        \"b3b3d391-a07b-47cc-9701-49228601b56f\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-0.003276954870671034,0.011968850158154964,-0.0024690343998372555,0.010896534658968449,0.02463439851999283,0.004504138603806496,-0.006314396392554045,-0.0034733156207948923,-0.015148083679378033,0.0013615190982818604,-0.015516772866249084,0.0010516748297959566],[-0.005969174671918154,-0.004053774289786816,0.002501831855624914,0.007648030761629343,-0.0024495257530361414,-0.0009063481702469289,0.001066094497218728,-0.0008278533350676298,0.0008194065885618329,0.0006284178816713393,-0.0006168257095851004,0.021516306325793266],[0.001977355219423771,-0.019867632538080215,-0.018490387126803398,-0.0006193349836394191,0.011022918857634068,-0.0007332778768613935,-0.00028708387981168926,0.014859833754599094,-0.0031067472882568836,-0.02006017602980137,-0.003509717993438244,0.0005305556114763021],[-0.0012852154904976487,-0.0019150824518874288,0.0038128825835883617,-0.0035368960816413164,-0.006892239674925804,-0.009536182507872581,-0.0075165932066738605,-0.039258427917957306,-0.0027596172876656055,-0.004137322772294283,0.01941341534256935,0.012854911386966705],[0.006432014983147383,-0.011862046085298061,-0.0005677001900039613,-0.009263553656637669,0.012522094883024693,-0.0025314840022474527,-0.013359843753278255,0.0537300668656826,-0.016537098214030266,0.0008152715745382011,0.0007700690184719861,-0.036513056606054306],[0.020648233592510223,-0.01483859308063984,0.007865596562623978,-0.007552606984972954,-0.01654876209795475,-0.003732797224074602,-0.0028303368017077446,0.0037496909499168396,0.025589412078261375,0.000553139194380492,-0.0036214697174727917,0.020459577441215515],[0.014516307972371578,0.09428074210882187,0.0035050527658313513,0.03303228318691254,-0.011870704591274261,-0.01131417229771614,0.014160907827317715,0.0406869575381279,-0.0001464930101064965,-0.01312644500285387,-0.004529903177171946,0.007631950080394745],[0.002730212640017271,0.008291362784802914,-0.013695170171558857,0.0009246907429769635,-0.0009956223657354712,0.02162834070622921,-0.007209505885839462,0.03835660591721535,-0.003117986023426056,0.04912640154361725,-0.01902836374938488,0.0061399126425385475],[-0.0026680459268391132,-0.048771779984235764,-0.011770756915211678,-0.041067201644182205,0.08109700679779053,0.04229746013879776,-0.014189994893968105,0.0030708396807312965,0.024279316887259483,-0.010509433224797249,0.031487975269556046,0.02164304256439209],[-0.001896492438390851,-0.02792789414525032,-0.012933335267007351,-0.004316684324294329,-0.001957280794158578,0.026262007653713226,-0.06146302446722984,1.0837187767028809,-0.004418645985424519,-0.04550956189632416,0.052880510687828064,-0.0008837291970849037],[-0.016890378668904305,-0.18157389760017395,-0.006509166676551104,-0.04653384909033775,-0.017392555251717567,0.007043786812573671,-0.00206768955104053,0.3293505609035492,0.0021621237974613905,0.00773369986563921,-0.01797000877559185,-0.02233058772981167],[0.006265730131417513,-0.007989965379238129,-0.014477149583399296,-0.010329541750252247,0.00427826726809144,-0.004690002650022507,-0.009365197271108627,-0.008876628242433071,-0.007687436882406473,-0.01776120811700821,0.11343598365783691,-0.03297686576843262]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003ecolor: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Patched Head Output Attention\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('b3b3d391-a07b-47cc-9701-49228601b56f');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This again shows how layer 10 is doing work at head 1 and head 7, but you can't see this at the layer level because their work cancels out. I'm curious to explore layer 9 and layer 10 attention heads the most."
      ],
      "metadata": {
        "id": "JQEnbpo9nwyr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Decomposing Heads: Values"
      ],
      "metadata": {
        "id": "hrwCFlMPiFpu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "patched_head_v_diff = torch.zeros(\n",
        "    model.cfg.n_layers, model.cfg.n_heads, device=device, dtype=torch.float32\n",
        ")\n",
        "for layer in range(model.cfg.n_layers):\n",
        "    for head_index in range(model.cfg.n_heads):\n",
        "        hook_fn = partial(patch_head_vector, head_index=head_index, clean_cache=cache)\n",
        "        patched_logits = model.run_with_hooks(\n",
        "            corrupted_tokens,\n",
        "            fwd_hooks=[(utils.get_act_name(\"v\", layer, \"attn\"), hook_fn)],\n",
        "            return_type=\"logits\",\n",
        "        )\n",
        "        patched_logit_diff = logits_to_ave_logit_diff(patched_logits, answer_tokens)\n",
        "\n",
        "        patched_head_v_diff[layer, head_index] = normalize_patched_logit_diff(\n",
        "            patched_logit_diff\n",
        "        )"
      ],
      "metadata": {
        "id": "DBN8TSsuiIsc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "imshow(\n",
        "    patched_head_v_diff,\n",
        "    title=\"Logit Difference From Patched Head Value\",\n",
        "    labels={\"x\": \"Head\", \"y\": \"Layer\"},\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "uPfH8FTXiMh-",
        "outputId": "83513a7c-a059-4420-ac32-41918ad608a4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"bc73d160-926f-49a2-8e0b-bac3939dcdf7\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"bc73d160-926f-49a2-8e0b-bac3939dcdf7\")) {                    Plotly.newPlot(                        \"bc73d160-926f-49a2-8e0b-bac3939dcdf7\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-0.0033113781828433275,0.012060986831784248,0.0008645031484775245,0.010729402303695679,0.023121897131204605,0.004399278666824102,-0.013467743992805481,-0.011770756915211678,-0.017555270344018936,0.0032133038621395826,-0.01360624935477972,0.002268680138513446],[-0.008639553561806679,-0.004421720746904612,0.0015616613673046231,-0.0008094401564449072,-0.0020340438932180405,-0.0012104315683245659,-0.007783709093928337,0.002858504420146346,0.0016261606942862272,0.0012236848706379533,-0.0008393041789531708,0.020728247240185738],[0.007223536726087332,-0.017885470762848854,-0.022150732576847076,0.003957467153668404,0.00859028659760952,-0.000766216719057411,-4.944359898217954e-05,0.01531765516847372,-0.0016403328627347946,-0.02366277389228344,-0.003998640459030867,0.00043990311678498983],[-0.0020593488588929176,-0.0028928215615451336,0.007919175550341606,-0.0062743183225393295,-0.008303378708660603,-0.016847295686602592,-0.0056588356383144855,-0.0396612249314785,0.0012681098887696862,-0.0019798290450125933,0.020602747797966003,0.008731584064662457],[0.002060974482446909,-0.005189916584640741,-0.0024065496399998665,-0.015918223187327385,0.012775215320289135,0.0006792398053221405,-0.007599364500492811,0.05653195083141327,-0.010987753048539162,0.009941344149410725,0.00301538803614676,-0.036514610052108765],[0.007631348911672831,-0.01926642842590809,0.0012342521222308278,-0.002438569674268365,-0.004988502245396376,-0.0050097075290977955,0.007892067544162273,0.006326660048216581,0.031534306704998016,0.004077948164194822,0.008075210265815258,0.012847984209656715],[0.01583651266992092,0.07786157727241516,0.0004433313151821494,0.03199096396565437,-0.01199100911617279,-0.009528477676212788,0.02822112664580345,0.04224776849150658,0.004297599662095308,-0.028549738228321075,-0.005256253760308027,0.001444042893126607],[0.004873075056821108,-0.004189488012343645,-0.019215606153011322,0.005228687077760696,0.002119501121342182,0.013561082072556019,-0.0013547333655878901,0.006687255576252937,-0.009675642475485802,0.05189482122659683,-0.029603250324726105,0.010614081285893917],[0.00019113007874693722,-0.05626847594976425,-0.012517500668764114,-0.045907940715551376,0.09821248054504395,0.04585040360689163,-0.03932780772447586,0.0005807060515508056,0.02133362367749214,-0.002150213345885277,0.04083097726106644,0.020184190943837166],[-0.006629188545048237,-0.03029736876487732,-0.030876273289322853,0.00012398009130265564,-0.0027012673672288656,0.019651656970381737,-0.07288795709609985,1.1177363395690918,-0.002539718756452203,-0.03701883554458618,0.04768352583050728,-0.0005549769848585129],[-0.023294012993574142,-0.16916868090629578,-0.00585364131256938,-0.027517182752490044,-0.02073786035180092,0.010228569619357586,-0.0021216922905296087,0.36011263728141785,0.0025426167994737625,0.007633787579834461,-0.019609104841947556,-0.024699075147509575],[0.0024477585684508085,-0.015226366929709911,-0.014063787646591663,-0.023901119828224182,0.0031162190716713667,0.0007741333683952689,-0.00992402620613575,-0.018597401678562164,-0.0033238185569643974,-0.019750438630580902,0.11346210539340973,-0.04057735949754715]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003ecolor: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Patched Head Value\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('bc73d160-926f-49a2-8e0b-bac3939dcdf7');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "head_labels = [\n",
        "    f\"L{l}H{h}\" for l in range(model.cfg.n_layers) for h in range(model.cfg.n_heads)\n",
        "]\n",
        "scatter(\n",
        "    x=utils.to_numpy(patched_head_v_diff.flatten()),\n",
        "    y=utils.to_numpy(patched_head_z_diff.flatten()),\n",
        "    xaxis=\"Value Patch\",\n",
        "    yaxis=\"Output Patch\",\n",
        "    caxis=\"Layer\",\n",
        "    hover_name=head_labels,\n",
        "    color=einops.repeat(\n",
        "        np.arange(model.cfg.n_layers), \"layer -> (layer head)\", head=model.cfg.n_heads\n",
        "    ),\n",
        "    range_x=(-2, 2),\n",
        "    range_y=(-2, 2),\n",
        "    title=\"Scatter plot of output patching vs value patching\",\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "jH2JSUXooPNF",
        "outputId": "7dad4298-7ac5-4ade-fdaa-9effccd4f19a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"b457a59f-9596-4d18-beb2-42c9e6d078ee\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"b457a59f-9596-4d18-beb2-42c9e6d078ee\")) {                    Plotly.newPlot(                        \"b457a59f-9596-4d18-beb2-42c9e6d078ee\",                        [{\"hovertemplate\":\"\\u003cb\\u003e%{hovertext}\\u003c\\u002fb\\u003e\\u003cbr\\u003e\\u003cbr\\u003eValue Patch=%{x}\\u003cbr\\u003eOutput Patch=%{y}\\u003cbr\\u003eLayer=%{marker.color}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"hovertext\":[\"L0H0\",\"L0H1\",\"L0H2\",\"L0H3\",\"L0H4\",\"L0H5\",\"L0H6\",\"L0H7\",\"L0H8\",\"L0H9\",\"L0H10\",\"L0H11\",\"L1H0\",\"L1H1\",\"L1H2\",\"L1H3\",\"L1H4\",\"L1H5\",\"L1H6\",\"L1H7\",\"L1H8\",\"L1H9\",\"L1H10\",\"L1H11\",\"L2H0\",\"L2H1\",\"L2H2\",\"L2H3\",\"L2H4\",\"L2H5\",\"L2H6\",\"L2H7\",\"L2H8\",\"L2H9\",\"L2H10\",\"L2H11\",\"L3H0\",\"L3H1\",\"L3H2\",\"L3H3\",\"L3H4\",\"L3H5\",\"L3H6\",\"L3H7\",\"L3H8\",\"L3H9\",\"L3H10\",\"L3H11\",\"L4H0\",\"L4H1\",\"L4H2\",\"L4H3\",\"L4H4\",\"L4H5\",\"L4H6\",\"L4H7\",\"L4H8\",\"L4H9\",\"L4H10\",\"L4H11\",\"L5H0\",\"L5H1\",\"L5H2\",\"L5H3\",\"L5H4\",\"L5H5\",\"L5H6\",\"L5H7\",\"L5H8\",\"L5H9\",\"L5H10\",\"L5H11\",\"L6H0\",\"L6H1\",\"L6H2\",\"L6H3\",\"L6H4\",\"L6H5\",\"L6H6\",\"L6H7\",\"L6H8\",\"L6H9\",\"L6H10\",\"L6H11\",\"L7H0\",\"L7H1\",\"L7H2\",\"L7H3\",\"L7H4\",\"L7H5\",\"L7H6\",\"L7H7\",\"L7H8\",\"L7H9\",\"L7H10\",\"L7H11\",\"L8H0\",\"L8H1\",\"L8H2\",\"L8H3\",\"L8H4\",\"L8H5\",\"L8H6\",\"L8H7\",\"L8H8\",\"L8H9\",\"L8H10\",\"L8H11\",\"L9H0\",\"L9H1\",\"L9H2\",\"L9H3\",\"L9H4\",\"L9H5\",\"L9H6\",\"L9H7\",\"L9H8\",\"L9H9\",\"L9H10\",\"L9H11\",\"L10H0\",\"L10H1\",\"L10H2\",\"L10H3\",\"L10H4\",\"L10H5\",\"L10H6\",\"L10H7\",\"L10H8\",\"L10H9\",\"L10H10\",\"L10H11\",\"L11H0\",\"L11H1\",\"L11H2\",\"L11H3\",\"L11H4\",\"L11H5\",\"L11H6\",\"L11H7\",\"L11H8\",\"L11H9\",\"L11H10\",\"L11H11\"],\"legendgroup\":\"\",\"marker\":{\"color\":[0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,1,1,1,1,1,2,2,2,2,2,2,2,2,2,2,2,2,3,3,3,3,3,3,3,3,3,3,3,3,4,4,4,4,4,4,4,4,4,4,4,4,5,5,5,5,5,5,5,5,5,5,5,5,6,6,6,6,6,6,6,6,6,6,6,6,7,7,7,7,7,7,7,7,7,7,7,7,8,8,8,8,8,8,8,8,8,8,8,8,9,9,9,9,9,9,9,9,9,9,9,9,10,10,10,10,10,10,10,10,10,10,10,10,11,11,11,11,11,11,11,11,11,11,11,11],\"coloraxis\":\"coloraxis\",\"symbol\":\"circle\"},\"mode\":\"markers\",\"name\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"x\":[-0.0033113781828433275,0.012060986831784248,0.0008645031484775245,0.010729402303695679,0.023121897131204605,0.004399278666824102,-0.013467743992805481,-0.011770756915211678,-0.017555270344018936,0.0032133038621395826,-0.01360624935477972,0.002268680138513446,-0.008639553561806679,-0.004421720746904612,0.0015616613673046231,-0.0008094401564449072,-0.0020340438932180405,-0.0012104315683245659,-0.007783709093928337,0.002858504420146346,0.0016261606942862272,0.0012236848706379533,-0.0008393041789531708,0.020728247240185738,0.007223536726087332,-0.017885470762848854,-0.022150732576847076,0.003957467153668404,0.00859028659760952,-0.000766216719057411,-4.944359898217954e-05,0.01531765516847372,-0.0016403328627347946,-0.02366277389228344,-0.003998640459030867,0.00043990311678498983,-0.0020593488588929176,-0.0028928215615451336,0.007919175550341606,-0.0062743183225393295,-0.008303378708660603,-0.016847295686602592,-0.0056588356383144855,-0.0396612249314785,0.0012681098887696862,-0.0019798290450125933,0.020602747797966003,0.008731584064662457,0.002060974482446909,-0.005189916584640741,-0.0024065496399998665,-0.015918223187327385,0.012775215320289135,0.0006792398053221405,-0.007599364500492811,0.05653195083141327,-0.010987753048539162,0.009941344149410725,0.00301538803614676,-0.036514610052108765,0.007631348911672831,-0.01926642842590809,0.0012342521222308278,-0.002438569674268365,-0.004988502245396376,-0.0050097075290977955,0.007892067544162273,0.006326660048216581,0.031534306704998016,0.004077948164194822,0.008075210265815258,0.012847984209656715,0.01583651266992092,0.07786157727241516,0.0004433313151821494,0.03199096396565437,-0.01199100911617279,-0.009528477676212788,0.02822112664580345,0.04224776849150658,0.004297599662095308,-0.028549738228321075,-0.005256253760308027,0.001444042893126607,0.004873075056821108,-0.004189488012343645,-0.019215606153011322,0.005228687077760696,0.002119501121342182,0.013561082072556019,-0.0013547333655878901,0.006687255576252937,-0.009675642475485802,0.05189482122659683,-0.029603250324726105,0.010614081285893917,0.00019113007874693722,-0.05626847594976425,-0.012517500668764114,-0.045907940715551376,0.09821248054504395,0.04585040360689163,-0.03932780772447586,0.0005807060515508056,0.02133362367749214,-0.002150213345885277,0.04083097726106644,0.020184190943837166,-0.006629188545048237,-0.03029736876487732,-0.030876273289322853,0.00012398009130265564,-0.0027012673672288656,0.019651656970381737,-0.07288795709609985,1.1177363395690918,-0.002539718756452203,-0.03701883554458618,0.04768352583050728,-0.0005549769848585129,-0.023294012993574142,-0.16916868090629578,-0.00585364131256938,-0.027517182752490044,-0.02073786035180092,0.010228569619357586,-0.0021216922905296087,0.36011263728141785,0.0025426167994737625,0.007633787579834461,-0.019609104841947556,-0.024699075147509575,0.0024477585684508085,-0.015226366929709911,-0.014063787646591663,-0.023901119828224182,0.0031162190716713667,0.0007741333683952689,-0.00992402620613575,-0.018597401678562164,-0.0033238185569643974,-0.019750438630580902,0.11346210539340973,-0.04057735949754715],\"xaxis\":\"x\",\"y\":[-0.003276954870671034,0.011968850158154964,-0.0024690343998372555,0.010896534658968449,0.02463439851999283,0.004504138603806496,-0.006314396392554045,-0.0034733156207948923,-0.015148083679378033,0.0013615190982818604,-0.015516772866249084,0.0010516748297959566,-0.005969174671918154,-0.004053774289786816,0.002501831855624914,0.007648030761629343,-0.0024495257530361414,-0.0009063481702469289,0.001066094497218728,-0.0008278533350676298,0.0008194065885618329,0.0006284178816713393,-0.0006168257095851004,0.021516306325793266,0.001977355219423771,-0.019867632538080215,-0.018490387126803398,-0.0006193349836394191,0.011022918857634068,-0.0007332778768613935,-0.00028708387981168926,0.014859833754599094,-0.0031067472882568836,-0.02006017602980137,-0.003509717993438244,0.0005305556114763021,-0.0012852154904976487,-0.0019150824518874288,0.0038128825835883617,-0.0035368960816413164,-0.006892239674925804,-0.009536182507872581,-0.0075165932066738605,-0.039258427917957306,-0.0027596172876656055,-0.004137322772294283,0.01941341534256935,0.012854911386966705,0.006432014983147383,-0.011862046085298061,-0.0005677001900039613,-0.009263553656637669,0.012522094883024693,-0.0025314840022474527,-0.013359843753278255,0.0537300668656826,-0.016537098214030266,0.0008152715745382011,0.0007700690184719861,-0.036513056606054306,0.020648233592510223,-0.01483859308063984,0.007865596562623978,-0.007552606984972954,-0.01654876209795475,-0.003732797224074602,-0.0028303368017077446,0.0037496909499168396,0.025589412078261375,0.000553139194380492,-0.0036214697174727917,0.020459577441215515,0.014516307972371578,0.09428074210882187,0.0035050527658313513,0.03303228318691254,-0.011870704591274261,-0.01131417229771614,0.014160907827317715,0.0406869575381279,-0.0001464930101064965,-0.01312644500285387,-0.004529903177171946,0.007631950080394745,0.002730212640017271,0.008291362784802914,-0.013695170171558857,0.0009246907429769635,-0.0009956223657354712,0.02162834070622921,-0.007209505885839462,0.03835660591721535,-0.003117986023426056,0.04912640154361725,-0.01902836374938488,0.0061399126425385475,-0.0026680459268391132,-0.048771779984235764,-0.011770756915211678,-0.041067201644182205,0.08109700679779053,0.04229746013879776,-0.014189994893968105,0.0030708396807312965,0.024279316887259483,-0.010509433224797249,0.031487975269556046,0.02164304256439209,-0.001896492438390851,-0.02792789414525032,-0.012933335267007351,-0.004316684324294329,-0.001957280794158578,0.026262007653713226,-0.06146302446722984,1.0837187767028809,-0.004418645985424519,-0.04550956189632416,0.052880510687828064,-0.0008837291970849037,-0.016890378668904305,-0.18157389760017395,-0.006509166676551104,-0.04653384909033775,-0.017392555251717567,0.007043786812573671,-0.00206768955104053,0.3293505609035492,0.0021621237974613905,0.00773369986563921,-0.01797000877559185,-0.02233058772981167,0.006265730131417513,-0.007989965379238129,-0.014477149583399296,-0.010329541750252247,0.00427826726809144,-0.004690002650022507,-0.009365197271108627,-0.008876628242433071,-0.007687436882406473,-0.01776120811700821,0.11343598365783691,-0.03297686576843262],\"yaxis\":\"y\",\"type\":\"scatter\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Value Patch\"},\"range\":[-2,2]},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Output Patch\"},\"range\":[-2,2]},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Layer\"}},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]},\"legend\":{\"tracegroupgap\":0},\"title\":{\"text\":\"Scatter plot of output patching vs value patching\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('b457a59f-9596-4d18-beb2-42c9e6d078ee');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Just as a recap: value patching compared to the output patching helps us figure out when figuring out what to move is especially important (the OV circuit).\n",
        "\n",
        "From this scatter plot we can see an interesting linear line, showing that all the heads seem to matter when considering the OV circuit, however L9H7 followed by L10H7 truly stand out as having significant contributions. I almost missed L9H7 because it was out of the graph's range!"
      ],
      "metadata": {
        "id": "2Chc-tY5ifhR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Decomposing Heads: Attention/ QK"
      ],
      "metadata": {
        "id": "G78B3-XJmHqf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def patch_head_pattern(\n",
        "    corrupted_head_pattern: Float[torch.Tensor, \"batch head_index query_pos d_head\"],\n",
        "    hook,\n",
        "    head_index,\n",
        "    clean_cache,\n",
        "):\n",
        "    corrupted_head_pattern[:, head_index, :, :] = clean_cache[hook.name][\n",
        "        :, head_index, :, :\n",
        "    ]\n",
        "    return corrupted_head_pattern\n",
        "\n",
        "\n",
        "patched_head_attn_diff = torch.zeros(\n",
        "    model.cfg.n_layers, model.cfg.n_heads, device=device, dtype=torch.float32\n",
        ")\n",
        "for layer in range(model.cfg.n_layers):\n",
        "    for head_index in range(model.cfg.n_heads):\n",
        "        hook_fn = partial(patch_head_pattern, head_index=head_index, clean_cache=cache)\n",
        "        patched_logits = model.run_with_hooks(\n",
        "            corrupted_tokens,\n",
        "            fwd_hooks=[(utils.get_act_name(\"attn\", layer, \"attn\"), hook_fn)],\n",
        "            return_type=\"logits\",\n",
        "        )\n",
        "        patched_logit_diff = logits_to_ave_logit_diff(patched_logits, answer_tokens)\n",
        "\n",
        "        patched_head_attn_diff[layer, head_index] = normalize_patched_logit_diff(\n",
        "            patched_logit_diff\n",
        "        )"
      ],
      "metadata": {
        "id": "JlpNWt83qi1q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "imshow(\n",
        "    patched_head_attn_diff,\n",
        "    title=\"Logit Difference From Patched Head Pattern\",\n",
        "    labels={\"x\": \"Head\", \"y\": \"Layer\"},\n",
        ")\n",
        "head_labels = [\n",
        "    f\"L{l}H{h}\" for l in range(model.cfg.n_layers) for h in range(model.cfg.n_heads)\n",
        "]\n",
        "scatter(\n",
        "    x=utils.to_numpy(patched_head_attn_diff.flatten()),\n",
        "    y=utils.to_numpy(patched_head_z_diff.flatten()),\n",
        "    hover_name=head_labels,\n",
        "    xaxis=\"Attention Patch\",\n",
        "    yaxis=\"Output Patch\",\n",
        "    title=\"Scatter plot of output patching vs attention patching\",\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "42SkIwifmFjb",
        "outputId": "d1e21d4d-bf63-472e-e8fc-00f14341fb3b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"71103ab8-496b-45dd-bbed-f43f0164f625\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"71103ab8-496b-45dd-bbed-f43f0164f625\")) {                    Plotly.newPlot(                        \"71103ab8-496b-45dd-bbed-f43f0164f625\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[0.0009426444885320961,9.514093835605308e-05,-0.0015488675562664866,0.0003437019186094403,-0.0009660056093707681,0.0003658614296000451,-0.005153585225343704,0.013546733185648918,0.0001831074187066406,-0.00088698067702353,-0.0009141234331764281,-0.001115290680900216],[-0.0007468492840416729,-0.0004383480700198561,0.0013769635697826743,0.008129142224788666,0.0012334039201959968,0.002106636529788375,0.0010569054866209626,-0.007528538815677166,-0.00024545087944716215,-0.0010706889443099499,0.0001961840025614947,0.0015130306128412485],[-0.0024322434328496456,-0.010977503843605518,0.004166409373283386,-0.004214616026729345,0.0028274741489440203,0.001366643700748682,-0.00012284914555493742,-0.000135961061459966,-0.00027496155234985054,0.002704836893826723,0.0001910947321448475,-4.269325654604472e-05],[0.00045799827785231173,0.0006786036537960172,-0.004250983241945505,0.002127559157088399,-0.0014173949602991343,0.004615148063749075,-0.0030375474598258734,0.0019918454345315695,-0.0030287120025604963,-0.0022658174857497215,-0.0007007631938904524,-0.002680026926100254],[0.004757188260555267,-0.0035628725308924913,0.00105955614708364,0.0065324571914970875,-0.004518664442002773,-0.003672362305223942,-0.009855533950030804,-0.002404535189270973,-0.00109083391726017,-0.0011459322413429618,-0.0021794766653329134,4.7004991756693926e-06],[0.012930402532219887,0.00211087753996253,0.0009747351286932826,-0.006272056605666876,-0.00913893710821867,0.0010564107215031981,-0.004276924300938845,0.0023160031996667385,-0.007413606625050306,-0.0007783391047269106,-0.01268382091075182,0.01114785298705101],[0.002466419246047735,-0.003144139191135764,0.0018313216278329492,0.0016449980903416872,-0.0012056957930326462,-0.0018440447747707367,-0.006244489923119545,0.01494080200791359,0.00043248129077255726,0.005018472205847502,-5.4285465012071654e-05,0.003481232328340411],[-0.003096957691013813,0.009485360234975815,-0.0010450305417180061,0.006258626468479633,-0.0016059096669778228,0.006989925168454647,-0.008994387462735176,0.033319152891635895,0.0018545060884207487,0.005744399037212133,0.004023062065243721,0.0010366191854700446],[-0.004199489951133728,-0.0035497606731951237,0.0020601970609277487,0.0013847388327121735,0.007245024666190147,-0.007857397198677063,0.02110343985259533,-3.569552063709125e-06,0.009662847965955734,-0.008406789973378181,-0.00231211562640965,-0.002066664630547166],[-0.007131046149879694,-0.004756905138492584,-0.008538192138075829,-0.004811614751815796,0.001973750302568078,0.007661990821361542,-0.004735381808131933,0.02941272035241127,-0.0013426817022264004,0.004126790910959244,-0.0012066146591678262,0.00021555146668106318],[0.003473527729511261,0.0016934521263465285,-0.0004835152649320662,0.000657610478810966,-0.00298594799824059,-0.0025072747375816107,-0.00046736394870094955,0.00863566529005766,-0.000456902664154768,2.1912102965870872e-05,-0.00681088212877512,0.004050522577017546],[0.004485442768782377,-0.00467081181704998,-0.0041924212127923965,0.0073399534448981285,0.0016161588719114661,-0.004227586556226015,-0.0033738629426807165,0.014578086324036121,-0.00203616451472044,-0.0029104927089065313,0.03582370653748512,0.001279242685995996]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003ecolor: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Patched Head Pattern\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('71103ab8-496b-45dd-bbed-f43f0164f625');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.24.1.min.js\"></script>                <div id=\"fceada07-2286-4de9-8861-70858b938427\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"fceada07-2286-4de9-8861-70858b938427\")) {                    Plotly.newPlot(                        \"fceada07-2286-4de9-8861-70858b938427\",                        [{\"hovertemplate\":\"\\u003cb\\u003e%{hovertext}\\u003c\\u002fb\\u003e\\u003cbr\\u003e\\u003cbr\\u003eAttention Patch=%{x}\\u003cbr\\u003eOutput Patch=%{y}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"hovertext\":[\"L0H0\",\"L0H1\",\"L0H2\",\"L0H3\",\"L0H4\",\"L0H5\",\"L0H6\",\"L0H7\",\"L0H8\",\"L0H9\",\"L0H10\",\"L0H11\",\"L1H0\",\"L1H1\",\"L1H2\",\"L1H3\",\"L1H4\",\"L1H5\",\"L1H6\",\"L1H7\",\"L1H8\",\"L1H9\",\"L1H10\",\"L1H11\",\"L2H0\",\"L2H1\",\"L2H2\",\"L2H3\",\"L2H4\",\"L2H5\",\"L2H6\",\"L2H7\",\"L2H8\",\"L2H9\",\"L2H10\",\"L2H11\",\"L3H0\",\"L3H1\",\"L3H2\",\"L3H3\",\"L3H4\",\"L3H5\",\"L3H6\",\"L3H7\",\"L3H8\",\"L3H9\",\"L3H10\",\"L3H11\",\"L4H0\",\"L4H1\",\"L4H2\",\"L4H3\",\"L4H4\",\"L4H5\",\"L4H6\",\"L4H7\",\"L4H8\",\"L4H9\",\"L4H10\",\"L4H11\",\"L5H0\",\"L5H1\",\"L5H2\",\"L5H3\",\"L5H4\",\"L5H5\",\"L5H6\",\"L5H7\",\"L5H8\",\"L5H9\",\"L5H10\",\"L5H11\",\"L6H0\",\"L6H1\",\"L6H2\",\"L6H3\",\"L6H4\",\"L6H5\",\"L6H6\",\"L6H7\",\"L6H8\",\"L6H9\",\"L6H10\",\"L6H11\",\"L7H0\",\"L7H1\",\"L7H2\",\"L7H3\",\"L7H4\",\"L7H5\",\"L7H6\",\"L7H7\",\"L7H8\",\"L7H9\",\"L7H10\",\"L7H11\",\"L8H0\",\"L8H1\",\"L8H2\",\"L8H3\",\"L8H4\",\"L8H5\",\"L8H6\",\"L8H7\",\"L8H8\",\"L8H9\",\"L8H10\",\"L8H11\",\"L9H0\",\"L9H1\",\"L9H2\",\"L9H3\",\"L9H4\",\"L9H5\",\"L9H6\",\"L9H7\",\"L9H8\",\"L9H9\",\"L9H10\",\"L9H11\",\"L10H0\",\"L10H1\",\"L10H2\",\"L10H3\",\"L10H4\",\"L10H5\",\"L10H6\",\"L10H7\",\"L10H8\",\"L10H9\",\"L10H10\",\"L10H11\",\"L11H0\",\"L11H1\",\"L11H2\",\"L11H3\",\"L11H4\",\"L11H5\",\"L11H6\",\"L11H7\",\"L11H8\",\"L11H9\",\"L11H10\",\"L11H11\"],\"legendgroup\":\"\",\"marker\":{\"color\":\"#636efa\",\"symbol\":\"circle\"},\"mode\":\"markers\",\"name\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"x\":[0.0009426444885320961,9.514093835605308e-05,-0.0015488675562664866,0.0003437019186094403,-0.0009660056093707681,0.0003658614296000451,-0.005153585225343704,0.013546733185648918,0.0001831074187066406,-0.00088698067702353,-0.0009141234331764281,-0.001115290680900216,-0.0007468492840416729,-0.0004383480700198561,0.0013769635697826743,0.008129142224788666,0.0012334039201959968,0.002106636529788375,0.0010569054866209626,-0.007528538815677166,-0.00024545087944716215,-0.0010706889443099499,0.0001961840025614947,0.0015130306128412485,-0.0024322434328496456,-0.010977503843605518,0.004166409373283386,-0.004214616026729345,0.0028274741489440203,0.001366643700748682,-0.00012284914555493742,-0.000135961061459966,-0.00027496155234985054,0.002704836893826723,0.0001910947321448475,-4.269325654604472e-05,0.00045799827785231173,0.0006786036537960172,-0.004250983241945505,0.002127559157088399,-0.0014173949602991343,0.004615148063749075,-0.0030375474598258734,0.0019918454345315695,-0.0030287120025604963,-0.0022658174857497215,-0.0007007631938904524,-0.002680026926100254,0.004757188260555267,-0.0035628725308924913,0.00105955614708364,0.0065324571914970875,-0.004518664442002773,-0.003672362305223942,-0.009855533950030804,-0.002404535189270973,-0.00109083391726017,-0.0011459322413429618,-0.0021794766653329134,4.7004991756693926e-06,0.012930402532219887,0.00211087753996253,0.0009747351286932826,-0.006272056605666876,-0.00913893710821867,0.0010564107215031981,-0.004276924300938845,0.0023160031996667385,-0.007413606625050306,-0.0007783391047269106,-0.01268382091075182,0.01114785298705101,0.002466419246047735,-0.003144139191135764,0.0018313216278329492,0.0016449980903416872,-0.0012056957930326462,-0.0018440447747707367,-0.006244489923119545,0.01494080200791359,0.00043248129077255726,0.005018472205847502,-5.4285465012071654e-05,0.003481232328340411,-0.003096957691013813,0.009485360234975815,-0.0010450305417180061,0.006258626468479633,-0.0016059096669778228,0.006989925168454647,-0.008994387462735176,0.033319152891635895,0.0018545060884207487,0.005744399037212133,0.004023062065243721,0.0010366191854700446,-0.004199489951133728,-0.0035497606731951237,0.0020601970609277487,0.0013847388327121735,0.007245024666190147,-0.007857397198677063,0.02110343985259533,-3.569552063709125e-06,0.009662847965955734,-0.008406789973378181,-0.00231211562640965,-0.002066664630547166,-0.007131046149879694,-0.004756905138492584,-0.008538192138075829,-0.004811614751815796,0.001973750302568078,0.007661990821361542,-0.004735381808131933,0.02941272035241127,-0.0013426817022264004,0.004126790910959244,-0.0012066146591678262,0.00021555146668106318,0.003473527729511261,0.0016934521263465285,-0.0004835152649320662,0.000657610478810966,-0.00298594799824059,-0.0025072747375816107,-0.00046736394870094955,0.00863566529005766,-0.000456902664154768,2.1912102965870872e-05,-0.00681088212877512,0.004050522577017546,0.004485442768782377,-0.00467081181704998,-0.0041924212127923965,0.0073399534448981285,0.0016161588719114661,-0.004227586556226015,-0.0033738629426807165,0.014578086324036121,-0.00203616451472044,-0.0029104927089065313,0.03582370653748512,0.001279242685995996],\"xaxis\":\"x\",\"y\":[-0.003276954870671034,0.011968850158154964,-0.0024690343998372555,0.010896534658968449,0.02463439851999283,0.004504138603806496,-0.006314396392554045,-0.0034733156207948923,-0.015148083679378033,0.0013615190982818604,-0.015516772866249084,0.0010516748297959566,-0.005969174671918154,-0.004053774289786816,0.002501831855624914,0.007648030761629343,-0.0024495257530361414,-0.0009063481702469289,0.001066094497218728,-0.0008278533350676298,0.0008194065885618329,0.0006284178816713393,-0.0006168257095851004,0.021516306325793266,0.001977355219423771,-0.019867632538080215,-0.018490387126803398,-0.0006193349836394191,0.011022918857634068,-0.0007332778768613935,-0.00028708387981168926,0.014859833754599094,-0.0031067472882568836,-0.02006017602980137,-0.003509717993438244,0.0005305556114763021,-0.0012852154904976487,-0.0019150824518874288,0.0038128825835883617,-0.0035368960816413164,-0.006892239674925804,-0.009536182507872581,-0.0075165932066738605,-0.039258427917957306,-0.0027596172876656055,-0.004137322772294283,0.01941341534256935,0.012854911386966705,0.006432014983147383,-0.011862046085298061,-0.0005677001900039613,-0.009263553656637669,0.012522094883024693,-0.0025314840022474527,-0.013359843753278255,0.0537300668656826,-0.016537098214030266,0.0008152715745382011,0.0007700690184719861,-0.036513056606054306,0.020648233592510223,-0.01483859308063984,0.007865596562623978,-0.007552606984972954,-0.01654876209795475,-0.003732797224074602,-0.0028303368017077446,0.0037496909499168396,0.025589412078261375,0.000553139194380492,-0.0036214697174727917,0.020459577441215515,0.014516307972371578,0.09428074210882187,0.0035050527658313513,0.03303228318691254,-0.011870704591274261,-0.01131417229771614,0.014160907827317715,0.0406869575381279,-0.0001464930101064965,-0.01312644500285387,-0.004529903177171946,0.007631950080394745,0.002730212640017271,0.008291362784802914,-0.013695170171558857,0.0009246907429769635,-0.0009956223657354712,0.02162834070622921,-0.007209505885839462,0.03835660591721535,-0.003117986023426056,0.04912640154361725,-0.01902836374938488,0.0061399126425385475,-0.0026680459268391132,-0.048771779984235764,-0.011770756915211678,-0.041067201644182205,0.08109700679779053,0.04229746013879776,-0.014189994893968105,0.0030708396807312965,0.024279316887259483,-0.010509433224797249,0.031487975269556046,0.02164304256439209,-0.001896492438390851,-0.02792789414525032,-0.012933335267007351,-0.004316684324294329,-0.001957280794158578,0.026262007653713226,-0.06146302446722984,1.0837187767028809,-0.004418645985424519,-0.04550956189632416,0.052880510687828064,-0.0008837291970849037,-0.016890378668904305,-0.18157389760017395,-0.006509166676551104,-0.04653384909033775,-0.017392555251717567,0.007043786812573671,-0.00206768955104053,0.3293505609035492,0.0021621237974613905,0.00773369986563921,-0.01797000877559185,-0.02233058772981167,0.006265730131417513,-0.007989965379238129,-0.014477149583399296,-0.010329541750252247,0.00427826726809144,-0.004690002650022507,-0.009365197271108627,-0.008876628242433071,-0.007687436882406473,-0.01776120811700821,0.11343598365783691,-0.03297686576843262],\"yaxis\":\"y\",\"type\":\"scatter\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Attention Patch\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Output Patch\"}},\"legend\":{\"tracegroupgap\":0},\"title\":{\"text\":\"Scatter plot of output patching vs attention patching\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('fceada07-2286-4de9-8861-70858b938427');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here we see L9H7 is once again doing a ton of work and L10H7 is doing a bit. However the interesting heads here are L7H7 and L11H10 which have the highest attention patching differences but near 0 output patching difference which is why we haven't looked closley at them until now. These are both at two different steps of the algorithm, before and after layer 9 where all the information is moved to the last token."
      ],
      "metadata": {
        "id": "_WYUreYJmuYY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Top Attention Patterns to Look At"
      ],
      "metadata": {
        "id": "_rAPv61mq5da"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "top_k = 10\n",
        "top_heads_by_output_patch = torch.topk(\n",
        "    patched_head_z_diff.abs().flatten(), k=top_k\n",
        ").indices\n",
        "first_mid_layer = 9\n",
        "first_late_layer = 10\n",
        "early_heads = top_heads_by_output_patch[\n",
        "    top_heads_by_output_patch < model.cfg.n_heads * first_mid_layer\n",
        "]\n",
        "mid_heads = top_heads_by_output_patch[\n",
        "    torch.logical_and(\n",
        "        model.cfg.n_heads * first_mid_layer <= top_heads_by_output_patch,\n",
        "        top_heads_by_output_patch < model.cfg.n_heads * first_late_layer,\n",
        "    )\n",
        "]\n",
        "late_heads = top_heads_by_output_patch[\n",
        "    model.cfg.n_heads * first_late_layer <= top_heads_by_output_patch\n",
        "]\n",
        "\n",
        "early = visualize_attention_patterns(\n",
        "    early_heads, cache, tokens[0], title=f\"Top Early Heads (Pre-layer 9)\"\n",
        ")\n",
        "mid = visualize_attention_patterns(\n",
        "    mid_heads, cache, tokens[0], title=f\"Layer 9 Heads\"\n",
        ")\n",
        "late = visualize_attention_patterns(\n",
        "    late_heads, cache, tokens[0], title=f\"Layer 10 and 11 Top Heads\"\n",
        ")\n",
        "\n",
        "HTML(early + mid+ late)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "uNCgtzEoq82-",
        "outputId": "c4663eb4-83df-48b6-9c88-14c8e17b2e31"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<div style='max-width: 700px;'><h2>Top Early Heads (Pre-layer 9)</h2><br/><div id=\"circuits-vis-a829f890-4f2b\" style=\"margin: 15px 0;\"/>\n",
              "    <script crossorigin type=\"module\">\n",
              "    import { render, AttentionHeads } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
              "    render(\n",
              "      \"circuits-vis-a829f890-4f2b\",\n",
              "      AttentionHeads,\n",
              "      {\"attention\": [[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.964978039264679, 0.03502202406525612, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9736860394477844, 0.012611649930477142, 0.01370230969041586, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5276022553443909, 0.14477655291557312, 0.29915332794189453, 0.028467876836657524, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9289711713790894, 0.008137783966958523, 0.025478897616267204, 0.003486581612378359, 0.033925533294677734, 0.0, 0.0, 0.0, 0.0], [0.23700562119483948, 0.19777897000312805, 0.0764579027891159, 0.02846609428524971, 0.37615475058555603, 0.08413666486740112, 0.0, 0.0, 0.0], [0.10975685715675354, 0.009542141109704971, 0.5365830063819885, 0.08856592327356339, 0.11783888190984726, 0.12533414363861084, 0.012379040941596031, 0.0, 0.0], [0.3708023130893707, 0.009788750670850277, 0.010584123432636261, 0.469735711812973, 0.0737849771976471, 0.00874222069978714, 0.037999510765075684, 0.01856233924627304, 0.0], [0.36382824182510376, 0.048604805022478104, 0.15031607449054718, 0.026363348588347435, 0.1105203703045845, 0.03854092210531235, 0.04886185750365257, 0.19599443674087524, 0.01696997880935669]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9931992888450623, 0.006800658069550991, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9138631820678711, 0.06185105815529823, 0.02428581938147545, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7841507792472839, 0.12281658500432968, 0.06913939863443375, 0.0238933265209198, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9331868886947632, 0.018937136977910995, 0.005671657156199217, 0.019392242655158043, 0.02281203866004944, 0.0, 0.0, 0.0, 0.0], [0.6068698763847351, 0.045783236622810364, 0.012801134027540684, 0.09728085994720459, 0.03564603626728058, 0.20161889493465424, 0.0, 0.0, 0.0], [0.8195384740829468, 0.004552430938929319, 0.029312722384929657, 0.013325835578143597, 0.013646610081195831, 0.09730537980794907, 0.022318528965115547, 0.0, 0.0], [0.5116273760795593, 0.06658565998077393, 0.016777895390987396, 0.07740448415279388, 0.021042609587311745, 0.11981628090143204, 0.16749291121959686, 0.01925274357199669, 0.0], [0.4468081593513489, 0.008176143281161785, 0.0062650712206959724, 0.002931079128757119, 0.009723193012177944, 0.19894132018089294, 0.17129847407341003, 0.12842556834220886, 0.027430986985564232]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.782707691192627, 0.21729232370853424, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3597410023212433, 0.0995279997587204, 0.5407310128211975, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.29054659605026245, 0.05059921368956566, 0.09629007428884506, 0.5625640749931335, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2415400594472885, 0.027116142213344574, 0.11136429756879807, 0.0995611846446991, 0.5204183459281921, 0.0, 0.0, 0.0, 0.0], [0.197573721408844, 0.014528391882777214, 0.06546825915575027, 0.08702148497104645, 0.2644214928150177, 0.3709866404533386, 0.0, 0.0, 0.0], [0.21950368583202362, 0.005088517442345619, 0.03149936720728874, 0.007485467009246349, 0.06844767928123474, 0.15236414968967438, 0.5156111121177673, 0.0, 0.0], [0.21485769748687744, 0.006036654114723206, 0.10478471964597702, 0.0031501688063144684, 0.018527556210756302, 0.011887716129422188, 0.028646990656852722, 0.6121084690093994, 0.0], [0.09201169013977051, 0.0013107822742313147, 0.015815753489732742, 0.08356527984142303, 0.03479959815740585, 0.014250626787543297, 0.050449784845113754, 0.11648599058389664, 0.5913105607032776]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9868621826171875, 0.013137818314135075, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9649261832237244, 0.022503796964883804, 0.012570105493068695, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8938682079315186, 0.02231936901807785, 0.041433706879615784, 0.0423787459731102, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7666069865226746, 0.005595801863819361, 0.05271056666970253, 0.09918360412120819, 0.07590307295322418, 0.0, 0.0, 0.0, 0.0], [0.6648315787315369, 0.02107170596718788, 0.029576247557997704, 0.08870969712734222, 0.11850114166736603, 0.07730958610773087, 0.0, 0.0, 0.0], [0.4009304642677307, 0.021049588918685913, 0.15090587735176086, 0.04355732724070549, 0.06900839507579803, 0.2844691574573517, 0.030079243704676628, 0.0, 0.0], [0.6246044635772705, 0.014333520084619522, 0.013596936129033566, 0.022188657894730568, 0.046098772436380386, 0.16726736724376678, 0.08161948621273041, 0.030290797352790833, 0.0], [0.33927738666534424, 0.008646586909890175, 0.010027729906141758, 0.010104580782353878, 0.08100951462984085, 0.3628728985786438, 0.05140048637986183, 0.0670248344540596, 0.06963589787483215]]], \"attentionHeadNames\": [\"L6H1\", \"L8H4\", \"L4H7\", \"L7H9\"], \"tokens\": [\"<|endoftext|>\", \"A\", \" mouse\", \" is\", \" small\", \" but\", \" an\", \" elephant\", \" is\"]}\n",
              "    )\n",
              "    </script></div><div style='max-width: 700px;'><h2>Layer 9 Heads</h2><br/><div id=\"circuits-vis-def0214d-84a8\" style=\"margin: 15px 0;\"/>\n",
              "    <script crossorigin type=\"module\">\n",
              "    import { render, AttentionHeads } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
              "    render(\n",
              "      \"circuits-vis-def0214d-84a8\",\n",
              "      AttentionHeads,\n",
              "      {\"attention\": [[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9860349893569946, 0.013964998535811901, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9551160931587219, 0.027135001495480537, 0.017748890444636345, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9366254806518555, 0.015255346894264221, 0.03175199776887894, 0.01636725477874279, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9476580023765564, 0.00735104363411665, 0.015432869084179401, 0.015446932055056095, 0.014111137017607689, 0.0, 0.0, 0.0, 0.0], [0.6686481833457947, 0.010562391020357609, 0.07823613286018372, 0.054229848086833954, 0.1783769279718399, 0.009946494363248348, 0.0, 0.0, 0.0], [0.37391355633735657, 0.003834002884104848, 0.30471310019493103, 0.009269432164728642, 0.30005908012390137, 0.00566705921664834, 0.002543787006288767, 0.0, 0.0], [0.797453761100769, 0.02370990812778473, 0.026920510455965996, 0.02839079685509205, 0.0921684056520462, 0.006652399897575378, 0.006916511803865433, 0.01778770610690117, 0.0], [0.07823030650615692, 0.0009519269806332886, 0.010218545794487, 0.0013520902721211314, 0.8983247876167297, 0.0010833048727363348, 0.0005594769027084112, 0.009186544455587864, 9.30116802919656e-05]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9991517066955566, 0.0008483462734147906, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9763867259025574, 0.002004027832299471, 0.021609224379062653, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9491040706634521, 0.003288817359134555, 0.03808357194066048, 0.009523462504148483, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9705836772918701, 0.0006270076264627278, 0.019056204706430435, 0.001989722717553377, 0.007743452675640583, 0.0, 0.0, 0.0, 0.0], [0.953666090965271, 0.00284485868178308, 0.007769260089844465, 0.002217331435531378, 0.027272069826722145, 0.006230365484952927, 0.0, 0.0, 0.0], [0.8848670721054077, 0.008865052834153175, 0.09450240433216095, 0.00225519435480237, 0.00651669641956687, 0.0010807545622810721, 0.0019128727726638317, 0.0, 0.0], [0.8574740886688232, 0.001894805347546935, 0.051362063735723495, 0.01619221270084381, 0.008171655237674713, 0.0024432905483990908, 0.0027271020226180553, 0.059734806418418884, 0.0], [0.7872770428657532, 0.0027225189842283726, 0.018669141456484795, 0.0032090358436107635, 0.14122553169727325, 0.0058607985265553, 0.004984687082469463, 0.032240722328424454, 0.003810395486652851]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9913721084594727, 0.008627929724752903, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8011402487754822, 0.16469566524028778, 0.034164104610681534, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7713150382041931, 0.04979563504457474, 0.11119654029607773, 0.06769274920225143, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8393492102622986, 0.006058461498469114, 0.015553104691207409, 0.09117545187473297, 0.04786376655101776, 0.0, 0.0, 0.0, 0.0], [0.739165186882019, 0.00655548507347703, 0.03426862508058548, 0.038614671677351, 0.059330057352781296, 0.12206603586673737, 0.0, 0.0, 0.0], [0.542241632938385, 0.0018716861959546804, 0.007934380322694778, 0.021533802151679993, 0.015422763302922249, 0.37715184688568115, 0.033843882381916046, 0.0, 0.0], [0.4709646701812744, 0.00516089890152216, 0.004534761421382427, 0.05688108131289482, 0.011267161928117275, 0.2579934298992157, 0.1691875010728836, 0.02401037886738777, 0.0], [0.46043065190315247, 0.004062458407133818, 0.012030668556690216, 0.013496359810233116, 0.017333772033452988, 0.16783007979393005, 0.0749841183423996, 0.10869032889604568, 0.14114150404930115]]], \"attentionHeadNames\": [\"L9H7\", \"L9H6\", \"L9H10\"], \"tokens\": [\"<|endoftext|>\", \"A\", \" mouse\", \" is\", \" small\", \" but\", \" an\", \" elephant\", \" is\"]}\n",
              "    )\n",
              "    </script></div><div style='max-width: 700px;'><h2>Layer 10 and 11 Top Heads</h2><br/><div id=\"circuits-vis-b74387ad-c98a\" style=\"margin: 15px 0;\"/>\n",
              "    <script crossorigin type=\"module\">\n",
              "    import { render, AttentionHeads } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
              "    render(\n",
              "      \"circuits-vis-b74387ad-c98a\",\n",
              "      AttentionHeads,\n",
              "      {\"attention\": [[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9711289405822754, 0.028871025890111923, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.962450385093689, 0.0038900121580809355, 0.03365957736968994, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.85878986120224, 0.015433978289365768, 0.07241909205913544, 0.05335712060332298, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7033705115318298, 0.01538859587162733, 0.11321327835321426, 0.059626255184412, 0.10840138047933578, 0.0, 0.0, 0.0, 0.0], [0.5705567598342896, 0.012362794950604439, 0.012375605292618275, 0.03007964789867401, 0.32326540350914, 0.05135980620980263, 0.0, 0.0, 0.0], [0.3255925178527832, 0.018225088715553284, 0.44416120648384094, 0.03159342333674431, 0.09710187464952469, 0.018276162445545197, 0.06504977494478226, 0.0, 0.0], [0.7573192119598389, 0.006619773339480162, 0.024497462436556816, 0.10662192106246948, 0.016511093825101852, 0.026536984369158745, 0.016377350315451622, 0.04551612585783005, 0.0], [0.09326817095279694, 0.0063356757164001465, 0.0017902236431837082, 0.005481854546815157, 0.7756001949310303, 0.014800457283854485, 0.030199307948350906, 0.049663979560136795, 0.022860107943415642]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9824504256248474, 0.01754951849579811, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9570852518081665, 0.0024241425562649965, 0.04049060493707657, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8858200907707214, 0.009392774663865566, 0.08570577204227448, 0.019081400707364082, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8826439380645752, 0.0028362898156046867, 0.0866345539689064, 0.002867253962904215, 0.025018051266670227, 0.0, 0.0, 0.0, 0.0], [0.8823736310005188, 0.005670490674674511, 0.009323567152023315, 0.008928066119551659, 0.059599556028842926, 0.03410470485687256, 0.0, 0.0, 0.0], [0.8217462301254272, 0.03471330553293228, 0.05128922685980797, 0.0036314306780695915, 0.053689006716012955, 0.012227464467287064, 0.02270328253507614, 0.0, 0.0], [0.5749940872192383, 0.009073768742382526, 0.11085585504770279, 0.04800565913319588, 0.10685914009809494, 0.041319381445646286, 0.02710012160241604, 0.08179203420877457, 0.0], [0.5314610004425049, 0.010446934029459953, 0.006295223254710436, 0.009326216764748096, 0.39668846130371094, 0.013130798004567623, 0.015078941360116005, 0.009814027696847916, 0.0077583217062056065]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9506229758262634, 0.04937697574496269, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8636487126350403, 0.0646706372499466, 0.07168063521385193, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7691375613212585, 0.029964299872517586, 0.12727312743663788, 0.07362502813339233, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6092661619186401, 0.02264225110411644, 0.1738659292459488, 0.11226367950439453, 0.08196202665567398, 0.0, 0.0, 0.0, 0.0], [0.3979811370372772, 0.007602523546665907, 0.07248664647340775, 0.26127785444259644, 0.15910464525222778, 0.10154716670513153, 0.0, 0.0, 0.0], [0.33622974157333374, 0.060418751090765, 0.2987986207008362, 0.16040712594985962, 0.05644509196281433, 0.039555490016937256, 0.04814521223306656, 0.0, 0.0], [0.535023033618927, 0.02326866425573826, 0.04021435230970383, 0.20968644320964813, 0.04370691999793053, 0.04488035663962364, 0.06491724401712418, 0.03830297663807869, 0.0], [0.48421016335487366, 0.01896764524281025, 0.030104398727416992, 0.09323664754629135, 0.16630412638187408, 0.06683766096830368, 0.05963825434446335, 0.03186190500855446, 0.048839252442121506]]], \"attentionHeadNames\": [\"L10H7\", \"L10H1\", \"L11H10\"], \"tokens\": [\"<|endoftext|>\", \"A\", \" mouse\", \" is\", \" small\", \" but\", \" an\", \" elephant\", \" is\"]}\n",
              "    )\n",
              "    </script></div>"
            ]
          },
          "metadata": {},
          "execution_count": 55
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Note: L4H7 looks like an induction head pattern"
      ],
      "metadata": {
        "id": "EbpLB_82yffh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can also try just swapping the adjectives, although this should result in much different answers than expected because now the sentence is constructed as if the answer should be an even stronger adjective rather than an opposite \"bigger\" vs \"small\", even though we are describing a \"mouse\" in this case."
      ],
      "metadata": {
        "id": "r1uABFGl2ECM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "corrupted_prompts_2 = [\n",
        "\"An elephant is small, but a mouse is\",\n",
        "\"A kitten is old, but a tortoise is\",\n",
        "\"A rainforest is dry, but a desert is\",\n",
        "\"A brick is soft, but a pillow is\",\n",
        "\"A sponge is hard, but a rock is\",\n",
        "\"A weed is beautiful, but a rose is\"\n",
        "]"
      ],
      "metadata": {
        "id": "rFqoKZux2Dpm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Current Conclusion/Thoughts\n",
        "\n",
        "Throughout this experiment, I investigated if I could find a circuit for the task of contrasting adjectives when presented with opposite nouns. I was able to isolate a large amount of computation happening in layer 9 and 10.\n",
        "\n",
        "Before layer 9, most of the computation is being done on the first adjective and the \"but\" token, mostly by MLPs which seem to be processing the adjective and figuring out a potential opposite.\n",
        "\n",
        "Then at layer 9, all of the information is moved from the first adjective to the final token and both attention and MLPs are fired in layer 9 and 10. Layer 9 head 7 makes the largest logit difference across all metrics except \"attention/QK\", which means I believe that is where the opposite is chosen based on the processing done in that layer and previous layers.\n",
        "\n",
        "Overall, this seems like just the start of trying to uncover a circuit and I would need to dig even deeper by analyzing way more prompts.\n",
        "\n",
        "I spent a lot of time trying to come up with better prompts and debugging errors in my prompts like different token positions etc. Throughout this process, I discovered gpt-2 small is very context dependent.\n",
        "\n",
        "For instance, for the majority of the research I had my prompt look like:\n",
        "\n",
        "    \"A mouse is small, but an elephant is\"\n",
        "\n",
        "However, I decided to remove the comma because it was weird to me how the comma was being attended to more than the 'but' and even more than 'mouse' most of the time. Then my new prompt was\n",
        "\n",
        "    \"A mouse is small but an elephant is\"\n",
        "\n",
        "and I had a 0.7 improvement in logit difference across prompts. Just by removing the comma!\n",
        "\n",
        "I believe this task wasn't the best because there are multiple things that go into the token prediction like noun adjective association and understanding the sentence is comparing opposite adjectives.\n",
        "\n",
        "### Areas of Improvement:\n",
        "\n",
        "I feel like I lacked a lot of the intuion to understand what different circuits were doing on a fundamental level. In the EAD colab, which I followed closely to get up to speed on this code, conclusions are made about which heads do deduplication and I wasn't able to understand how they figured that out and then apply that methodology here. In this research, I felt like I was doing a lot of guessing about what's going on, which is necessary to some extend. I'm looking forward to having my intution improve from either mentorship or just putting reps in.\n",
        "\n",
        "In addition, I'm excited to learn about how these type of experiment results (or more robust results like the Interpretability in the Wild Paper) get unified in a broader understanding of the gpt-2 small model and LLMs in general. Part of me feels like there might be all this specific, niche work going on but no way to unify all these learnings so that all researchers might better get a global view of the model. Maybe an idea for this could be implemented in the TransformerLens Library.\n",
        "\n",
        "\n",
        "## Exciting other experiments\n",
        "\n",
        "After reading the *Beyond A: Better Planning with Transformers via Search Dynamics Bootstrapping* paper, extremely motivated to see if we can use these mech interp techniques to uncover the types of algorithms that the model is learning.\n",
        "\n",
        "After doing this excersise and immersing myself within alginment and interpretability for a short time, I believe the key to many scientific and mathematical discoveries are hidden within LLMs and their internal algorithms. Sometimes these algorithms are probably not optimal or correct, however other times, in the case of search, it seems that large language models are innovating on our best approaches like A*.\n",
        "\n",
        "I would first like to start by investigating how models think about sorting, CS 101. I think because it is multi token output, it will be a lot harder than the excersise above or IOI but I'm sure we can reveal some interesting dynamics about how different sized models sort (I assume Gpt2-small is either incapable or does some algorithm with similar efficancy to bogo sort lol)."
      ],
      "metadata": {
        "id": "erZVkLr0Cxd7"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "cS0bEtBqEAnK"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}