iit_equality.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interchange Intervention Training: Equality learning tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"Atticus Geiger\"\n",
    "__version__ = \"CS224u, Stanford, Spring 2023\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Contents\n",
    "\n",
    "1. [Overview](#Overview)\n",
    "1. [Set-up](#Set-up)\n",
    "1. [The hierarchical equality task](#The-hierarchical-equality-task)\n",
    "1. [The high-level causal model](#The-high-level-causal-model)\n",
    "    1. [The algorithm with no intervention](#The-algorithm-with-no-intervention)\n",
    "    1. [The algorithm with an intervention](#The-algorithm-with-an-intervention)\n",
    "    1. [The algorithm with an interchange intervention](#The-algorithm-with-an-interchange-intervention)\n",
    "1. [A fully-connected feed-forward neural network](#A-fully-connected-feed-forward-neural-network)\n",
    "    1. [Basic intervention: zeroing out part of a hidden layer](#Basic-intervention:-zeroing-out-part-of-a-hidden-layer)\n",
    "    1. [An interchange intervention](#An-interchange-intervention)\n",
    "1. [Causal abstraction](#Causal-abstraction)\n",
    "    1. [Alignment](#Alignment)\n",
    "    1. [Interchange intervention](#Interchange-intervention)\n",
    "    1. [Evaluation](#Evaluation)\n",
    "1. [Interchange Intervention Training (IIT)](#Interchange-Intervention-Training-(IIT))\n",
    "    1. [IIT on variable V1](#IIT-on-variable-V1)\n",
    "    1. [IIT on variables V1 and V2](#IIT-on-variables-V1-and-V2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This notebook is a hands-on introduction to __causal abstraction analysis__ and __interchange intervention training__ with neural networks.\n",
    "\n",
    "In causal abstraction analysis, we assess whether trained models conform to high-level causal models that we specify, not just in terms of their input–output behavior, but also in terms of their internal dynamics. \n",
    "\n",
    "The core technique is the __interchange intervention__, in which we actively manipulate internal states in the high-level causal model and in the neural network to see whether the two models show the same behavior in these counterfactual states.\n",
    "\n",
    "In interchange intervention training, we go beyond passive analysis by actively training networks to conform to the high-level causal model.\n",
    "\n",
    "To motivate and illustrate these concepts, we're going to focus on a challenging hierarchical equality task, building on work by [Geiger, Carstensen, Frank, and Potts (2020)](https://arxiv.org/abs/2006.07968)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import random\n",
    "import copy\n",
    "import itertools\n",
    "import numpy as np\n",
    "from sklearn.metrics import classification_report\n",
    "from torch_deep_neural_classifier import TorchDeepNeuralClassifier\n",
    "from torch_deep_neural_classifier_iit import TorchDeepNeuralClassifierIIT\n",
    "import iit\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.fix_random_seeds()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The hierarchical equality task\n",
    "\n",
    "This section builds on results presented in [Geiger, Carstensen, Frank, and Potts (2020)](https://arxiv.org/abs/2006.07968). We will use a hierarchical equality task ([Premack 1983](https://www.cambridge.org/core/services/aop-cambridge-core/content/view/7DF6F2D22838F7546AF7279679F3571D/S0140525X00015077a.pdf/div-class-title-the-codes-of-man-and-beasts-div.pdf)) to present interchange intervention training (IIT). \n",
    "\n",
    "We define the hierarchical equality task as follows: The input is two pairs of objects and the output is **True** if both pairs contain the same object or if both pairs contain different objects and **False** otherwise.  For example, `AABB` and `ABCD` are both labeled **True**, while `ABCC` and `BBCD` are both labeled **False**. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The high-level causal model\n",
    "\n",
    "Let $\\mathcal{A}$ be the simple tree-structured algorithm that solves this task by applying a simple equality relation three times: Compute whether the first two inputs are equal, compute whether the second two inputs are equal, then compute whether the truth-valued outputs of these first two computations are equal. Here's a visual depiction of the algorithm:\n",
    "\n",
    "<img src=\"fig/IIT/PremackFunctions.png\" width=\"500\"/>\n",
    "<img src=\"fig/IIT/PremackGraph.png\" width=\"500\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's a Python implementation of $\\mathcal{A}$ that supports the interventions we'll want to do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_A(ex, intervention):\n",
    "    graph = {}\n",
    "    for i, obj in enumerate(ex):\n",
    "        graph[\"input\" + str(i+1)] = obj\n",
    "    if \"V1\" in intervention:\n",
    "        graph[\"V1\"] = intervention[\"V1\"]\n",
    "    else:\n",
    "        graph[\"V1\"] = graph[\"input1\"] == graph[\"input2\"]\n",
    "    if \"V2\" in intervention:\n",
    "        graph[\"V2\"] = intervention[\"V2\"]\n",
    "    else:\n",
    "        graph[\"V2\"] = graph[\"input3\"] == graph[\"input4\"]\n",
    "    graph[\"output\"] = graph[\"V1\"] == graph[\"V2\"]\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The algorithm with no intervention\n",
    "\n",
    "Let's first observe the behavior of the algorithm when we provide the input **(pentagon,pentagon, triangle, square)** with no interventions. Here is a visual depiction:\n",
    "\n",
    "<img src=\"fig/IIT/PremackNoIntervention.png\" width=\"500\"/>\n",
    "\n",
    "And here is the computation using `compute_A`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input1': 'pentagon',\n",
       " 'input2': 'pentagon',\n",
       " 'input3': 'triangle',\n",
       " 'input4': 'square',\n",
       " 'V1': True,\n",
       " 'V2': False,\n",
       " 'output': False}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_A(\n",
    "    (\"pentagon\", \"pentagon\", \"triangle\", \"square\"), \n",
    "    intervention={})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The algorithm with an intervention\n",
    "\n",
    "Let's now see the behavior of the algorithm when we provide the input **(square,pentagon,triangle, triangle)** with an intervention setting **V1** to **False**. First, a visual depiction:\n",
    "\n",
    "<img src=\"fig/IIT/PremackIntervention.png\" width=\"500\"/>\n",
    "\n",
    "And then the same computation with `compute_A`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input1': 'square',\n",
       " 'input2': 'pentagon',\n",
       " 'input3': 'triangle',\n",
       " 'input4': 'triangle',\n",
       " 'V1': True,\n",
       " 'V2': True,\n",
       " 'output': True}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_A(\n",
    "    (\"square\", \"pentagon\", \"triangle\", \"triangle\"), \n",
    "    intervention={\"V1\": True})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that, in this example, even though the left two inputs are not the same, the intervention has changed the intermediate prediction for those two inputs from **False** to **True**, and thus the algorithm outputs **True**, since its output is determined by **V1** and **V2**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The algorithm with an interchange intervention\n",
    "\n",
    "Finally, let's observe the behavior of the algorithm when we provide the base input **(square, pentagon, triangle, triangle)** with an intervention setting **V1** to be the value it would be for the source input **(pentagon, pentagon, triangle, square)**. Here's a diagram in which the dashed line indicates the interchange intervention:\n",
    "\n",
    "<img src=\"fig/IIT/algorithmII.png\" width=\"600\"/>\n",
    "\n",
    "And here is the corresponding interchange intervention in code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_interchange_A(base, source, variable):\n",
    "    # Run the algorithm on `source`:\n",
    "    src_output = compute_A(source, intervention={})\n",
    "    # Get the source value for `variable`:\n",
    "    val = src_output[variable]\n",
    "    # Process `base` with the intervention setting `variable`\n",
    "    # to the value it had in `source`:        \n",
    "    return compute_A(base, intervention={variable: val})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input1': 'pentagon',\n",
       " 'input2': 'pentagon',\n",
       " 'input3': 'triangle',\n",
       " 'input4': 'square',\n",
       " 'V1': False,\n",
       " 'V2': False,\n",
       " 'output': True}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_interchange_A(\n",
    "    base=(\"pentagon\", \"pentagon\", \"triangle\", \"square\"),    # base: T F ==> F\n",
    "    source=(\"square\", \"pentagon\", \"triangle\", \"triangle\"),  # source: F T ==> F\n",
    "    variable=\"V1\") # Will set base V1 to be source V1, leading to F F ==> T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A fully-connected feed-forward neural network\n",
    "\n",
    "We've now seen how interventions work in our high-level causal model. We turn now to doing parallel work in our neural network, which will be a fully-connected feed-forward neural network with three hidden layers. The following code simply extends `TorchDeepNeuralClassifier` with a method `retrieve_activations` that supports interventions on PyTorch computation graphs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class InterventionableTorchDeepNeuralClassifier(TorchDeepNeuralClassifier):\n",
    "    def __init__(self, **base_kwargs):\n",
    "        super().__init__(**base_kwargs)\n",
    "        \n",
    "    def make_hook(self, gets, sets, layer):\n",
    "        def hook(model, input, output):\n",
    "            layer_gets, layer_sets = [], []\n",
    "            if gets is not None and layer in gets:\n",
    "                layer_gets = gets[layer]\n",
    "            if sets is not None and layer in sets:\n",
    "                layer_sets = sets[layer]\n",
    "            for set in layer_sets:\n",
    "                output = torch.cat(\n",
    "                    [output[:, :set[\"start\"]], \n",
    "                     set[\"intervention\"], \n",
    "                     output[:, set[\"end\"]: ]], \n",
    "                    dim=1)\n",
    "            for get in layer_gets:\n",
    "                k = f'{get[\"layer\"]}-{get[\"start\"]}-{get[\"end\"]}'\n",
    "                self.activation[k] = output[:, get[\"start\"]: get[\"end\"] ]\n",
    "            return output\n",
    "        return hook\n",
    "\n",
    "    def _gets_sets(self, gets=None, sets=None):\n",
    "        handlers = []\n",
    "        for layer in range(len(self.layers)):\n",
    "            hook = self.make_hook(gets, sets, layer)\n",
    "            both_handler = self.layers[layer].register_forward_hook(hook)\n",
    "            handlers.append(both_handler)\n",
    "        return handlers\n",
    "\n",
    "    def retrieve_activations(self, X, get, sets):\n",
    "        if sets is not None and \"intervention\" in sets:\n",
    "            sets[\"intervention\"] = sets[\"intervention\"].type(torch.FloatTensor).to(self.device)\n",
    "        X = X.type(torch.FloatTensor).to(self.device)\n",
    "        self.activation = {}\n",
    "        get_val = {get[\"layer\"]: [get]} if get is not None else None\n",
    "        set_val = {sets[\"layer\"]: [sets]} if sets is not None else None\n",
    "        handlers = self._gets_sets(get_val, set_val)\n",
    "        logits = self.model(X)\n",
    "        for handler in handlers:\n",
    "            handler.remove()\n",
    "        return self.activation[f'{get[\"layer\"]}-{get[\"start\"]}-{get[\"end\"]}']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The module `iit` provides some dataset functions for equality learning. Here we define a simple an equality dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "embedding_dim = 4\n",
    "\n",
    "n_examples = 10000\n",
    "\n",
    "X_train, X_test, y_train, y_test, test_dataset = iit.get_equality_dataset(\n",
    "    embedding_dim, n_examples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The examples in this dataset are 16-dimensional vectors: the concatenation of 4 4-dimensional vectors. Here's the first example with its label:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 0.4365, -0.0289, -0.2255,  0.3735,  0.4365, -0.0289, -0.2255,  0.3735,\n",
       "          0.3634,  0.1236,  0.0243, -0.1352,  0.4411, -0.2673, -0.4682, -0.3980],\n",
       "        dtype=torch.float64),\n",
       " 0)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train[0], y_train[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The label for this example is determined by whether the equality value for the first two inputs matches the equality value for the second two inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "left = torch.equal(\n",
    "    X_train[0][: embedding_dim],\n",
    "    X_train[0][embedding_dim: embedding_dim*2])\n",
    "\n",
    "left"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "right = torch.equal(\n",
    "    X_train[0][embedding_dim*2: embedding_dim*3],\n",
    "    X_train[0][embedding_dim*3: ])\n",
    "\n",
    "right"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "int(left == right)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how our model does out-of-the-box on this task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping after epoch 450. Training loss did not improve more than tol=1e-05. Final error is 0.027576118358410895."
     ]
    }
   ],
   "source": [
    "model = InterventionableTorchDeepNeuralClassifier(\n",
    "    hidden_dim=embedding_dim * 4,\n",
    "    hidden_activation=torch.nn.ReLU(), \n",
    "    num_layers=3)\n",
    "\n",
    "_ = model.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This neural network achieves near perfect performance on its train set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Results\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       1.00      1.00      1.00      5000\n",
      "           1       1.00      1.00      1.00      5000\n",
      "\n",
      "    accuracy                           1.00     10000\n",
      "   macro avg       1.00      1.00      1.00     10000\n",
      "weighted avg       1.00      1.00      1.00     10000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_preds = model.predict(X_train)\n",
    "\n",
    "print(\"Train Results\")\n",
    "print(classification_report(y_train, train_preds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And it generalizes perfectly to the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Results\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.99      1.00      0.99      5000\n",
      "           1       1.00      0.99      0.99      5000\n",
      "\n",
      "    accuracy                           0.99     10000\n",
      "   macro avg       0.99      0.99      0.99     10000\n",
      "weighted avg       0.99      0.99      0.99     10000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Test Results\")\n",
    "\n",
    "test_preds = model.predict(X_test)\n",
    "\n",
    "print(classification_report(y_test, test_preds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Does it implement our high-level model of the problem, though?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic intervention: zeroing out part of a hidden layer\n",
    "\n",
    "To begin to build towards the full interchange intervention, let's consider a simpler intervention, where we zero out the first `embedding_dim` neurons in the first hidden layer.\n",
    "\n",
    "Our basic inputs are random vectors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = X_train[0][: embedding_dim]\n",
    "b = X_train[1][: embedding_dim]\n",
    "c = X_train[2][: embedding_dim]\n",
    "\n",
    "X_same_different = torch.cat((a, a, b, c)).unsqueeze(0)\n",
    "\n",
    "X_different_same = torch.cat((a, b, c, c)).unsqueeze(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here we define two different inputs for use in later examples. We'll use training examples so that we are sure to see the full logic of these interventions; the next section will consider test examples in the context of a full abstraction analysis:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the intervention, we first specify that we want it target layer 1 (the outer key), and then we give a list of specifications for interventions at that layer (here we will do just one). So that we can study the full layer before and after the intervention, we specify the entire layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeroing_get_coord = {\n",
    "    \"layer\": 1,\n",
    "    \"start\": 0, \n",
    "    \"end\": \n",
    "    embedding_dim*4\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we specify the intervention itself: in layer 1, the first `embedding_layer` inputs will be turned into 0s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeroing_intervention = {\n",
    "    \"layer\": 1,\n",
    "    \"start\": 0,  \n",
    "    \"end\": embedding_dim, \n",
    "    \"intervention\": torch.zeros((1,embedding_dim))\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the `X_same_different` input, the network computes the following values at our intervention site, without any intervention:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0000, 0.0655, 0.4071, 0.0000, 0.1467, 0.5594, 0.1825, 0.0000, 0.1477,\n",
       "         0.0000, 0.1222, 0.0000, 0.0333, 0.6937, 0.0000, 0.0000]],\n",
       "       grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, zeroing_get_coord, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here are the values computed with the intervention:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.1467, 0.5594, 0.1825, 0.0000, 0.1477,\n",
       "         0.0000, 0.1222, 0.0000, 0.0333, 0.6937, 0.0000, 0.0000]],\n",
       "       grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, zeroing_get_coord, zeroing_intervention)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also see how the intervention affects outputs. To that, we specify the final layer (the two logits) as the coordinate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeroing_output_coord = {\n",
    "    \"layer\": 3, \n",
    "    \"start\": 0, \n",
    "    \"end\": 2}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the outputs without an intervention:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 3.0205, -3.0677]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, zeroing_output_coord, sets=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And with the intervention we specified above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-2.9687,  3.2119]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, zeroing_output_coord, zeroing_intervention)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An interchange intervention\n",
    "\n",
    "We're now ready to do a full intervention. The only change from the above is that, instead of simply zeroing out some neurons, we'll replace them with the corresponding values determined by a distinct input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll again target the first `embedding_dim` units in the first hidden layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "ii_coord = {\"layer\": 1, \"start\": 0, \"end\": embedding_dim}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our **source** input, we'll use `X_different_same`. The first step is to get the activations for this input at our coordinate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0784, 0.0000, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intervention_get = model.retrieve_activations(X_different_same, ii_coord, None)\n",
    "\n",
    "intervention_get"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we define the intervention using these values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "ii_set = {\n",
    "    \"layer\": 1, \n",
    "    \"start\": 0, \n",
    "    \"end\": embedding_dim, \n",
    "    \"intervention\": intervention_get}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now turn to our __base__ input, which will be `X_same_different`. With no intervention, this has the following values at our intervention site:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0000, 0.0655, 0.4071, 0.0000]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, ii_coord, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then we can verify that the intervention works as we intended it to; these values should be the same as `intervention_get` above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0784, 0.0000, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, ii_coord, ii_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can see what the intervention does to the network's predictions. We specify the coordinates of the output logits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "ii_output_coord = {\"layer\": 3, \"start\": 0, \"end\": 2}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With no intervention, the input `X_same_different` delivers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 3.0205, -3.0677]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, ii_output_coord, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the intervention, that same input delivers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-3.0118,  3.1399]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.retrieve_activations(X_same_different, ii_output_coord, ii_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If our target coordinates for the intervention were a modular encoding of the value for the first two inputs, then this intervention would have change the network's prediction from `0` to `1`, since we would have effectively created a **(different, different)** input. It's unlikely that this happened, suggesting that our hypothesis about where this information is encoded is false. A full-fledged causal abstraction analysis will allow us to assess this more comprehensively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Causal abstraction\n",
    "\n",
    "To recap:\n",
    "\n",
    "1. We defined a **high-level causal model** (a tree-structured algorithm) that solves the hierarchical equality task.\n",
    "\n",
    "1. We trained a **low-level fully connected neural network** that seeks to solv the hierarchical equality task.\n",
    "\n",
    "1. We peformed illustrative interventions on both these networks to begin to get a feel for whether the high-level model is an abstraction of the lower-level neural one.\n",
    "\n",
    "The formal theory of **causal abstraction** describes the conditions that must hold for the high-level tree structured algorithm to be a **simplified and faithful description** of the neural network. \n",
    "\n",
    "In essence: an high-level model is a causal abstraction of a neural network if and only if for all base and source inputs, the algorithm and network provides the same output, for some alignment between these two models.\n",
    "\n",
    "Below, we define an alignment between the neural network and the algorithm and a function to compute the **interchange intervention accuracy** (II accuracy) for a high-level variable: the percentage of aligned interchange interventions that the network and algorithm produce the same output on. When the II accuracy is 100%, the causal abstraction relation holds between the network and a simplified version of the algorithm where only one high-level variable exists."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is to specify an alignment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "alignment = {\n",
    "    \"V1\": {\"layer\": 1, \"start\": 0, \"end\": embedding_dim}, \n",
    "    \"V2\": {\"layer\": 1, \"start\": embedding_dim, \"end\": embedding_dim*2}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In essence, this reflects a hypothesis that we will find the equality label for the first two inputs in the first four neurons in layer 1, and that we'll find the equality label for the second two inputs in the next four neurons in layer 1. This is of course just one of a great many hypotheses we could state. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interchange intervention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `interchange_intervention` packages up the multi-step process we walked through above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def interchange_intervention(model, base, source, get_coord, output_coord):\n",
    "    intervention = model.retrieve_activations(source, get_coord, None)\n",
    "    get_coord[\"intervention\"] = intervention\n",
    "    return model.retrieve_activations(base, output_coord, get_coord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_coord = {\"layer\": 3, \"start\": 0, \"end\": 2}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-3.0118,  3.1399]], grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "interchange_intervention(\n",
    "    model, \n",
    "    base=X_same_different, \n",
    "    source=X_different_same, \n",
    "    get_coord=ii_coord, \n",
    "    output_coord=output_coord)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So that we can run out high-level model on our vector examples, we define a helper function to parse them into their component inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_input(tensor, embedding_dim):\n",
    "    return [tuple(tensor[0, embedding_dim*k:embedding_dim*(k+1)].flatten().tolist()) \n",
    "            for k in range(4)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Illustration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_A(convert_input(X_same_different, embedding_dim), {})['output']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_A(convert_input(X_different_same, embedding_dim), {})['output']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now if we perform an intervention on the **V1** variable with `X_different_same` as the base and `X_same_different` as the source, we effectively create a **(False, False)** example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_interchange_A(\n",
    "    convert_input(X_different_same, embedding_dim),\n",
    "    convert_input(X_same_different, embedding_dim),\n",
    "    variable=\"V1\")['output']    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `ii_evaluation` puts these pieces together in the context of a full evaluation on a set of examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ii_evaluation(X_assess, model, variable, output_coord):\n",
    "    labels = []\n",
    "    predictions = []\n",
    "    for base, source in itertools.product(X_assess, repeat=2):\n",
    "        base = base.unsqueeze(0)\n",
    "        source = source.unsqueeze(0)\n",
    "        # Run the high-level model with the intervention:\n",
    "        algorithm_output = compute_interchange_A(\n",
    "            convert_input(base, embedding_dim), \n",
    "            convert_input(source, embedding_dim), \n",
    "            variable)\n",
    "        # Get the high-level model's label:\n",
    "        labels.append(int(algorithm_output[\"output\"]))\n",
    "        # Run the neural model with the intervention:\n",
    "        network_output = interchange_intervention(\n",
    "            model, \n",
    "            base,\n",
    "            source,\n",
    "            alignment[variable],\n",
    "            output_coord)\n",
    "        # Get the neural model's prediction with the intervention:\n",
    "        pred = network_output.argmax(axis=1)\n",
    "        predictions.append(int(pred))\n",
    "    return labels, predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's assess the hypothesis that **V1** is encoded at our chosen site, using a sample of test cases for efficiency:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.50      0.42      0.46      5016\n",
      "           1       0.50      0.58      0.54      4984\n",
      "\n",
      "    accuracy                           0.50     10000\n",
      "   macro avg       0.50      0.50      0.50     10000\n",
      "weighted avg       0.50      0.50      0.50     10000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(*ii_evaluation(X_test[: 100], model, \"V1\", output_coord)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then the corresponding assessment for **V2**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.55      0.44      0.49      5016\n",
      "           1       0.53      0.64      0.58      4984\n",
      "\n",
      "    accuracy                           0.54     10000\n",
      "   macro avg       0.54      0.54      0.53     10000\n",
      "weighted avg       0.54      0.54      0.53     10000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(*ii_evaluation(X_test[: 100], model, \"V2\", output_coord)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have low accuracy for both **V1** and **V2**, meaning that under this alignment the neural network does not compute either variable. In other words, we have no evidence that this network computes simple equality relations to solve this hierarchical equality task. The goal of interchange intervention training is to change this. We turn to that method next."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interchange Intervention Training (IIT)\n",
    "\n",
    "Interchange Intervention Training (IIT) is a method for training a neural network to conform to the causal structure of a high-level algorithm. Conceptually, it is a direct extension of the causal abstraction analysis we just performed, except instead of **evaluating** whether the neural network and algorithm produce the same outputs under aligned interchange interventions, we are now **training** the neural network to produce the output of the algorithm under aligned interchange interventions.\n",
    "\n",
    "IIT was developed by [Geiger\\*, Wu\\*, Lu\\*, Rozner, Kreiss, Icard, Goodman, and Potts (2021)](https://arxiv.org/abs/2112.00826), and it is used for model distillation [ Wu\\*, Geiger\\*, Rozner, Kreiss, Lu, Icard, Goodman, and Potts (2021)](https://arxiv.org/abs/2112.02505)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IIT on variable V1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our first intervention will target the first four dimensions of layer 1 in the network, but now we will be training the network in part to play the role of the **V1** variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "V1 = 0\n",
    "\n",
    "V1_id_to_coords = {\n",
    "    V1: [{\"layer\": 1, \"start\": 0, \"end\": embedding_dim}]    \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we create an equality dataset that includes examples for IIT training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_size = 10000\n",
    "\n",
    "iit_equality_dataset = iit.get_IIT_equality_dataset(\"V1\", embedding_dim, data_size)\n",
    "\n",
    "X_base_train, X_sources_train, y_base_train, y_IIT_train, interventions = iit_equality_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This dataset has the following components:\n",
    "\n",
    "* `X_base_train`: a regular set of train examples\n",
    "* `y_base_train`: a regular set of train labels\n",
    "* `X_sources_train`: a list additional train sets (here, a singleton list of them) for counterfactuals\n",
    "* `y_IIT_train`: a list of labels for the examples in `X_sources_train`.\n",
    "* `interventions`: a list of intervention sites (here, all `0` corresponding to our key for \"V1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our model is a deep classifier like the one we used above, but now one that can do IIT:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "iit_model = TorchDeepNeuralClassifierIIT(\n",
    "    hidden_dim=embedding_dim*4, \n",
    "    hidden_activation=torch.nn.ReLU(), \n",
    "    num_layers=3,\n",
    "    id_to_coords=V1_id_to_coords)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model is fit using our IIT dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping after epoch 642. Training loss did not improve more than tol=1e-05. Final error is 0.00104618080513319."
     ]
    }
   ],
   "source": [
    "_ = iit_model.fit(\n",
    "    X_base_train, \n",
    "    X_sources_train, \n",
    "    y_base_train, \n",
    "    y_IIT_train, \n",
    "    interventions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To evaluate this model, we create a fresh IIT equality dataset consisting of 100 examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_base_test, X_sources_test, y_base_test, y_IIT_test, interventions = iit.get_IIT_equality_dataset(\n",
    "    \"V1\", embedding_dim, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "IIT_preds, base_preds = iit_model.iit_predict(X_base_test, X_sources_test, interventions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This IIT-trained model does well in terms of a standard behavioral tests:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       1.00      1.00      1.00        48\n",
      "           1       1.00      1.00      1.00        48\n",
      "\n",
      "    accuracy                           1.00        96\n",
      "   macro avg       1.00      1.00      1.00        96\n",
      "weighted avg       1.00      1.00      1.00        96\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(y_base_test, base_preds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importantly, it _also_ performs perfectly on counterfactual examples – certainly a marked improvement over the model we studied above that did no IIT:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       1.00      1.00      1.00        48\n",
      "           1       1.00      1.00      1.00        48\n",
      "\n",
      "    accuracy                           1.00        96\n",
      "   macro avg       1.00      1.00      1.00        96\n",
      "weighted avg       1.00      1.00      1.00        96\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(y_IIT_test, IIT_preds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course, we did only one kind of IIT: we pushed the first `embedding_dim` neurons in layer 1 to conform to **V1** in the causal model. As a results, we still have low counterfactual accuracy for **V2**, meaning that, under this alignment, the neural network doesn't compute whether the second pair of inputs are equal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.50      0.50      0.50      5000\n",
      "           1       0.50      0.50      0.50      5000\n",
      "\n",
      "    accuracy                           0.50     10000\n",
      "   macro avg       0.50      0.50      0.50     10000\n",
      "weighted avg       0.50      0.50      0.50     10000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "X_base_test_V2, X_sources_test_V2, y_base_test_V2, y_IIT_test_V2, interventions_V2 = iit.get_IIT_equality_dataset(\n",
    "    \"V2\", embedding_dim,data_size)\n",
    "\n",
    "IIT_preds_V2, base_preds_V2 = iit_model.iit_predict(X_base_test_V2, X_sources_test_V2, interventions_V2)\n",
    "\n",
    "print(classification_report(y_IIT_test_V2, IIT_preds_V2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IIT on variables V1 and V2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To address this, we can simply train the network to compute both **V1** and **V2**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "V2 = 1\n",
    "\n",
    "id_to_coords = {\n",
    "    V1: [{\"layer\": 1, \"start\": 0, \"end\": embedding_dim}], \n",
    "    V2: [{\"layer\": 1, \"start\": embedding_dim, \"end\": embedding_dim*2}]    \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "both_model = TorchDeepNeuralClassifierIIT(\n",
    "    hidden_dim=embedding_dim*4, \n",
    "    hidden_activation=torch.nn.ReLU(), \n",
    "    num_layers=3, \n",
    "    id_to_coords=id_to_coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "v1data = iit.get_IIT_equality_dataset(\"V1\", embedding_dim, data_size)\n",
    "\n",
    "v2data = iit.get_IIT_equality_dataset(\"V2\", embedding_dim, data_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stopping after epoch 715. Training loss did not improve more than tol=1e-05. Final error is 0.8061347478069365."
     ]
    }
   ],
   "source": [
    "X_base_train_both = torch.cat([v1data[0], v2data[0]], dim=0)\n",
    "\n",
    "X_sources_train_both = [torch.cat([v1data[1][i], v2data[1][i]], dim=0) \n",
    "                        for i in range(len(v1data[1]))] \n",
    "\n",
    "y_base_train_both = torch.cat([v1data[2], v2data[2]])\n",
    "y_IIT_train_both = torch.cat([v1data[3], v2data[3]])\n",
    "\n",
    "interventions_both = torch.cat([v1data[4], v2data[4]])\n",
    "\n",
    "_ = both_model.fit(\n",
    "    X_base_train_both, \n",
    "    X_sources_train_both, \n",
    "    y_base_train_both, \n",
    "    y_IIT_train_both, \n",
    "    interventions_both)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Standard evaluation\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       1.00      1.00      1.00        48\n",
      "           1       1.00      1.00      1.00        48\n",
      "\n",
      "    accuracy                           1.00        96\n",
      "   macro avg       1.00      1.00      1.00        96\n",
      "weighted avg       1.00      1.00      1.00        96\n",
      "\n",
      "V1 counterfactual evaluation\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       1.00      1.00      1.00        48\n",
      "           1       1.00      1.00      1.00        48\n",
      "\n",
      "    accuracy                           1.00        96\n",
      "   macro avg       1.00      1.00      1.00        96\n",
      "weighted avg       1.00      1.00      1.00        96\n",
      "\n"
     ]
    }
   ],
   "source": [
    "IIT_preds, base_preds = both_model.iit_predict(\n",
    "    X_base_test, X_sources_test, interventions)\n",
    "\n",
    "print(\"Standard evaluation\")\n",
    "print(classification_report(y_base_test, base_preds))\n",
    "      \n",
    "print(\"V1 counterfactual evaluation\")\n",
    "print(classification_report(y_IIT_test, IIT_preds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "V2 counterfactual evaluation\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       1.00      1.00      1.00      5000\n",
      "           1       1.00      1.00      1.00      5000\n",
      "\n",
      "    accuracy                           1.00     10000\n",
      "   macro avg       1.00      1.00      1.00     10000\n",
      "weighted avg       1.00      1.00      1.00     10000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "IIT_preds_V2, base_preds_V2 = both_model.iit_predict(\n",
    "    X_base_test_V2, X_sources_test_V2, interventions_V2)\n",
    "\n",
    "print(\"V2 counterfactual evaluation\")\n",
    "print(classification_report(y_IIT_test_V2, IIT_preds_V2))"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "933b0a94e0d88ac80a17cb26ca3d8d36930c12815b02a2885c1925c2b1ae3c33"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}