train_polysemantic_benchmarks.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from iit.model_pairs.strict_iit_model_pair import StrictIITModelPair\n",
    "\n",
    "from poly_bench.cases.paren_checker import HighLevelParensBalanceChecker, BalancedParensDataset\n",
    "from poly_bench.cases.left_greater import HighLevelLeftGreater, LeftGreaterDataset\n",
    "from poly_bench.cases.duplicate_remover import HighLevelDuplicateRemover, DuplicateRemoverDataset\n",
    "from poly_bench.cases.unique_extractor import HighLevelUniqueExtractor, UniqueExtractorDataset\n",
    "from poly_bench.poly_hl_model import PolyHLModel, PolyModelDataset\n",
    "from poly_bench.io import save_poly_model_to_dir, save_to_hf\n",
    "\n",
    "Case0 = HighLevelDuplicateRemover\n",
    "Case1 = HighLevelLeftGreater\n",
    "Case2 = HighLevelParensBalanceChecker\n",
    "Case3 = HighLevelUniqueExtractor\n",
    "\n",
    "dataset_mapping = {\n",
    "    HighLevelDuplicateRemover: DuplicateRemoverDataset,\n",
    "    HighLevelLeftGreater: LeftGreaterDataset,\n",
    "    HighLevelParensBalanceChecker: BalancedParensDataset,\n",
    "    HighLevelUniqueExtractor: UniqueExtractorDataset\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 100\n",
    "n_samples = 10_000\n",
    "training_args = {\n",
    "    \"batch_size\": 256,\n",
    "    \"num_workers\": 0,\n",
    "    \"use_single_loss\": True,\n",
    "    \"behavior_weight\": 0.4, #basically doubles the strict weight's job.\n",
    "    \"iit_weight\": 1.,\n",
    "    \"strict_weight\": 0.4,\n",
    "    \"clip_grad_norm\": 1.0,\n",
    "    \"early_stop\" : True,\n",
    "    \"lr_scheduler\": torch.optim.lr_scheduler.LinearLR,\n",
    "    \"scheduler_kwargs\": dict(start_factor=1, end_factor=0.2, total_iters=int(n_epochs)),\n",
    "    \"optimizer_kwargs\": dict(lr=1e-3, betas=(0.9, 0.9)),\n",
    "    \"scheduler_val_metric\": [\"val/accuracy\", \"val/IIA\"], #for ReduceLRonPlateau\n",
    "    \"scheduler_mode\": \"max\", #for ReduceLRonPlateau\n",
    "    \"siit_sampling\" : \"sample_all\",\n",
    "    \"val_IIA_sampling\": \"all\", # random or all\n",
    "    \"seed\" : 42\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_ctx = 15\n",
    "seed = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (0) DuplicateRemover + (1) LeftGreater"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[prev_equal_hook], [mlp0_hook]]\n",
      "blocks.0.attn.hook_z.0 [[prev_token_hook], None]\n",
      "blocks.0.attn.hook_z.1 [None, [paren_counts_hook]]\n",
      "blocks.0.attn.hook_z.2 [None, None]\n",
      "blocks.0.attn.hook_z.3 [None, None]\n",
      "blocks.1.mlp.hook_post [[output_hook], None]\n",
      "blocks.1.attn.hook_z.0 [None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case0, Case1]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:14: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.inputs = t.tensor(inputs).to(t.int)\n",
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:15: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.targets = t.tensor(targets).to(t.float32)\n",
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:19: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.markers = t.tensor(markers).to(t.int)\n"
     ]
    }
   ],
   "source": [
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "[LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c58a312f9da4abd9b5125ff4eb0e677",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d981c8fa56cf46d1aba916a5a4a72ac0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/54 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 1.1303, train/behavior_loss: 0.4246, train/strict_loss: 0.1767, val/iit_loss: 0.8554, val/IIA: 68.64%, val/accuracy: 76.78%, val/strict_accuracy: 75.68%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.6910, train/behavior_loss: 0.2209, train/strict_loss: 0.1072, val/iit_loss: 0.5357, val/IIA: 79.27%, val/accuracy: 85.36%, val/strict_accuracy: 84.45%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.4386, train/behavior_loss: 0.1223, train/strict_loss: 0.0726, val/iit_loss: 0.3746, val/IIA: 85.13%, val/accuracy: 91.02%, val/strict_accuracy: 89.75%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3199, train/behavior_loss: 0.0829, train/strict_loss: 0.0583, val/iit_loss: 0.3103, val/IIA: 87.58%, val/accuracy: 93.06%, val/strict_accuracy: 91.35%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2798, train/behavior_loss: 0.0636, train/strict_loss: 0.0494, val/iit_loss: 0.2529, val/IIA: 89.70%, val/accuracy: 94.13%, val/strict_accuracy: 92.21%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2106, train/behavior_loss: 0.0437, train/strict_loss: 0.0477, val/iit_loss: 0.1763, val/IIA: 93.30%, val/accuracy: 97.22%, val/strict_accuracy: 94.06%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1759, train/behavior_loss: 0.0239, train/strict_loss: 0.0500, val/iit_loss: 0.1457, val/IIA: 94.34%, val/accuracy: 98.80%, val/strict_accuracy: 95.96%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1255, train/behavior_loss: 0.0091, train/strict_loss: 0.0424, val/iit_loss: 0.1259, val/IIA: 95.69%, val/accuracy: 99.94%, val/strict_accuracy: 97.54%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1236, train/behavior_loss: 0.0069, train/strict_loss: 0.0347, val/iit_loss: 0.1007, val/IIA: 96.43%, val/accuracy: 99.92%, val/strict_accuracy: 97.75%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1073, train/behavior_loss: 0.0040, train/strict_loss: 0.0336, val/iit_loss: 0.0987, val/IIA: 96.44%, val/accuracy: 100.00%, val/strict_accuracy: 98.16%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1030, train/behavior_loss: 0.0034, train/strict_loss: 0.0273, val/iit_loss: 0.1197, val/IIA: 95.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.00%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1210, train/behavior_loss: 0.0046, train/strict_loss: 0.0284, val/iit_loss: 0.1009, val/IIA: 96.58%, val/accuracy: 99.93%, val/strict_accuracy: 97.89%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0936, train/behavior_loss: 0.0029, train/strict_loss: 0.0240, val/iit_loss: 0.1008, val/IIA: 96.38%, val/accuracy: 99.95%, val/strict_accuracy: 98.73%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0770, train/behavior_loss: 0.0023, train/strict_loss: 0.0242, val/iit_loss: 0.0772, val/IIA: 97.13%, val/accuracy: 100.00%, val/strict_accuracy: 98.97%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0752, train/behavior_loss: 0.0020, train/strict_loss: 0.0215, val/iit_loss: 0.0864, val/IIA: 96.31%, val/accuracy: 100.00%, val/strict_accuracy: 99.09%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0705, train/behavior_loss: 0.0011, train/strict_loss: 0.0139, val/iit_loss: 0.1230, val/IIA: 96.23%, val/accuracy: 100.00%, val/strict_accuracy: 99.73%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1034, train/behavior_loss: 0.0014, train/strict_loss: 0.0096, val/iit_loss: 0.0626, val/IIA: 97.80%, val/accuracy: 99.99%, val/strict_accuracy: 98.92%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0582, train/behavior_loss: 0.0011, train/strict_loss: 0.0172, val/iit_loss: 0.0513, val/IIA: 98.24%, val/accuracy: 99.98%, val/strict_accuracy: 99.30%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0551, train/behavior_loss: 0.0010, train/strict_loss: 0.0170, val/iit_loss: 0.0636, val/IIA: 97.71%, val/accuracy: 100.00%, val/strict_accuracy: 99.24%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0523, train/behavior_loss: 0.0008, train/strict_loss: 0.0103, val/iit_loss: 0.0482, val/IIA: 98.31%, val/accuracy: 100.00%, val/strict_accuracy: 99.54%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0694, train/behavior_loss: 0.0005, train/strict_loss: 0.0089, val/iit_loss: 0.0529, val/IIA: 97.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.53%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0552, train/behavior_loss: 0.0005, train/strict_loss: 0.0093, val/iit_loss: 0.0609, val/IIA: 97.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.72%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0484, train/behavior_loss: 0.0006, train/strict_loss: 0.0086, val/iit_loss: 0.0402, val/IIA: 98.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.57%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0501, train/behavior_loss: 0.0003, train/strict_loss: 0.0057, val/iit_loss: 0.0457, val/IIA: 98.31%, val/accuracy: 100.00%, val/strict_accuracy: 99.59%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0449, train/behavior_loss: 0.0004, train/strict_loss: 0.0058, val/iit_loss: 0.0459, val/IIA: 98.04%, val/accuracy: 99.98%, val/strict_accuracy: 99.69%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0578, train/behavior_loss: 0.0004, train/strict_loss: 0.0066, val/iit_loss: 0.0355, val/IIA: 98.71%, val/accuracy: 99.99%, val/strict_accuracy: 99.67%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0361, train/behavior_loss: 0.0003, train/strict_loss: 0.0041, val/iit_loss: 0.0359, val/IIA: 98.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0434, train/behavior_loss: 0.0002, train/strict_loss: 0.0043, val/iit_loss: 0.0333, val/IIA: 98.80%, val/accuracy: 99.97%, val/strict_accuracy: 99.56%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0373, train/behavior_loss: 0.0001, train/strict_loss: 0.0035, val/iit_loss: 0.0358, val/IIA: 98.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.83%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0365, train/behavior_loss: 0.0003, train/strict_loss: 0.0033, val/iit_loss: 0.0468, val/IIA: 98.42%, val/accuracy: 100.00%, val/strict_accuracy: 99.88%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0420, train/behavior_loss: 0.0002, train/strict_loss: 0.0036, val/iit_loss: 0.0282, val/IIA: 98.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.58%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0388, train/behavior_loss: 0.0001, train/strict_loss: 0.0046, val/iit_loss: 0.0823, val/IIA: 97.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.40%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0338, train/behavior_loss: 0.0001, train/strict_loss: 0.0043, val/iit_loss: 0.0573, val/IIA: 98.16%, val/accuracy: 100.00%, val/strict_accuracy: 99.77%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0413, train/behavior_loss: 0.0001, train/strict_loss: 0.0021, val/iit_loss: 0.0607, val/IIA: 98.09%, val/accuracy: 100.00%, val/strict_accuracy: 99.84%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0383, train/behavior_loss: 0.0001, train/strict_loss: 0.0029, val/iit_loss: 0.0303, val/IIA: 98.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.82%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0637, train/behavior_loss: 0.0001, train/strict_loss: 0.0037, val/iit_loss: 0.0456, val/IIA: 98.32%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0475, train/behavior_loss: 0.0001, train/strict_loss: 0.0033, val/iit_loss: 0.0301, val/IIA: 98.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0361, train/behavior_loss: 0.0001, train/strict_loss: 0.0026, val/iit_loss: 0.0265, val/IIA: 98.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.83%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0356, train/behavior_loss: 0.0001, train/strict_loss: 0.0022, val/iit_loss: 0.0368, val/IIA: 98.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0514, train/behavior_loss: 0.0002, train/strict_loss: 0.0030, val/iit_loss: 0.0446, val/IIA: 98.14%, val/accuracy: 99.76%, val/strict_accuracy: 99.35%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0369, train/behavior_loss: 0.0002, train/strict_loss: 0.0033, val/iit_loss: 0.0276, val/IIA: 98.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0276, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0593, val/IIA: 98.07%, val/accuracy: 99.99%, val/strict_accuracy: 99.73%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0695, train/behavior_loss: 0.0001, train/strict_loss: 0.0023, val/iit_loss: 0.0467, val/IIA: 98.16%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0444, train/behavior_loss: 0.0001, train/strict_loss: 0.0030, val/iit_loss: 0.0469, val/IIA: 98.29%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0527, train/behavior_loss: 0.0001, train/strict_loss: 0.0025, val/iit_loss: 0.0324, val/IIA: 98.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0332, train/behavior_loss: 0.0000, train/strict_loss: 0.0021, val/iit_loss: 0.0239, val/IIA: 99.09%, val/accuracy: 100.00%, val/strict_accuracy: 99.88%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0277, train/behavior_loss: 0.0000, train/strict_loss: 0.0011, val/iit_loss: 0.0254, val/IIA: 99.08%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0368, train/behavior_loss: 0.0000, train/strict_loss: 0.0021, val/iit_loss: 0.0246, val/IIA: 99.04%, val/accuracy: 100.00%, val/strict_accuracy: 99.79%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0366, train/behavior_loss: 0.0000, train/strict_loss: 0.0017, val/iit_loss: 0.0303, val/IIA: 98.83%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0355, train/behavior_loss: 0.0000, train/strict_loss: 0.0014, val/iit_loss: 0.0331, val/IIA: 98.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0349, train/behavior_loss: 0.0000, train/strict_loss: 0.0017, val/iit_loss: 0.0251, val/IIA: 99.05%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0376, train/behavior_loss: 0.0001, train/strict_loss: 0.0018, val/iit_loss: 0.0266, val/IIA: 99.01%, val/accuracy: 100.00%, val/strict_accuracy: 99.88%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0318, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0529, val/IIA: 98.00%, val/accuracy: 100.00%, val/strict_accuracy: 99.80%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0393, train/behavior_loss: 0.0001, train/strict_loss: 0.0015, val/iit_loss: 0.0255, val/IIA: 98.97%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0335, train/behavior_loss: 0.0001, train/strict_loss: 0.0018, val/iit_loss: 0.0241, val/IIA: 99.02%, val/accuracy: 100.00%, val/strict_accuracy: 99.75%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0253, train/behavior_loss: 0.0001, train/strict_loss: 0.0021, val/iit_loss: 0.0308, val/IIA: 98.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.89%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0256, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0234, val/IIA: 99.07%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0314, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0206, val/IIA: 99.21%, val/accuracy: 100.00%, val/strict_accuracy: 99.88%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0342, train/behavior_loss: 0.0000, train/strict_loss: 0.0016, val/iit_loss: 0.0229, val/IIA: 99.14%, val/accuracy: 100.00%, val/strict_accuracy: 99.83%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0330, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0399, val/IIA: 98.67%, val/accuracy: 99.83%, val/strict_accuracy: 99.81%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0525, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0538, val/IIA: 97.99%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0366, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0175, val/IIA: 99.37%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0228, train/behavior_loss: 0.0000, train/strict_loss: 0.0011, val/iit_loss: 0.0292, val/IIA: 99.00%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0260, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0403, val/IIA: 98.39%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0273, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0313, val/IIA: 98.72%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0286, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0257, val/IIA: 99.02%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0251, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0220, val/IIA: 99.19%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0292, train/behavior_loss: 0.0000, train/strict_loss: 0.0011, val/iit_loss: 0.0245, val/IIA: 98.99%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0281, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0261, val/IIA: 99.04%, val/accuracy: 100.00%, val/strict_accuracy: 99.67%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0324, train/behavior_loss: 0.0000, train/strict_loss: 0.0018, val/iit_loss: 0.0224, val/IIA: 99.13%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0213, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0215, val/IIA: 99.18%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0201, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0174, val/IIA: 99.32%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0177, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0156, val/IIA: 99.43%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0226, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0227, val/IIA: 99.14%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0179, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0162, val/IIA: 99.35%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0211, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0311, val/IIA: 98.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.81%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0269, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0200, val/IIA: 99.23%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0204, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0203, val/IIA: 99.12%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0193, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0195, val/IIA: 99.22%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0234, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0266, val/IIA: 99.02%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0190, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0144, val/IIA: 99.47%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0500, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.1058, val/IIA: 97.35%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0291, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0151, val/IIA: 99.44%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0139, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0160, val/IIA: 99.40%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0267, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0273, val/IIA: 98.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0219, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0246, val/IIA: 99.12%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0199, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0318, val/IIA: 98.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0184, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0273, val/IIA: 98.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0190, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0168, val/IIA: 99.27%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0215, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0170, val/IIA: 99.33%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0118, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0128, val/IIA: 99.51%, val/accuracy: 100.00%, val/strict_accuracy: 99.82%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0166, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0123, val/IIA: 99.57%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0145, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0129, val/IIA: 99.48%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0148, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0128, val/IIA: 99.55%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0158, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0188, val/IIA: 99.34%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0110, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0113, val/IIA: 99.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0121, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0130, val/IIA: 99.55%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0164, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0126, val/IIA: 99.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0132, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0117, val/IIA: 99.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.80%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0107, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0156, val/IIA: 99.34%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_0+1')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_0+1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (0) DuplicateRemover + (2) ParenChecker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[prev_equal_hook], [mlp0_hook]]\n",
      "blocks.0.attn.hook_z.0 [[prev_token_hook], None]\n",
      "blocks.0.attn.hook_z.1 [None, None]\n",
      "blocks.0.attn.hook_z.2 [None, None]\n",
      "blocks.0.attn.hook_z.3 [None, [paren_counts_hook]]\n",
      "blocks.1.mlp.hook_post [[output_hook], [mlp1_hook]]\n",
      "blocks.1.attn.hook_z.0 [None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None]\n",
      "blocks.2.mlp.hook_post [None, [mlp2_hook]]\n",
      "blocks.2.attn.hook_z.0 [None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None]\n",
      "blocks.2.attn.hook_z.3 [None, [horizon_lookback_hook]]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case0, Case2]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:14: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.inputs = t.tensor(inputs).to(t.int)\n",
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:15: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.targets = t.tensor(targets).to(t.float32)\n",
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:19: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.markers = t.tensor(markers).to(t.int)\n"
     ]
    }
   ],
   "source": [
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "[LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "585a7213e5e6493889537580f20bcd70",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b56ce169ee1341d8b50eed11486fb85e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/54 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 1.0867, train/behavior_loss: 0.4002, train/strict_loss: 0.1730, val/iit_loss: 0.7787, val/IIA: 76.15%, val/accuracy: 81.04%, val/strict_accuracy: 80.95%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.6278, train/behavior_loss: 0.2130, train/strict_loss: 0.1016, val/iit_loss: 0.4354, val/IIA: 86.25%, val/accuracy: 89.27%, val/strict_accuracy: 88.56%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3528, train/behavior_loss: 0.1151, train/strict_loss: 0.0593, val/iit_loss: 0.2729, val/IIA: 90.48%, val/accuracy: 93.35%, val/strict_accuracy: 92.29%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2344, train/behavior_loss: 0.0627, train/strict_loss: 0.0458, val/iit_loss: 0.1807, val/IIA: 93.68%, val/accuracy: 96.56%, val/strict_accuracy: 94.75%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1731, train/behavior_loss: 0.0318, train/strict_loss: 0.0452, val/iit_loss: 0.1211, val/IIA: 96.13%, val/accuracy: 98.80%, val/strict_accuracy: 96.68%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0885, train/behavior_loss: 0.0143, train/strict_loss: 0.0467, val/iit_loss: 0.0857, val/IIA: 97.32%, val/accuracy: 99.46%, val/strict_accuracy: 97.80%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0608, train/behavior_loss: 0.0092, train/strict_loss: 0.0356, val/iit_loss: 0.0749, val/IIA: 97.69%, val/accuracy: 99.41%, val/strict_accuracy: 98.13%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0548, train/behavior_loss: 0.0064, train/strict_loss: 0.0247, val/iit_loss: 0.0485, val/IIA: 98.49%, val/accuracy: 99.69%, val/strict_accuracy: 98.89%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0522, train/behavior_loss: 0.0060, train/strict_loss: 0.0208, val/iit_loss: 0.0464, val/IIA: 98.54%, val/accuracy: 99.71%, val/strict_accuracy: 99.28%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0385, train/behavior_loss: 0.0044, train/strict_loss: 0.0141, val/iit_loss: 0.0410, val/IIA: 98.56%, val/accuracy: 99.62%, val/strict_accuracy: 99.35%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0367, train/behavior_loss: 0.0039, train/strict_loss: 0.0089, val/iit_loss: 0.0442, val/IIA: 98.43%, val/accuracy: 99.32%, val/strict_accuracy: 99.27%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0452, train/behavior_loss: 0.0037, train/strict_loss: 0.0064, val/iit_loss: 0.0329, val/IIA: 98.84%, val/accuracy: 99.81%, val/strict_accuracy: 99.70%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0317, train/behavior_loss: 0.0024, train/strict_loss: 0.0050, val/iit_loss: 0.0265, val/IIA: 99.12%, val/accuracy: 99.91%, val/strict_accuracy: 99.78%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0437, train/behavior_loss: 0.0021, train/strict_loss: 0.0043, val/iit_loss: 0.0367, val/IIA: 98.73%, val/accuracy: 99.66%, val/strict_accuracy: 99.53%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0277, train/behavior_loss: 0.0021, train/strict_loss: 0.0052, val/iit_loss: 0.0263, val/IIA: 99.03%, val/accuracy: 99.92%, val/strict_accuracy: 99.75%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0241, train/behavior_loss: 0.0014, train/strict_loss: 0.0038, val/iit_loss: 0.0199, val/IIA: 99.35%, val/accuracy: 99.97%, val/strict_accuracy: 99.84%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0210, train/behavior_loss: 0.0011, train/strict_loss: 0.0037, val/iit_loss: 0.0142, val/IIA: 99.53%, val/accuracy: 99.95%, val/strict_accuracy: 99.84%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0155, train/behavior_loss: 0.0009, train/strict_loss: 0.0041, val/iit_loss: 0.0182, val/IIA: 99.37%, val/accuracy: 99.98%, val/strict_accuracy: 99.74%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0150, train/behavior_loss: 0.0005, train/strict_loss: 0.0033, val/iit_loss: 0.0180, val/IIA: 99.41%, val/accuracy: 99.93%, val/strict_accuracy: 99.77%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0189, train/behavior_loss: 0.0005, train/strict_loss: 0.0041, val/iit_loss: 0.0109, val/IIA: 99.66%, val/accuracy: 99.99%, val/strict_accuracy: 99.88%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0140, train/behavior_loss: 0.0004, train/strict_loss: 0.0030, val/iit_loss: 0.0094, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0127, train/behavior_loss: 0.0004, train/strict_loss: 0.0030, val/iit_loss: 0.0146, val/IIA: 99.52%, val/accuracy: 99.88%, val/strict_accuracy: 99.67%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0135, train/behavior_loss: 0.0004, train/strict_loss: 0.0030, val/iit_loss: 0.0114, val/IIA: 99.61%, val/accuracy: 99.98%, val/strict_accuracy: 99.88%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0119, train/behavior_loss: 0.0003, train/strict_loss: 0.0029, val/iit_loss: 0.0115, val/IIA: 99.61%, val/accuracy: 99.94%, val/strict_accuracy: 99.84%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0109, train/behavior_loss: 0.0003, train/strict_loss: 0.0026, val/iit_loss: 0.0065, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0063, train/behavior_loss: 0.0002, train/strict_loss: 0.0024, val/iit_loss: 0.0081, val/IIA: 99.74%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0078, train/behavior_loss: 0.0002, train/strict_loss: 0.0021, val/iit_loss: 0.0071, val/IIA: 99.80%, val/accuracy: 99.98%, val/strict_accuracy: 99.90%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0085, train/behavior_loss: 0.0001, train/strict_loss: 0.0019, val/iit_loss: 0.0087, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0069, train/behavior_loss: 0.0002, train/strict_loss: 0.0019, val/iit_loss: 0.0095, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0091, train/behavior_loss: 0.0002, train/strict_loss: 0.0017, val/iit_loss: 0.0084, val/IIA: 99.69%, val/accuracy: 99.98%, val/strict_accuracy: 99.92%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0002, train/strict_loss: 0.0017, val/iit_loss: 0.0069, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0082, train/behavior_loss: 0.0001, train/strict_loss: 0.0018, val/iit_loss: 0.0192, val/IIA: 99.49%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0076, train/behavior_loss: 0.0001, train/strict_loss: 0.0018, val/iit_loss: 0.0057, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0121, val/IIA: 99.64%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0066, train/behavior_loss: 0.0001, train/strict_loss: 0.0015, val/iit_loss: 0.0055, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0081, val/IIA: 99.74%, val/accuracy: 99.98%, val/strict_accuracy: 99.89%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0001, train/strict_loss: 0.0016, val/iit_loss: 0.0040, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0060, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0048, val/IIA: 99.85%, val/accuracy: 99.99%, val/strict_accuracy: 99.90%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0053, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0000, train/strict_loss: 0.0011, val/iit_loss: 0.0038, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0046, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0046, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.89%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0052, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0054, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0058, val/IIA: 99.82%, val/accuracy: 99.99%, val/strict_accuracy: 99.91%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0033, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0037, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0065, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0042, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0047, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0036, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0045, val/IIA: 99.84%, val/accuracy: 99.98%, val/strict_accuracy: 99.86%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0036, val/IIA: 99.88%, val/accuracy: 99.96%, val/strict_accuracy: 99.82%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0048, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0033, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0074, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0036, val/IIA: 99.88%, val/accuracy: 99.97%, val/strict_accuracy: 99.93%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0031, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0031, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0039, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0033, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0050, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0027, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0029, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0031, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0026, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0037, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0028, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0031, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0028, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0020, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0032, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0029, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0029, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0027, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0024, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0025, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0025, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0029, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0025, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0022, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0035, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0031, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0027, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0014, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0028, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0025, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0067, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0033, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0024, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0023, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0023, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0016, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0022, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0025, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0027, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0021, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0021, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0024, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0022, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0023, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0022, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0029, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0021, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0020, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0022, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0014, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0022, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0022, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0022, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0017, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0021, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0018, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0022, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0021, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_0+2')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_0+2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (0) DuplicateRemover + (3) UniqueExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[prev_equal_hook], [appeared_mlp]]\n",
      "blocks.0.attn.hook_z.0 [[prev_token_hook], None]\n",
      "blocks.0.attn.hook_z.1 [None, None]\n",
      "blocks.0.attn.hook_z.2 [None, [counter_head]]\n",
      "blocks.0.attn.hook_z.3 [None, None]\n",
      "blocks.1.mlp.hook_post [[output_hook], [mask_mlp]]\n",
      "blocks.1.attn.hook_z.0 [None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None]\n",
      "blocks.2.mlp.hook_post [None, [output_mlp]]\n",
      "blocks.2.attn.hook_z.0 [None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None]\n",
      "blocks.2.attn.hook_z.3 [None, None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case0, Case3]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "[LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6afba2badd56484f86a2fc348e98320a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7927e6ca7cca4bc0b47856db7cb71ab1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/63 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.7631, train/behavior_loss: 0.2748, train/strict_loss: 0.1209, val/iit_loss: 0.4155, val/IIA: 86.74%, val/accuracy: 90.95%, val/strict_accuracy: 91.26%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3254, train/behavior_loss: 0.0832, train/strict_loss: 0.0493, val/iit_loss: 0.2580, val/IIA: 89.51%, val/accuracy: 95.46%, val/strict_accuracy: 94.97%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2232, train/behavior_loss: 0.0387, train/strict_loss: 0.0358, val/iit_loss: 0.1836, val/IIA: 92.62%, val/accuracy: 98.21%, val/strict_accuracy: 97.30%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1380, train/behavior_loss: 0.0111, train/strict_loss: 0.0315, val/iit_loss: 0.1416, val/IIA: 94.67%, val/accuracy: 99.84%, val/strict_accuracy: 99.01%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0931, train/behavior_loss: 0.0044, train/strict_loss: 0.0258, val/iit_loss: 0.1074, val/IIA: 96.01%, val/accuracy: 99.93%, val/strict_accuracy: 99.50%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1354, train/behavior_loss: 0.0028, train/strict_loss: 0.0148, val/iit_loss: 0.0855, val/IIA: 96.81%, val/accuracy: 99.97%, val/strict_accuracy: 99.65%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0926, train/behavior_loss: 0.0019, train/strict_loss: 0.0115, val/iit_loss: 0.1126, val/IIA: 96.13%, val/accuracy: 100.00%, val/strict_accuracy: 99.78%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1022, train/behavior_loss: 0.0013, train/strict_loss: 0.0083, val/iit_loss: 0.0659, val/IIA: 97.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0955, train/behavior_loss: 0.0011, train/strict_loss: 0.0073, val/iit_loss: 0.0808, val/IIA: 97.20%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1042, train/behavior_loss: 0.0008, train/strict_loss: 0.0063, val/iit_loss: 0.0594, val/IIA: 97.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0663, train/behavior_loss: 0.0006, train/strict_loss: 0.0060, val/iit_loss: 0.0677, val/IIA: 97.60%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0719, train/behavior_loss: 0.0005, train/strict_loss: 0.0045, val/iit_loss: 0.0499, val/IIA: 98.15%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0677, train/behavior_loss: 0.0006, train/strict_loss: 0.0057, val/iit_loss: 0.0561, val/IIA: 97.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0698, train/behavior_loss: 0.0007, train/strict_loss: 0.0042, val/iit_loss: 0.0917, val/IIA: 97.39%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0557, train/behavior_loss: 0.0003, train/strict_loss: 0.0036, val/iit_loss: 0.0614, val/IIA: 97.57%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0435, train/behavior_loss: 0.0003, train/strict_loss: 0.0045, val/iit_loss: 0.0429, val/IIA: 98.20%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0419, train/behavior_loss: 0.0002, train/strict_loss: 0.0044, val/iit_loss: 0.0351, val/IIA: 98.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0508, train/behavior_loss: 0.0002, train/strict_loss: 0.0037, val/iit_loss: 0.0469, val/IIA: 98.13%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0266, train/behavior_loss: 0.0002, train/strict_loss: 0.0043, val/iit_loss: 0.0425, val/IIA: 98.38%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0478, train/behavior_loss: 0.0002, train/strict_loss: 0.0034, val/iit_loss: 0.0230, val/IIA: 99.29%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0237, train/behavior_loss: 0.0002, train/strict_loss: 0.0034, val/iit_loss: 0.0244, val/IIA: 99.16%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0330, train/behavior_loss: 0.0002, train/strict_loss: 0.0033, val/iit_loss: 0.0280, val/IIA: 98.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0280, train/behavior_loss: 0.0001, train/strict_loss: 0.0031, val/iit_loss: 0.0242, val/IIA: 99.13%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0162, train/behavior_loss: 0.0001, train/strict_loss: 0.0033, val/iit_loss: 0.0526, val/IIA: 98.23%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0368, train/behavior_loss: 0.0001, train/strict_loss: 0.0025, val/iit_loss: 0.0320, val/IIA: 98.73%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0268, train/behavior_loss: 0.0001, train/strict_loss: 0.0018, val/iit_loss: 0.0600, val/IIA: 98.03%, val/accuracy: 100.00%, val/strict_accuracy: 99.82%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0240, train/behavior_loss: 0.0001, train/strict_loss: 0.0026, val/iit_loss: 0.0121, val/IIA: 99.58%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0228, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0662, val/IIA: 98.10%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0195, train/behavior_loss: 0.0001, train/strict_loss: 0.0021, val/iit_loss: 0.0141, val/IIA: 99.49%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0305, train/behavior_loss: 0.0001, train/strict_loss: 0.0020, val/iit_loss: 0.0147, val/IIA: 99.48%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0163, train/behavior_loss: 0.0001, train/strict_loss: 0.0024, val/iit_loss: 0.0319, val/IIA: 98.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0247, train/behavior_loss: 0.0001, train/strict_loss: 0.0019, val/iit_loss: 0.0163, val/IIA: 99.39%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0241, train/behavior_loss: 0.0001, train/strict_loss: 0.0021, val/iit_loss: 0.0211, val/IIA: 99.33%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0179, train/behavior_loss: 0.0000, train/strict_loss: 0.0015, val/iit_loss: 0.0197, val/IIA: 99.33%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0226, train/behavior_loss: 0.0000, train/strict_loss: 0.0018, val/iit_loss: 0.0183, val/IIA: 99.34%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0179, train/behavior_loss: 0.0000, train/strict_loss: 0.0016, val/iit_loss: 0.0144, val/IIA: 99.50%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0332, train/behavior_loss: 0.0000, train/strict_loss: 0.0017, val/iit_loss: 0.0266, val/IIA: 99.24%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0140, train/behavior_loss: 0.0000, train/strict_loss: 0.0014, val/iit_loss: 0.0196, val/IIA: 99.30%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0158, train/behavior_loss: 0.0000, train/strict_loss: 0.0020, val/iit_loss: 0.0161, val/IIA: 99.43%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0176, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0161, val/IIA: 99.37%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0205, train/behavior_loss: 0.0000, train/strict_loss: 0.0015, val/iit_loss: 0.0178, val/IIA: 99.28%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0205, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0249, val/IIA: 99.03%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0258, train/behavior_loss: 0.0001, train/strict_loss: 0.0015, val/iit_loss: 0.0135, val/IIA: 99.55%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0115, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0247, val/IIA: 99.12%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0187, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0123, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0172, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0162, val/IIA: 99.38%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0232, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0256, val/IIA: 99.09%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0186, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0164, val/IIA: 99.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0089, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0099, val/IIA: 99.66%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0092, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0106, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0181, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0093, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0136, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0116, val/IIA: 99.58%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0180, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0106, val/IIA: 99.66%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0173, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0151, val/IIA: 99.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0118, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0189, val/IIA: 99.35%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0163, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0112, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0203, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0137, val/IIA: 99.40%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0116, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0125, val/IIA: 99.59%, val/accuracy: 99.99%, val/strict_accuracy: 99.91%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0104, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0087, val/IIA: 99.65%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0131, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0102, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0103, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0102, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0118, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0083, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0118, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0081, val/IIA: 99.72%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0085, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0198, val/IIA: 99.31%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0104, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0091, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0121, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0292, val/IIA: 99.16%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0081, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0090, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0131, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0167, val/IIA: 99.39%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0130, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0134, val/IIA: 99.49%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0081, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0080, val/IIA: 99.72%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0122, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0108, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0085, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0085, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0112, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0069, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0085, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0098, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0082, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0063, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0067, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0091, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0103, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0054, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0069, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0061, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0122, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0106, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0081, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0116, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0068, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0068, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0063, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0099, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0066, val/IIA: 99.74%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0064, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0100, val/IIA: 99.65%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0073, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0061, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0053, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0077, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0054, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0054, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0052, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0058, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0067, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0048, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0059, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0055, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0049, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0049, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0052, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0053, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0036, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0062, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_0+3')"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_0+3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (1) LeftGreater + (2) ParenChecker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[mlp0_hook], [mlp0_hook]]\n",
      "blocks.0.attn.hook_z.0 [None, None]\n",
      "blocks.0.attn.hook_z.1 [[paren_counts_hook], None]\n",
      "blocks.0.attn.hook_z.2 [None, None]\n",
      "blocks.0.attn.hook_z.3 [None, [paren_counts_hook]]\n",
      "blocks.1.mlp.hook_post [None, [mlp1_hook]]\n",
      "blocks.1.attn.hook_z.0 [None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None]\n",
      "blocks.2.mlp.hook_post [None, [mlp2_hook]]\n",
      "blocks.2.attn.hook_z.0 [None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None]\n",
      "blocks.2.attn.hook_z.3 [None, [horizon_lookback_hook]]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case1, Case2]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "[LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c13bfdb382cb4524bb1785f7e05a3936",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d0c12329d5424f61bca99928d0880fe9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/44 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.7672, train/behavior_loss: 0.2989, train/strict_loss: 0.1218, val/iit_loss: 0.4668, val/IIA: 91.08%, val/accuracy: 95.88%, val/strict_accuracy: 95.39%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3462, train/behavior_loss: 0.1083, train/strict_loss: 0.0493, val/iit_loss: 0.2613, val/IIA: 94.42%, val/accuracy: 97.74%, val/strict_accuracy: 97.59%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1935, train/behavior_loss: 0.0568, train/strict_loss: 0.0245, val/iit_loss: 0.1835, val/IIA: 94.67%, val/accuracy: 97.90%, val/strict_accuracy: 97.81%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1577, train/behavior_loss: 0.0314, train/strict_loss: 0.0166, val/iit_loss: 0.1350, val/IIA: 95.67%, val/accuracy: 98.73%, val/strict_accuracy: 98.48%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1116, train/behavior_loss: 0.0209, train/strict_loss: 0.0132, val/iit_loss: 0.1356, val/IIA: 95.41%, val/accuracy: 99.01%, val/strict_accuracy: 98.70%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1335, train/behavior_loss: 0.0157, train/strict_loss: 0.0099, val/iit_loss: 0.1040, val/IIA: 96.37%, val/accuracy: 98.80%, val/strict_accuracy: 98.56%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1089, train/behavior_loss: 0.0120, train/strict_loss: 0.0112, val/iit_loss: 0.0864, val/IIA: 96.61%, val/accuracy: 99.22%, val/strict_accuracy: 98.82%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0951, train/behavior_loss: 0.0092, train/strict_loss: 0.0099, val/iit_loss: 0.0842, val/IIA: 97.28%, val/accuracy: 99.54%, val/strict_accuracy: 99.18%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0742, train/behavior_loss: 0.0057, train/strict_loss: 0.0070, val/iit_loss: 0.0801, val/IIA: 97.36%, val/accuracy: 99.90%, val/strict_accuracy: 99.69%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0753, train/behavior_loss: 0.0051, train/strict_loss: 0.0049, val/iit_loss: 0.0709, val/IIA: 97.33%, val/accuracy: 99.86%, val/strict_accuracy: 99.68%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0605, train/behavior_loss: 0.0037, train/strict_loss: 0.0061, val/iit_loss: 0.0580, val/IIA: 97.89%, val/accuracy: 99.84%, val/strict_accuracy: 99.61%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0600, train/behavior_loss: 0.0039, train/strict_loss: 0.0068, val/iit_loss: 0.0408, val/IIA: 98.71%, val/accuracy: 99.91%, val/strict_accuracy: 99.53%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0407, train/behavior_loss: 0.0024, train/strict_loss: 0.0064, val/iit_loss: 0.0313, val/IIA: 99.01%, val/accuracy: 99.96%, val/strict_accuracy: 99.56%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0362, train/behavior_loss: 0.0017, train/strict_loss: 0.0054, val/iit_loss: 0.0433, val/IIA: 98.74%, val/accuracy: 99.97%, val/strict_accuracy: 99.70%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0336, train/behavior_loss: 0.0015, train/strict_loss: 0.0054, val/iit_loss: 0.0184, val/IIA: 99.43%, val/accuracy: 99.95%, val/strict_accuracy: 99.53%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0212, train/behavior_loss: 0.0009, train/strict_loss: 0.0051, val/iit_loss: 0.0202, val/IIA: 99.35%, val/accuracy: 100.00%, val/strict_accuracy: 99.67%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0176, train/behavior_loss: 0.0004, train/strict_loss: 0.0045, val/iit_loss: 0.0118, val/IIA: 99.64%, val/accuracy: 100.00%, val/strict_accuracy: 99.68%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0125, train/behavior_loss: 0.0003, train/strict_loss: 0.0041, val/iit_loss: 0.0117, val/IIA: 99.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.65%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0124, train/behavior_loss: 0.0002, train/strict_loss: 0.0038, val/iit_loss: 0.0130, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.83%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0115, train/behavior_loss: 0.0002, train/strict_loss: 0.0036, val/iit_loss: 0.0110, val/IIA: 99.57%, val/accuracy: 100.00%, val/strict_accuracy: 99.68%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0159, train/behavior_loss: 0.0002, train/strict_loss: 0.0037, val/iit_loss: 0.0088, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.79%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0086, train/behavior_loss: 0.0001, train/strict_loss: 0.0032, val/iit_loss: 0.0065, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.88%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0086, train/behavior_loss: 0.0001, train/strict_loss: 0.0035, val/iit_loss: 0.0056, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.77%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0072, train/behavior_loss: 0.0001, train/strict_loss: 0.0028, val/iit_loss: 0.0071, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 99.77%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0093, train/behavior_loss: 0.0001, train/strict_loss: 0.0025, val/iit_loss: 0.0071, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0081, train/behavior_loss: 0.0001, train/strict_loss: 0.0021, val/iit_loss: 0.0057, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0019, val/iit_loss: 0.0039, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0016, val/iit_loss: 0.0038, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0014, val/iit_loss: 0.0068, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0036, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0052, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0025, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0029, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0023, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0020, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0016, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0020, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0023, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0014, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0021, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0032, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0043, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0013, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0014, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0026, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0013, val/IIA: 99.97%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0025, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0028, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0020, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0025, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0017, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0018, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0017, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0011, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0013, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0017, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0010, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0011, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0018, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0012, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0011, val/IIA: 99.97%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0098, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0019, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0012, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0009, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0012, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0014, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0008, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0012, val/IIA: 99.97%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0014, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0010, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0012, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0010, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0009, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0009, val/IIA: 99.97%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0027, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0006, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0007, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0007, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0010, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0008, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0004, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0007, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0008, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0012, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0014, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0013, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0006, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0005, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0018, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0010, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0038, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0010, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0017, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0016, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0008, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0007, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0019, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0012, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0007, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0007, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0017, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0006, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0005, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0004, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0005, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0005, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0004, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0008, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0010, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0007, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0006, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0004, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0009, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0007, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0015, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0006, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0003, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0005, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0006, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0007, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0004, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0004, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0002, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0005, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0003, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0004, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0003, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0005, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0001, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0006, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0004, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0003, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0002, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0007, val/IIA: 99.97%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0005, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0005, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0002, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0003, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0003, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0003, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0001, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0004, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0004, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0004, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0005, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0003, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0003, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0003, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0002, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0002, val/IIA: 100.00%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_1+2')"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_1+2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (1) LeftGreater + (3) UniqueExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[mlp0_hook], [appeared_mlp]]\n",
      "blocks.0.attn.hook_z.0 [None, None]\n",
      "blocks.0.attn.hook_z.1 [[paren_counts_hook], None]\n",
      "blocks.0.attn.hook_z.2 [None, [counter_head]]\n",
      "blocks.0.attn.hook_z.3 [None, None]\n",
      "blocks.1.mlp.hook_post [None, [mask_mlp]]\n",
      "blocks.1.attn.hook_z.0 [None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None]\n",
      "blocks.2.mlp.hook_post [None, [output_mlp]]\n",
      "blocks.2.attn.hook_z.0 [None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None]\n",
      "blocks.2.attn.hook_z.3 [None, None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case1, Case3]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "[LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "977be68efed64fc58406e2941bdb0360",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8be28973c8394c25a01c06507bc57ae9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/54 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.7667, train/behavior_loss: 0.2909, train/strict_loss: 0.1265, val/iit_loss: 0.4739, val/IIA: 86.81%, val/accuracy: 92.44%, val/strict_accuracy: 91.41%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3308, train/behavior_loss: 0.0943, train/strict_loss: 0.0470, val/iit_loss: 0.2141, val/IIA: 93.61%, val/accuracy: 98.58%, val/strict_accuracy: 98.21%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1541, train/behavior_loss: 0.0219, train/strict_loss: 0.0185, val/iit_loss: 0.1194, val/IIA: 96.25%, val/accuracy: 99.64%, val/strict_accuracy: 99.25%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1118, train/behavior_loss: 0.0097, train/strict_loss: 0.0167, val/iit_loss: 0.0813, val/IIA: 97.33%, val/accuracy: 99.95%, val/strict_accuracy: 99.62%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0615, train/behavior_loss: 0.0037, train/strict_loss: 0.0106, val/iit_loss: 0.0859, val/IIA: 96.94%, val/accuracy: 99.95%, val/strict_accuracy: 99.74%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0632, train/behavior_loss: 0.0026, train/strict_loss: 0.0107, val/iit_loss: 0.0772, val/IIA: 96.93%, val/accuracy: 99.99%, val/strict_accuracy: 99.85%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0593, train/behavior_loss: 0.0019, train/strict_loss: 0.0068, val/iit_loss: 0.0637, val/IIA: 97.94%, val/accuracy: 99.99%, val/strict_accuracy: 99.73%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0448, train/behavior_loss: 0.0011, train/strict_loss: 0.0056, val/iit_loss: 0.0363, val/IIA: 98.75%, val/accuracy: 99.86%, val/strict_accuracy: 99.68%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0252, train/behavior_loss: 0.0006, train/strict_loss: 0.0055, val/iit_loss: 0.0294, val/IIA: 99.13%, val/accuracy: 99.99%, val/strict_accuracy: 99.77%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0396, train/behavior_loss: 0.0006, train/strict_loss: 0.0060, val/iit_loss: 0.0205, val/IIA: 99.22%, val/accuracy: 99.99%, val/strict_accuracy: 99.73%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0278, train/behavior_loss: 0.0004, train/strict_loss: 0.0048, val/iit_loss: 0.0214, val/IIA: 99.29%, val/accuracy: 100.00%, val/strict_accuracy: 99.80%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0218, train/behavior_loss: 0.0004, train/strict_loss: 0.0042, val/iit_loss: 0.0172, val/IIA: 99.33%, val/accuracy: 99.99%, val/strict_accuracy: 99.83%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0184, train/behavior_loss: 0.0003, train/strict_loss: 0.0042, val/iit_loss: 0.0316, val/IIA: 98.99%, val/accuracy: 99.99%, val/strict_accuracy: 99.90%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0227, train/behavior_loss: 0.0002, train/strict_loss: 0.0044, val/iit_loss: 0.0136, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0100, train/behavior_loss: 0.0001, train/strict_loss: 0.0022, val/iit_loss: 0.0093, val/IIA: 99.64%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0175, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0219, val/IIA: 99.35%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0163, train/behavior_loss: 0.0001, train/strict_loss: 0.0016, val/iit_loss: 0.0107, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.82%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0226, train/behavior_loss: 0.0001, train/strict_loss: 0.0018, val/iit_loss: 0.0104, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.83%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0228, train/behavior_loss: 0.0001, train/strict_loss: 0.0021, val/iit_loss: 0.0662, val/IIA: 98.02%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0248, train/behavior_loss: 0.0001, train/strict_loss: 0.0020, val/iit_loss: 0.0058, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.76%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0211, train/behavior_loss: 0.0001, train/strict_loss: 0.0016, val/iit_loss: 0.0103, val/IIA: 99.65%, val/accuracy: 100.00%, val/strict_accuracy: 99.80%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0064, train/behavior_loss: 0.0001, train/strict_loss: 0.0015, val/iit_loss: 0.0123, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0114, train/behavior_loss: 0.0000, train/strict_loss: 0.0014, val/iit_loss: 0.0114, val/IIA: 99.51%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0096, train/behavior_loss: 0.0001, train/strict_loss: 0.0013, val/iit_loss: 0.0060, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0202, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0049, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0072, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0040, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0085, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0086, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0115, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0056, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0284, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0117, val/IIA: 99.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0114, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0081, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0105, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0103, val/IIA: 99.55%, val/accuracy: 100.00%, val/strict_accuracy: 99.88%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0104, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0222, val/IIA: 99.29%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0167, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0045, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0125, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0320, val/IIA: 98.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0083, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0079, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0140, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0050, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0137, val/IIA: 99.51%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0078, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0036, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0060, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0039, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0156, val/IIA: 99.39%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0066, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0048, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0123, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0043, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0129, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0114, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0071, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0191, val/IIA: 99.30%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0121, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0032, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0054, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0039, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0027, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0036, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0067, val/IIA: 99.72%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0179, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0061, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0053, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0033, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0050, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0102, val/IIA: 99.60%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0084, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0038, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0053, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0085, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0058, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0047, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0044, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0067, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0027, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0081, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0054, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0065, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0163, val/IIA: 99.46%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0050, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0061, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0039, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0112, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0104, val/IIA: 99.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0050, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0072, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0039, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0053, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0042, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0037, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0035, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0025, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0031, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0018, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0046, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0019, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0081, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0042, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0048, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0043, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0035, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0029, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0033, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0032, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0019, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0016, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0011, val/IIA: 99.97%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0168, val/IIA: 99.57%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0102, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0064, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0026, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0125, val/IIA: 99.57%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0046, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0021, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0045, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0023, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0057, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0012, val/IIA: 99.97%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0027, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0016, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0017, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0015, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0011, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0007, val/IIA: 99.99%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0012, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0015, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0009, val/IIA: 99.98%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0014, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0011, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0010, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0010, val/IIA: 99.96%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0013, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0016, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_1+3')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_1+3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (2) ParenChecker + (3) UniqueExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[mlp0_hook], [appeared_mlp]]\n",
      "blocks.0.attn.hook_z.0 [None, None]\n",
      "blocks.0.attn.hook_z.1 [None, None]\n",
      "blocks.0.attn.hook_z.2 [None, [counter_head]]\n",
      "blocks.0.attn.hook_z.3 [[paren_counts_hook], None]\n",
      "blocks.1.mlp.hook_post [[mlp1_hook], [mask_mlp]]\n",
      "blocks.1.attn.hook_z.0 [None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None]\n",
      "blocks.2.mlp.hook_post [[mlp2_hook], [output_mlp]]\n",
      "blocks.2.attn.hook_z.0 [None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None]\n",
      "blocks.2.attn.hook_z.3 [[horizon_lookback_hook], None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case2, Case3]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "[LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a804e19bfedd457dafbfa05076ddcd22",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a929ac27d814832b883239a96a8aa32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/54 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.7994, train/behavior_loss: 0.3052, train/strict_loss: 0.1345, val/iit_loss: 0.4381, val/IIA: 88.08%, val/accuracy: 90.95%, val/strict_accuracy: 91.12%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2983, train/behavior_loss: 0.0870, train/strict_loss: 0.0364, val/iit_loss: 0.1716, val/IIA: 95.63%, val/accuracy: 97.45%, val/strict_accuracy: 97.38%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1292, train/behavior_loss: 0.0267, train/strict_loss: 0.0153, val/iit_loss: 0.0976, val/IIA: 97.23%, val/accuracy: 98.83%, val/strict_accuracy: 98.76%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0744, train/behavior_loss: 0.0099, train/strict_loss: 0.0068, val/iit_loss: 0.0720, val/IIA: 97.78%, val/accuracy: 99.23%, val/strict_accuracy: 99.24%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0660, train/behavior_loss: 0.0071, train/strict_loss: 0.0049, val/iit_loss: 0.0472, val/IIA: 98.36%, val/accuracy: 99.29%, val/strict_accuracy: 99.25%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0365, train/behavior_loss: 0.0049, train/strict_loss: 0.0050, val/iit_loss: 0.0377, val/IIA: 98.68%, val/accuracy: 99.34%, val/strict_accuracy: 99.30%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0304, train/behavior_loss: 0.0030, train/strict_loss: 0.0041, val/iit_loss: 0.0373, val/IIA: 98.76%, val/accuracy: 99.61%, val/strict_accuracy: 99.52%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0279, train/behavior_loss: 0.0019, train/strict_loss: 0.0037, val/iit_loss: 0.0273, val/IIA: 99.14%, val/accuracy: 99.93%, val/strict_accuracy: 99.77%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0227, train/behavior_loss: 0.0012, train/strict_loss: 0.0038, val/iit_loss: 0.0187, val/IIA: 99.36%, val/accuracy: 99.99%, val/strict_accuracy: 99.81%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0168, train/behavior_loss: 0.0007, train/strict_loss: 0.0040, val/iit_loss: 0.0213, val/IIA: 99.28%, val/accuracy: 99.55%, val/strict_accuracy: 99.47%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0188, train/behavior_loss: 0.0010, train/strict_loss: 0.0042, val/iit_loss: 0.0129, val/IIA: 99.60%, val/accuracy: 100.00%, val/strict_accuracy: 99.82%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0161, train/behavior_loss: 0.0005, train/strict_loss: 0.0034, val/iit_loss: 0.0131, val/IIA: 99.60%, val/accuracy: 99.86%, val/strict_accuracy: 99.69%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0120, train/behavior_loss: 0.0005, train/strict_loss: 0.0030, val/iit_loss: 0.0174, val/IIA: 99.36%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0162, train/behavior_loss: 0.0004, train/strict_loss: 0.0031, val/iit_loss: 0.0111, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0097, train/behavior_loss: 0.0002, train/strict_loss: 0.0030, val/iit_loss: 0.0095, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0152, train/behavior_loss: 0.0003, train/strict_loss: 0.0024, val/iit_loss: 0.0132, val/IIA: 99.52%, val/accuracy: 100.00%, val/strict_accuracy: 99.88%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0132, train/behavior_loss: 0.0002, train/strict_loss: 0.0023, val/iit_loss: 0.0107, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0087, train/behavior_loss: 0.0003, train/strict_loss: 0.0022, val/iit_loss: 0.0100, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.85%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0086, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0087, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0124, train/behavior_loss: 0.0001, train/strict_loss: 0.0023, val/iit_loss: 0.0068, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.86%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0082, train/behavior_loss: 0.0001, train/strict_loss: 0.0019, val/iit_loss: 0.0071, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0070, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0070, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0057, train/behavior_loss: 0.0001, train/strict_loss: 0.0018, val/iit_loss: 0.0056, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.92%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0051, train/behavior_loss: 0.0001, train/strict_loss: 0.0013, val/iit_loss: 0.0041, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0069, train/behavior_loss: 0.0001, train/strict_loss: 0.0016, val/iit_loss: 0.0047, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0059, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.87%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0051, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0052, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0001, train/strict_loss: 0.0016, val/iit_loss: 0.0033, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0033, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0054, train/behavior_loss: 0.0001, train/strict_loss: 0.0013, val/iit_loss: 0.0048, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0040, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0013, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0067, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0032, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0011, val/iit_loss: 0.0032, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0033, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0030, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0031, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0035, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0041, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0045, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0029, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0037, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0044, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0028, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0027, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0026, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0063, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0025, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0036, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0032, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0029, val/IIA: 99.92%, val/accuracy: 99.99%, val/strict_accuracy: 99.96%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0031, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0027, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0028, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0033, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0031, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0113, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0025, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0033, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0020, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0025, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0080, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0036, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0024, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0031, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0022, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0016, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 99.99%, val/strict_accuracy: 99.98%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0042, val/IIA: 99.88%, val/accuracy: 99.99%, val/strict_accuracy: 99.98%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0037, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0019, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0158, val/IIA: 99.51%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0020, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0030, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0019, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0019, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0024, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0025, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0026, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0036, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0043, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0029, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0025, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0023, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0013, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0024, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0027, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0021, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0027, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0020, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0013, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0019, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0020, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0016, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0021, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0021, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0023, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0022, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0021, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0024, val/IIA: 99.94%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0027, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0018, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0020, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0019, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0042, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0019, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0019, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0011, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0022, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0019, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0017, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0025, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0015, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0015, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0015, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0017, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0015, val/IIA: 99.95%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_2+3')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_2+3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (0) DuplicateRemover + (1) LeftGreater + (2) ParenChecker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[prev_equal_hook], [mlp0_hook], [mlp0_hook]]\n",
      "blocks.0.attn.hook_z.0 [[prev_token_hook], None, None]\n",
      "blocks.0.attn.hook_z.1 [None, [paren_counts_hook], None]\n",
      "blocks.0.attn.hook_z.2 [None, None, None]\n",
      "blocks.0.attn.hook_z.3 [None, None, [paren_counts_hook]]\n",
      "blocks.1.mlp.hook_post [[output_hook], None, [mlp1_hook]]\n",
      "blocks.1.attn.hook_z.0 [None, None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None, None]\n",
      "blocks.2.mlp.hook_post [None, None, [mlp2_hook]]\n",
      "blocks.2.attn.hook_z.0 [None, None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None, None]\n",
      "blocks.2.attn.hook_z.3 [None, None, [horizon_lookback_hook]]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case0, Case1, Case2]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "[LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d47fab17e6984d26bc4e58f28ae00b47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "497ce5bcc60d47cd85f32b65077adb44",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.9956, train/behavior_loss: 0.3630, train/strict_loss: 0.1667, val/iit_loss: 0.6228, val/IIA: 79.09%, val/accuracy: 85.76%, val/strict_accuracy: 83.54%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.4591, train/behavior_loss: 0.1503, train/strict_loss: 0.0895, val/iit_loss: 0.3790, val/IIA: 85.87%, val/accuracy: 90.76%, val/strict_accuracy: 90.17%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3343, train/behavior_loss: 0.0841, train/strict_loss: 0.0491, val/iit_loss: 0.2625, val/IIA: 89.47%, val/accuracy: 93.54%, val/strict_accuracy: 92.62%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2175, train/behavior_loss: 0.0557, train/strict_loss: 0.0412, val/iit_loss: 0.1919, val/IIA: 92.36%, val/accuracy: 95.96%, val/strict_accuracy: 94.98%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1598, train/behavior_loss: 0.0284, train/strict_loss: 0.0360, val/iit_loss: 0.1333, val/IIA: 95.06%, val/accuracy: 98.86%, val/strict_accuracy: 97.33%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1154, train/behavior_loss: 0.0117, train/strict_loss: 0.0316, val/iit_loss: 0.1069, val/IIA: 95.96%, val/accuracy: 99.35%, val/strict_accuracy: 98.29%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0891, train/behavior_loss: 0.0066, train/strict_loss: 0.0242, val/iit_loss: 0.0822, val/IIA: 97.12%, val/accuracy: 99.76%, val/strict_accuracy: 99.10%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0698, train/behavior_loss: 0.0042, train/strict_loss: 0.0180, val/iit_loss: 0.0824, val/IIA: 97.04%, val/accuracy: 99.67%, val/strict_accuracy: 99.25%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0771, train/behavior_loss: 0.0032, train/strict_loss: 0.0136, val/iit_loss: 0.0828, val/IIA: 97.37%, val/accuracy: 99.70%, val/strict_accuracy: 99.37%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0963, train/behavior_loss: 0.0021, train/strict_loss: 0.0095, val/iit_loss: 0.0557, val/IIA: 97.99%, val/accuracy: 99.97%, val/strict_accuracy: 99.76%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0575, train/behavior_loss: 0.0014, train/strict_loss: 0.0068, val/iit_loss: 0.0597, val/IIA: 97.94%, val/accuracy: 99.82%, val/strict_accuracy: 99.75%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0673, train/behavior_loss: 0.0009, train/strict_loss: 0.0039, val/iit_loss: 0.0863, val/IIA: 97.54%, val/accuracy: 99.98%, val/strict_accuracy: 99.81%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0511, train/behavior_loss: 0.0011, train/strict_loss: 0.0058, val/iit_loss: 0.0610, val/IIA: 98.07%, val/accuracy: 99.98%, val/strict_accuracy: 99.86%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0541, train/behavior_loss: 0.0007, train/strict_loss: 0.0042, val/iit_loss: 0.0386, val/IIA: 98.80%, val/accuracy: 99.98%, val/strict_accuracy: 99.94%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0531, train/behavior_loss: 0.0010, train/strict_loss: 0.0037, val/iit_loss: 0.0490, val/IIA: 98.07%, val/accuracy: 99.96%, val/strict_accuracy: 99.68%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0439, train/behavior_loss: 0.0007, train/strict_loss: 0.0042, val/iit_loss: 0.0609, val/IIA: 98.05%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0366, train/behavior_loss: 0.0003, train/strict_loss: 0.0025, val/iit_loss: 0.0413, val/IIA: 98.46%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0375, train/behavior_loss: 0.0004, train/strict_loss: 0.0024, val/iit_loss: 0.0628, val/IIA: 98.09%, val/accuracy: 99.92%, val/strict_accuracy: 99.79%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0384, train/behavior_loss: 0.0006, train/strict_loss: 0.0030, val/iit_loss: 0.0475, val/IIA: 98.36%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0295, train/behavior_loss: 0.0003, train/strict_loss: 0.0019, val/iit_loss: 0.0522, val/IIA: 98.04%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0357, train/behavior_loss: 0.0003, train/strict_loss: 0.0016, val/iit_loss: 0.0286, val/IIA: 99.05%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0297, train/behavior_loss: 0.0003, train/strict_loss: 0.0021, val/iit_loss: 0.0298, val/IIA: 98.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0309, train/behavior_loss: 0.0003, train/strict_loss: 0.0018, val/iit_loss: 0.0398, val/IIA: 98.81%, val/accuracy: 100.00%, val/strict_accuracy: 99.93%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0329, train/behavior_loss: 0.0002, train/strict_loss: 0.0015, val/iit_loss: 0.0406, val/IIA: 98.65%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0269, train/behavior_loss: 0.0002, train/strict_loss: 0.0017, val/iit_loss: 0.0231, val/IIA: 99.21%, val/accuracy: 100.00%, val/strict_accuracy: 99.90%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0351, train/behavior_loss: 0.0002, train/strict_loss: 0.0015, val/iit_loss: 0.0202, val/IIA: 99.31%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0274, train/behavior_loss: 0.0003, train/strict_loss: 0.0018, val/iit_loss: 0.0245, val/IIA: 98.96%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0260, train/behavior_loss: 0.0002, train/strict_loss: 0.0015, val/iit_loss: 0.0294, val/IIA: 99.09%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0216, train/behavior_loss: 0.0002, train/strict_loss: 0.0014, val/iit_loss: 0.0165, val/IIA: 99.45%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0204, train/behavior_loss: 0.0001, train/strict_loss: 0.0013, val/iit_loss: 0.0286, val/IIA: 98.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0254, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0231, val/IIA: 99.27%, val/accuracy: 99.98%, val/strict_accuracy: 99.95%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0260, train/behavior_loss: 0.0002, train/strict_loss: 0.0013, val/iit_loss: 0.0177, val/IIA: 99.39%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0189, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0214, val/IIA: 99.23%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0176, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0417, val/IIA: 98.83%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0192, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0244, val/IIA: 99.32%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0281, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0196, val/IIA: 99.34%, val/accuracy: 99.98%, val/strict_accuracy: 99.92%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0141, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0181, val/IIA: 99.33%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0236, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0179, val/IIA: 99.34%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0224, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0248, val/IIA: 99.20%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0134, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0111, val/IIA: 99.65%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0150, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0428, val/IIA: 98.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0160, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0465, val/IIA: 98.72%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0268, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0157, val/IIA: 99.51%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0138, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0246, val/IIA: 99.07%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0182, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0152, val/IIA: 99.44%, val/accuracy: 99.98%, val/strict_accuracy: 99.91%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0120, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0141, val/IIA: 99.49%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0124, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0083, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0099, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0091, val/IIA: 99.70%, val/accuracy: 99.99%, val/strict_accuracy: 99.96%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0194, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0107, val/IIA: 99.60%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0110, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0173, val/IIA: 99.29%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0117, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0079, val/IIA: 99.74%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0113, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0076, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0131, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0067, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0093, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0084, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0153, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0080, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0098, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0145, val/IIA: 99.36%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0116, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0106, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0090, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0082, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0079, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0162, val/IIA: 99.51%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0122, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0087, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0082, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0080, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0074, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0081, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0087, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0049, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0072, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0054, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0087, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0054, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0094, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0082, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0070, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0066, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0070, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0094, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0079, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0133, val/IIA: 99.47%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0062, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0063, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0121, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0099, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0046, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0051, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0054, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0050, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0056, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0115, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0075, val/IIA: 99.72%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0043, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0041, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0036, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0052, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0063, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0061, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0034, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0082, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0053, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0053, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0033, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0083, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0050, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0034, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0033, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0032, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0034, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0036, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0027, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0033, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0048, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0035, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0026, val/IIA: 99.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0062, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0029, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0049, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0027, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0042, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_0+1+2')"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_0+1+2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (0) DuplicateRemover + (1) LeftGreater + (3) UniqueExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[prev_equal_hook], [mlp0_hook], [appeared_mlp]]\n",
      "blocks.0.attn.hook_z.0 [[prev_token_hook], None, None]\n",
      "blocks.0.attn.hook_z.1 [None, [paren_counts_hook], None]\n",
      "blocks.0.attn.hook_z.2 [None, None, [counter_head]]\n",
      "blocks.0.attn.hook_z.3 [None, None, None]\n",
      "blocks.1.mlp.hook_post [[output_hook], None, [mask_mlp]]\n",
      "blocks.1.attn.hook_z.0 [None, None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None, None]\n",
      "blocks.2.mlp.hook_post [None, None, [output_mlp]]\n",
      "blocks.2.attn.hook_z.0 [None, None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None, None]\n",
      "blocks.2.attn.hook_z.3 [None, None, None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case0, Case1, Case3]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "[LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4af62cc23945468b9c37847c369428a8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b15e06600df462e9a2e69cadf610401",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/85 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.7619, train/behavior_loss: 0.2693, train/strict_loss: 0.1164, val/iit_loss: 0.4183, val/IIA: 85.40%, val/accuracy: 90.46%, val/strict_accuracy: 89.81%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3279, train/behavior_loss: 0.0696, train/strict_loss: 0.0427, val/iit_loss: 0.2384, val/IIA: 90.93%, val/accuracy: 97.18%, val/strict_accuracy: 96.29%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1841, train/behavior_loss: 0.0224, train/strict_loss: 0.0327, val/iit_loss: 0.1624, val/IIA: 94.10%, val/accuracy: 99.76%, val/strict_accuracy: 98.80%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1386, train/behavior_loss: 0.0087, train/strict_loss: 0.0284, val/iit_loss: 0.1228, val/IIA: 95.61%, val/accuracy: 99.77%, val/strict_accuracy: 99.08%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1260, train/behavior_loss: 0.0059, train/strict_loss: 0.0262, val/iit_loss: 0.1101, val/IIA: 95.89%, val/accuracy: 99.86%, val/strict_accuracy: 99.30%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1138, train/behavior_loss: 0.0038, train/strict_loss: 0.0221, val/iit_loss: 0.0952, val/IIA: 96.47%, val/accuracy: 100.00%, val/strict_accuracy: 99.57%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1069, train/behavior_loss: 0.0027, train/strict_loss: 0.0176, val/iit_loss: 0.0999, val/IIA: 96.25%, val/accuracy: 99.99%, val/strict_accuracy: 99.87%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0942, train/behavior_loss: 0.0025, train/strict_loss: 0.0118, val/iit_loss: 0.0789, val/IIA: 97.18%, val/accuracy: 99.90%, val/strict_accuracy: 99.82%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0916, train/behavior_loss: 0.0017, train/strict_loss: 0.0070, val/iit_loss: 0.1075, val/IIA: 96.65%, val/accuracy: 99.98%, val/strict_accuracy: 99.93%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0732, train/behavior_loss: 0.0012, train/strict_loss: 0.0060, val/iit_loss: 0.0828, val/IIA: 97.20%, val/accuracy: 99.85%, val/strict_accuracy: 99.84%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0688, train/behavior_loss: 0.0012, train/strict_loss: 0.0048, val/iit_loss: 0.1096, val/IIA: 96.43%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0694, train/behavior_loss: 0.0011, train/strict_loss: 0.0058, val/iit_loss: 0.0518, val/IIA: 98.13%, val/accuracy: 99.99%, val/strict_accuracy: 99.94%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0585, train/behavior_loss: 0.0007, train/strict_loss: 0.0042, val/iit_loss: 0.0753, val/IIA: 97.48%, val/accuracy: 99.99%, val/strict_accuracy: 99.92%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0571, train/behavior_loss: 0.0009, train/strict_loss: 0.0040, val/iit_loss: 0.0427, val/IIA: 98.48%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0447, train/behavior_loss: 0.0004, train/strict_loss: 0.0034, val/iit_loss: 0.0541, val/IIA: 98.00%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0493, train/behavior_loss: 0.0004, train/strict_loss: 0.0029, val/iit_loss: 0.0318, val/IIA: 98.93%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0543, train/behavior_loss: 0.0004, train/strict_loss: 0.0031, val/iit_loss: 0.0293, val/IIA: 98.84%, val/accuracy: 99.99%, val/strict_accuracy: 99.94%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0323, train/behavior_loss: 0.0003, train/strict_loss: 0.0023, val/iit_loss: 0.0338, val/IIA: 98.69%, val/accuracy: 99.98%, val/strict_accuracy: 99.95%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0369, train/behavior_loss: 0.0002, train/strict_loss: 0.0018, val/iit_loss: 0.0290, val/IIA: 99.00%, val/accuracy: 99.93%, val/strict_accuracy: 99.89%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0418, train/behavior_loss: 0.0003, train/strict_loss: 0.0025, val/iit_loss: 0.0369, val/IIA: 98.65%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0355, train/behavior_loss: 0.0002, train/strict_loss: 0.0022, val/iit_loss: 0.0291, val/IIA: 98.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0409, train/behavior_loss: 0.0002, train/strict_loss: 0.0019, val/iit_loss: 0.0266, val/IIA: 99.03%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0382, train/behavior_loss: 0.0004, train/strict_loss: 0.0018, val/iit_loss: 0.0404, val/IIA: 98.49%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0286, train/behavior_loss: 0.0002, train/strict_loss: 0.0017, val/iit_loss: 0.0175, val/IIA: 99.38%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0280, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0329, val/IIA: 98.97%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0277, train/behavior_loss: 0.0001, train/strict_loss: 0.0017, val/iit_loss: 0.0171, val/IIA: 99.40%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0217, train/behavior_loss: 0.0001, train/strict_loss: 0.0015, val/iit_loss: 0.0342, val/IIA: 98.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0250, train/behavior_loss: 0.0002, train/strict_loss: 0.0018, val/iit_loss: 0.0179, val/IIA: 99.39%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0236, train/behavior_loss: 0.0002, train/strict_loss: 0.0017, val/iit_loss: 0.0310, val/IIA: 99.01%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0252, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0174, val/IIA: 99.40%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0311, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0187, val/IIA: 99.32%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0253, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0225, val/IIA: 99.17%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0193, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0128, val/IIA: 99.58%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0162, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0165, val/IIA: 99.43%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0146, train/behavior_loss: 0.0000, train/strict_loss: 0.0012, val/iit_loss: 0.0465, val/IIA: 98.64%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0216, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0106, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0130, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0377, val/IIA: 98.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0225, train/behavior_loss: 0.0000, train/strict_loss: 0.0010, val/iit_loss: 0.0162, val/IIA: 99.47%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0233, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0259, val/IIA: 99.23%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0158, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0131, val/IIA: 99.55%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0174, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0134, val/IIA: 99.53%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0200, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0140, val/IIA: 99.50%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0135, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0125, val/IIA: 99.52%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0152, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0093, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0143, train/behavior_loss: 0.0000, train/strict_loss: 0.0009, val/iit_loss: 0.0154, val/IIA: 99.46%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0135, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0094, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0138, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0094, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0103, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0109, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0168, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0294, val/IIA: 99.20%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0116, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0124, val/IIA: 99.60%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0143, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0089, val/IIA: 99.66%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0088, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0154, val/IIA: 99.41%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0199, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0174, val/IIA: 99.42%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0110, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0112, val/IIA: 99.60%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0178, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0079, val/IIA: 99.72%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0101, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0114, val/IIA: 99.57%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0108, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0063, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0122, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0102, val/IIA: 99.64%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0160, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0085, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0106, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0157, val/IIA: 99.44%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0077, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0078, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0109, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0066, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0102, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0120, val/IIA: 99.56%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0079, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0084, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0102, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0094, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0077, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0129, val/IIA: 99.60%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0096, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0063, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0070, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0090, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0094, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0081, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0101, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0086, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0087, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0073, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0077, val/IIA: 99.72%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0116, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0085, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0048, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0057, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0094, val/IIA: 99.64%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0093, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0069, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0053, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0056, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0049, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0057, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0087, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0051, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0099, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0078, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0109, val/IIA: 99.58%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0072, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0058, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0048, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0041, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0052, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0064, val/IIA: 99.74%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0047, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0080, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0045, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0055, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0045, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0033, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0090, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0036, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0082, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0052, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0037, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0059, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0038, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0038, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0033, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_0+1+3')"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_0+1+3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (0) DuplicateRemover + (2) ParenChecker + (3) UniqueExtractor "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[prev_equal_hook], [mlp0_hook], [appeared_mlp]]\n",
      "blocks.0.attn.hook_z.0 [[prev_token_hook], None, None]\n",
      "blocks.0.attn.hook_z.1 [None, None, None]\n",
      "blocks.0.attn.hook_z.2 [None, None, [counter_head]]\n",
      "blocks.0.attn.hook_z.3 [None, [paren_counts_hook], None]\n",
      "blocks.1.mlp.hook_post [[output_hook], [mlp1_hook], [mask_mlp]]\n",
      "blocks.1.attn.hook_z.0 [None, None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None, None]\n",
      "blocks.2.mlp.hook_post [None, [mlp2_hook], [output_mlp]]\n",
      "blocks.2.attn.hook_z.0 [None, None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None, None]\n",
      "blocks.2.attn.hook_z.3 [None, [horizon_lookback_hook], None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case0, Case2, Case3]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "[LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81c0883316804c5a9faf7ced497ed72e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d4ae7aeb61b42359ed0e92e84994867",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/85 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.8056, train/behavior_loss: 0.2910, train/strict_loss: 0.1275, val/iit_loss: 0.4016, val/IIA: 85.89%, val/accuracy: 89.19%, val/strict_accuracy: 88.75%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2906, train/behavior_loss: 0.0781, train/strict_loss: 0.0423, val/iit_loss: 0.2135, val/IIA: 92.17%, val/accuracy: 96.62%, val/strict_accuracy: 95.61%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1896, train/behavior_loss: 0.0274, train/strict_loss: 0.0284, val/iit_loss: 0.1267, val/IIA: 95.87%, val/accuracy: 99.26%, val/strict_accuracy: 98.33%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0948, train/behavior_loss: 0.0085, train/strict_loss: 0.0222, val/iit_loss: 0.1018, val/IIA: 96.38%, val/accuracy: 99.54%, val/strict_accuracy: 98.59%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0785, train/behavior_loss: 0.0047, train/strict_loss: 0.0161, val/iit_loss: 0.0807, val/IIA: 97.37%, val/accuracy: 99.53%, val/strict_accuracy: 99.33%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0732, train/behavior_loss: 0.0038, train/strict_loss: 0.0113, val/iit_loss: 0.0664, val/IIA: 97.72%, val/accuracy: 99.76%, val/strict_accuracy: 99.46%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0597, train/behavior_loss: 0.0035, train/strict_loss: 0.0077, val/iit_loss: 0.0744, val/IIA: 97.27%, val/accuracy: 99.67%, val/strict_accuracy: 99.53%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0677, train/behavior_loss: 0.0032, train/strict_loss: 0.0061, val/iit_loss: 0.0497, val/IIA: 98.27%, val/accuracy: 99.87%, val/strict_accuracy: 99.71%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0647, train/behavior_loss: 0.0024, train/strict_loss: 0.0052, val/iit_loss: 0.0482, val/IIA: 98.15%, val/accuracy: 99.65%, val/strict_accuracy: 99.52%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0441, train/behavior_loss: 0.0018, train/strict_loss: 0.0045, val/iit_loss: 0.0501, val/IIA: 98.26%, val/accuracy: 99.93%, val/strict_accuracy: 99.74%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0482, train/behavior_loss: 0.0015, train/strict_loss: 0.0040, val/iit_loss: 0.0337, val/IIA: 98.80%, val/accuracy: 99.93%, val/strict_accuracy: 99.76%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0283, train/behavior_loss: 0.0008, train/strict_loss: 0.0032, val/iit_loss: 0.0443, val/IIA: 98.30%, val/accuracy: 99.95%, val/strict_accuracy: 99.86%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0431, train/behavior_loss: 0.0009, train/strict_loss: 0.0031, val/iit_loss: 0.0271, val/IIA: 99.09%, val/accuracy: 99.94%, val/strict_accuracy: 99.83%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0351, train/behavior_loss: 0.0011, train/strict_loss: 0.0029, val/iit_loss: 0.0379, val/IIA: 98.58%, val/accuracy: 99.91%, val/strict_accuracy: 99.81%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0376, train/behavior_loss: 0.0005, train/strict_loss: 0.0028, val/iit_loss: 0.0320, val/IIA: 98.93%, val/accuracy: 99.92%, val/strict_accuracy: 99.81%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0428, train/behavior_loss: 0.0008, train/strict_loss: 0.0031, val/iit_loss: 0.0407, val/IIA: 98.53%, val/accuracy: 99.82%, val/strict_accuracy: 99.75%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0226, train/behavior_loss: 0.0006, train/strict_loss: 0.0025, val/iit_loss: 0.0263, val/IIA: 99.01%, val/accuracy: 99.96%, val/strict_accuracy: 99.90%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0199, train/behavior_loss: 0.0004, train/strict_loss: 0.0021, val/iit_loss: 0.0297, val/IIA: 98.90%, val/accuracy: 99.99%, val/strict_accuracy: 99.90%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0303, train/behavior_loss: 0.0003, train/strict_loss: 0.0020, val/iit_loss: 0.0268, val/IIA: 99.05%, val/accuracy: 99.99%, val/strict_accuracy: 99.93%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0237, train/behavior_loss: 0.0004, train/strict_loss: 0.0020, val/iit_loss: 0.0193, val/IIA: 99.33%, val/accuracy: 99.97%, val/strict_accuracy: 99.91%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0216, train/behavior_loss: 0.0003, train/strict_loss: 0.0018, val/iit_loss: 0.0137, val/IIA: 99.56%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0221, train/behavior_loss: 0.0003, train/strict_loss: 0.0017, val/iit_loss: 0.0372, val/IIA: 98.85%, val/accuracy: 99.97%, val/strict_accuracy: 99.91%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0177, train/behavior_loss: 0.0002, train/strict_loss: 0.0017, val/iit_loss: 0.0211, val/IIA: 99.22%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0236, train/behavior_loss: 0.0002, train/strict_loss: 0.0016, val/iit_loss: 0.0154, val/IIA: 99.47%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0200, train/behavior_loss: 0.0002, train/strict_loss: 0.0015, val/iit_loss: 0.0224, val/IIA: 99.22%, val/accuracy: 99.99%, val/strict_accuracy: 99.95%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0204, train/behavior_loss: 0.0002, train/strict_loss: 0.0014, val/iit_loss: 0.0265, val/IIA: 99.03%, val/accuracy: 99.99%, val/strict_accuracy: 99.96%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0225, train/behavior_loss: 0.0002, train/strict_loss: 0.0015, val/iit_loss: 0.0163, val/IIA: 99.43%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0193, train/behavior_loss: 0.0003, train/strict_loss: 0.0014, val/iit_loss: 0.0140, val/IIA: 99.55%, val/accuracy: 99.96%, val/strict_accuracy: 99.94%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0132, train/behavior_loss: 0.0002, train/strict_loss: 0.0012, val/iit_loss: 0.0204, val/IIA: 99.20%, val/accuracy: 99.99%, val/strict_accuracy: 99.97%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0162, train/behavior_loss: 0.0002, train/strict_loss: 0.0012, val/iit_loss: 0.0211, val/IIA: 99.28%, val/accuracy: 99.97%, val/strict_accuracy: 99.92%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0155, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0154, val/IIA: 99.48%, val/accuracy: 99.98%, val/strict_accuracy: 99.96%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0183, train/behavior_loss: 0.0001, train/strict_loss: 0.0013, val/iit_loss: 0.0147, val/IIA: 99.51%, val/accuracy: 99.96%, val/strict_accuracy: 99.95%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0109, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0111, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0137, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0136, val/IIA: 99.49%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0103, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0132, val/IIA: 99.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0145, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0109, val/IIA: 99.60%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0140, train/behavior_loss: 0.0002, train/strict_loss: 0.0010, val/iit_loss: 0.0303, val/IIA: 98.97%, val/accuracy: 99.96%, val/strict_accuracy: 99.94%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0163, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0111, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0095, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0197, val/IIA: 99.33%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0155, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0120, val/IIA: 99.56%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0154, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0129, val/IIA: 99.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0129, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0097, val/IIA: 99.65%, val/accuracy: 99.99%, val/strict_accuracy: 99.99%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0114, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0136, val/IIA: 99.47%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0071, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0076, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0087, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0087, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0087, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0097, val/IIA: 99.66%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0115, train/behavior_loss: 0.0000, train/strict_loss: 0.0008, val/iit_loss: 0.0088, val/IIA: 99.69%, val/accuracy: 99.99%, val/strict_accuracy: 99.96%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0094, train/behavior_loss: 0.0002, train/strict_loss: 0.0008, val/iit_loss: 0.0092, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0128, val/IIA: 99.63%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0072, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0105, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0106, train/behavior_loss: 0.0001, train/strict_loss: 0.0006, val/iit_loss: 0.0085, val/IIA: 99.72%, val/accuracy: 99.97%, val/strict_accuracy: 99.91%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0082, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0072, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0098, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0094, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0092, val/IIA: 99.66%, val/accuracy: 99.99%, val/strict_accuracy: 99.96%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0080, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0067, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0092, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0070, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0001, train/strict_loss: 0.0005, val/iit_loss: 0.0068, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0070, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0112, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0090, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0083, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0074, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0056, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0088, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0063, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0080, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0095, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0095, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0133, val/IIA: 99.62%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0076, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0070, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0066, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0063, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0063, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0066, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0066, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0052, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0048, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0049, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0050, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0067, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0048, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0057, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0064, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0056, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0053, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0064, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0054, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0071, val/IIA: 99.76%, val/accuracy: 99.99%, val/strict_accuracy: 99.97%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0060, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0052, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0073, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0066, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0050, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0071, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0075, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0077, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0051, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0053, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0064, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0045, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0061, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0050, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0047, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0058, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0066, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0047, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0056, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0053, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0061, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0047, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0067, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0045, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0013, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0045, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0042, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0045, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0041, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0041, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0052, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0050, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0044, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_0+2+3')"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_0+2+3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (1) LeftGreater + (2) ParenChecker + (3) UniqueExtractor "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[mlp0_hook], [mlp0_hook], [appeared_mlp]]\n",
      "blocks.0.attn.hook_z.0 [None, None, None]\n",
      "blocks.0.attn.hook_z.1 [[paren_counts_hook], None, None]\n",
      "blocks.0.attn.hook_z.2 [None, None, [counter_head]]\n",
      "blocks.0.attn.hook_z.3 [None, [paren_counts_hook], None]\n",
      "blocks.1.mlp.hook_post [None, [mlp1_hook], [mask_mlp]]\n",
      "blocks.1.attn.hook_z.0 [None, None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None, None]\n",
      "blocks.2.mlp.hook_post [None, [mlp2_hook], [output_mlp]]\n",
      "blocks.2.attn.hook_z.0 [None, None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None, None]\n",
      "blocks.2.attn.hook_z.3 [None, [horizon_lookback_hook], None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case1, Case2, Case3]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "[LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc801924670749f59a3ecc4d52d6f365",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5578d30746bd4f5fbc7c89463395491f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.7057, train/behavior_loss: 0.2560, train/strict_loss: 0.1102, val/iit_loss: 0.3788, val/IIA: 89.01%, val/accuracy: 95.11%, val/strict_accuracy: 95.00%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2848, train/behavior_loss: 0.0660, train/strict_loss: 0.0299, val/iit_loss: 0.1911, val/IIA: 94.42%, val/accuracy: 98.38%, val/strict_accuracy: 98.34%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1570, train/behavior_loss: 0.0237, train/strict_loss: 0.0133, val/iit_loss: 0.1371, val/IIA: 95.40%, val/accuracy: 98.64%, val/strict_accuracy: 98.49%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1056, train/behavior_loss: 0.0111, train/strict_loss: 0.0085, val/iit_loss: 0.0870, val/IIA: 97.03%, val/accuracy: 99.63%, val/strict_accuracy: 99.24%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0679, train/behavior_loss: 0.0053, train/strict_loss: 0.0063, val/iit_loss: 0.0680, val/IIA: 97.76%, val/accuracy: 99.61%, val/strict_accuracy: 99.46%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0679, train/behavior_loss: 0.0035, train/strict_loss: 0.0055, val/iit_loss: 0.0817, val/IIA: 97.11%, val/accuracy: 99.81%, val/strict_accuracy: 99.63%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0586, train/behavior_loss: 0.0022, train/strict_loss: 0.0045, val/iit_loss: 0.0387, val/IIA: 98.74%, val/accuracy: 99.99%, val/strict_accuracy: 99.93%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0537, train/behavior_loss: 0.0016, train/strict_loss: 0.0033, val/iit_loss: 0.0460, val/IIA: 98.42%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0411, train/behavior_loss: 0.0009, train/strict_loss: 0.0025, val/iit_loss: 0.0515, val/IIA: 98.55%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0465, train/behavior_loss: 0.0007, train/strict_loss: 0.0020, val/iit_loss: 0.0313, val/IIA: 98.88%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0313, train/behavior_loss: 0.0005, train/strict_loss: 0.0020, val/iit_loss: 0.0245, val/IIA: 99.21%, val/accuracy: 99.98%, val/strict_accuracy: 99.94%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0291, train/behavior_loss: 0.0004, train/strict_loss: 0.0019, val/iit_loss: 0.0259, val/IIA: 99.01%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0242, train/behavior_loss: 0.0003, train/strict_loss: 0.0020, val/iit_loss: 0.0233, val/IIA: 99.15%, val/accuracy: 99.99%, val/strict_accuracy: 99.96%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0232, train/behavior_loss: 0.0002, train/strict_loss: 0.0015, val/iit_loss: 0.0308, val/IIA: 98.85%, val/accuracy: 99.96%, val/strict_accuracy: 99.91%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0245, train/behavior_loss: 0.0003, train/strict_loss: 0.0021, val/iit_loss: 0.0269, val/IIA: 99.02%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0239, train/behavior_loss: 0.0004, train/strict_loss: 0.0024, val/iit_loss: 0.0147, val/IIA: 99.53%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0147, train/behavior_loss: 0.0001, train/strict_loss: 0.0014, val/iit_loss: 0.0159, val/IIA: 99.44%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0154, train/behavior_loss: 0.0002, train/strict_loss: 0.0012, val/iit_loss: 0.0270, val/IIA: 99.08%, val/accuracy: 99.99%, val/strict_accuracy: 99.93%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0183, train/behavior_loss: 0.0002, train/strict_loss: 0.0016, val/iit_loss: 0.0189, val/IIA: 99.34%, val/accuracy: 99.99%, val/strict_accuracy: 99.86%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0155, train/behavior_loss: 0.0002, train/strict_loss: 0.0012, val/iit_loss: 0.0161, val/IIA: 99.44%, val/accuracy: 100.00%, val/strict_accuracy: 99.94%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0145, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0155, val/IIA: 99.46%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0143, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0188, val/IIA: 99.30%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0125, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0140, val/IIA: 99.50%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0097, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0102, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0119, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0105, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0116, train/behavior_loss: 0.0001, train/strict_loss: 0.0006, val/iit_loss: 0.0172, val/IIA: 99.44%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0109, train/behavior_loss: 0.0001, train/strict_loss: 0.0007, val/iit_loss: 0.0118, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0078, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0105, val/IIA: 99.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0094, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0117, val/IIA: 99.58%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0100, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0090, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0121, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0109, val/IIA: 99.66%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0082, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0089, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0099, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0080, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0084, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0111, val/IIA: 99.57%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0086, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0087, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0122, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0081, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0078, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0061, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0091, val/IIA: 99.65%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0080, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0110, val/IIA: 99.64%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0069, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0072, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0126, val/IIA: 99.57%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0073, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0096, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0069, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0052, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0076, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0047, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0068, val/IIA: 99.74%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0082, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0105, val/IIA: 99.56%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0104, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0059, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0143, val/IIA: 99.49%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0102, val/IIA: 99.64%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0070, val/IIA: 99.71%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0051, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0055, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0047, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0048, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0054, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0044, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0054, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0063, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0060, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0047, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0064, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0052, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0049, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0051, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0030, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0043, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0074, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0081, val/IIA: 99.74%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0049, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0042, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0041, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0040, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0051, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0038, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0068, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0056, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0047, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0040, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0050, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0067, val/IIA: 99.73%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0041, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0042, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0039, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0053, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0045, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0036, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0038, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0050, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0038, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0054, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0040, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0040, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0030, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0036, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0042, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0046, val/IIA: 99.84%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0024, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0039, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0062, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0039, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0045, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0054, val/IIA: 99.80%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0036, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0043, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0036, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0066, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0037, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0038, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0040, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0054, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0036, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0028, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0058, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0029, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0033, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0033, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0032, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0042, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0032, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0029, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_1+2+3')"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_1+2+3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (0) DuplicateRemover + (1) LeftGreater + (2) ParenChecker + (3) UniqueExtractor "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "\n",
      "blocks.0.mlp.hook_post [[prev_equal_hook], [mlp0_hook], [mlp0_hook], [appeared_mlp]]\n",
      "blocks.0.attn.hook_z.0 [[prev_token_hook], None, None, None]\n",
      "blocks.0.attn.hook_z.1 [None, [paren_counts_hook], None, None]\n",
      "blocks.0.attn.hook_z.2 [None, None, None, [counter_head]]\n",
      "blocks.0.attn.hook_z.3 [None, None, [paren_counts_hook], None]\n",
      "blocks.1.mlp.hook_post [[output_hook], None, [mlp1_hook], [mask_mlp]]\n",
      "blocks.1.attn.hook_z.0 [None, None, None, None]\n",
      "blocks.1.attn.hook_z.1 [None, None, None, None]\n",
      "blocks.1.attn.hook_z.2 [None, None, None, None]\n",
      "blocks.1.attn.hook_z.3 [None, None, None, None]\n",
      "blocks.2.mlp.hook_post [None, None, [mlp2_hook], [output_mlp]]\n",
      "blocks.2.attn.hook_z.0 [None, None, None, None]\n",
      "blocks.2.attn.hook_z.1 [None, None, None, None]\n",
      "blocks.2.attn.hook_z.2 [None, None, None, None]\n",
      "blocks.2.attn.hook_z.3 [None, None, [horizon_lookback_hook], None]\n"
     ]
    }
   ],
   "source": [
    "cases = [Case0, Case1, Case2, Case3]\n",
    "poly_hl_model = PolyHLModel(hl_classes=cases, size_expansion=1)\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "\n",
    "print()\n",
    "for k, v in poly_hl_model.corr_mapping.items():\n",
    "    print(k, v)\n",
    "\n",
    "dataset_cases = [dataset_mapping[case] for case in cases]\n",
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=seed) for dsetcase in dataset_cases]\n",
    "poly_dataset = PolyModelDataset(dsets, n_ctx=poly_hl_model.cfg.n_ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  mps\n",
      "input_hook {LLNode(name='blocks.0.hook_resid_pre', index=[:], subspace=None)}\n",
      "mlp_hooks.0 {LLNode(name='blocks.0.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.0.0 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "attn_hooks.0.1 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 1, :], subspace=None)}\n",
      "attn_hooks.0.2 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 2, :], subspace=None)}\n",
      "attn_hooks.0.3 {LLNode(name='blocks.0.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "mlp_hooks.1 {LLNode(name='blocks.1.mlp.hook_post', index=[:], subspace=None)}\n",
      "task_hook {LLNode(name='blocks.1.attn.hook_z', index=[:, :, 0, :], subspace=None)}\n",
      "mlp_hooks.2 {LLNode(name='blocks.2.mlp.hook_post', index=[:], subspace=None)}\n",
      "attn_hooks.2.3 {LLNode(name='blocks.2.attn.hook_z', index=[:, :, 3, :], subspace=None)}\n",
      "[LLNode(name='blocks.1.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 2, :], subspace=None), LLNode(name='blocks.1.attn.hook_z', index=[:, :, 3, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 0, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 1, :], subspace=None), LLNode(name='blocks.2.attn.hook_z', index=[:, :, 2, :], subspace=None)]\n",
      "training_args={'batch_size': 256, 'num_workers': 0, 'early_stop': True, 'lr_scheduler': <class 'torch.optim.lr_scheduler.LinearLR'>, 'scheduler_val_metric': ['val/accuracy', 'val/IIA'], 'scheduler_mode': 'max', 'scheduler_kwargs': {'start_factor': 1, 'end_factor': 0.2, 'total_iters': 100}, 'clip_grad_norm': 1.0, 'seed': 42, 'detach_while_caching': True, 'optimizer_cls': <class 'torch.optim.adam.Adam'>, 'optimizer_kwargs': {'lr': 0.001, 'betas': (0.9, 0.9)}, 'atol': 0.05, 'use_single_loss': True, 'iit_weight': 1.0, 'behavior_weight': 0.4, 'val_IIA_sampling': 'all', 'strict_weight': 0.4, 'siit_sampling': 'sample_all'}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "26ccc94c46884c1e850137bb12489a08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d086dc310ba74d31a6ae9c1e477eba6c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Batches:   0%|          | 0/107 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: lr: 9.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.7860, train/behavior_loss: 0.2727, train/strict_loss: 0.1158, val/iit_loss: 0.4426, val/IIA: 84.27%, val/accuracy: 92.07%, val/strict_accuracy: 90.99%\n",
      "Epoch 2: lr: 9.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.3287, train/behavior_loss: 0.0676, train/strict_loss: 0.0345, val/iit_loss: 0.2746, val/IIA: 90.16%, val/accuracy: 97.41%, val/strict_accuracy: 96.47%\n",
      "Epoch 3: lr: 9.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.2217, train/behavior_loss: 0.0269, train/strict_loss: 0.0218, val/iit_loss: 0.1841, val/IIA: 93.73%, val/accuracy: 99.13%, val/strict_accuracy: 98.59%\n",
      "Epoch 4: lr: 9.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1529, train/behavior_loss: 0.0120, train/strict_loss: 0.0178, val/iit_loss: 0.1357, val/IIA: 95.25%, val/accuracy: 99.56%, val/strict_accuracy: 98.93%\n",
      "Epoch 5: lr: 9.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1231, train/behavior_loss: 0.0081, train/strict_loss: 0.0126, val/iit_loss: 0.1267, val/IIA: 95.59%, val/accuracy: 99.68%, val/strict_accuracy: 99.55%\n",
      "Epoch 6: lr: 9.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1261, train/behavior_loss: 0.0060, train/strict_loss: 0.0092, val/iit_loss: 0.1044, val/IIA: 96.22%, val/accuracy: 99.31%, val/strict_accuracy: 98.98%\n",
      "Epoch 7: lr: 9.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.1008, train/behavior_loss: 0.0037, train/strict_loss: 0.0061, val/iit_loss: 0.1018, val/IIA: 96.60%, val/accuracy: 99.83%, val/strict_accuracy: 99.56%\n",
      "Epoch 8: lr: 9.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0853, train/behavior_loss: 0.0028, train/strict_loss: 0.0046, val/iit_loss: 0.0863, val/IIA: 96.80%, val/accuracy: 99.83%, val/strict_accuracy: 99.81%\n",
      "Epoch 9: lr: 9.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0807, train/behavior_loss: 0.0019, train/strict_loss: 0.0033, val/iit_loss: 0.0818, val/IIA: 97.23%, val/accuracy: 99.93%, val/strict_accuracy: 99.80%\n",
      "Epoch 10: lr: 9.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0662, train/behavior_loss: 0.0016, train/strict_loss: 0.0033, val/iit_loss: 0.0661, val/IIA: 97.77%, val/accuracy: 99.94%, val/strict_accuracy: 99.86%\n",
      "Epoch 11: lr: 9.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0626, train/behavior_loss: 0.0015, train/strict_loss: 0.0034, val/iit_loss: 0.0858, val/IIA: 97.07%, val/accuracy: 99.93%, val/strict_accuracy: 99.83%\n",
      "Epoch 12: lr: 9.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0680, train/behavior_loss: 0.0011, train/strict_loss: 0.0032, val/iit_loss: 0.0546, val/IIA: 98.06%, val/accuracy: 99.93%, val/strict_accuracy: 99.88%\n",
      "Epoch 13: lr: 8.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0493, train/behavior_loss: 0.0009, train/strict_loss: 0.0029, val/iit_loss: 0.0509, val/IIA: 98.16%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 14: lr: 8.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0466, train/behavior_loss: 0.0007, train/strict_loss: 0.0032, val/iit_loss: 0.0488, val/IIA: 98.15%, val/accuracy: 99.92%, val/strict_accuracy: 99.82%\n",
      "Epoch 15: lr: 8.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0447, train/behavior_loss: 0.0007, train/strict_loss: 0.0029, val/iit_loss: 0.0367, val/IIA: 98.77%, val/accuracy: 99.99%, val/strict_accuracy: 99.93%\n",
      "Epoch 16: lr: 8.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0408, train/behavior_loss: 0.0005, train/strict_loss: 0.0025, val/iit_loss: 0.0437, val/IIA: 98.52%, val/accuracy: 99.96%, val/strict_accuracy: 99.87%\n",
      "Epoch 17: lr: 8.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0386, train/behavior_loss: 0.0005, train/strict_loss: 0.0024, val/iit_loss: 0.0377, val/IIA: 98.64%, val/accuracy: 99.99%, val/strict_accuracy: 99.95%\n",
      "Epoch 18: lr: 8.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0409, train/behavior_loss: 0.0004, train/strict_loss: 0.0024, val/iit_loss: 0.0342, val/IIA: 98.75%, val/accuracy: 99.99%, val/strict_accuracy: 99.94%\n",
      "Epoch 19: lr: 8.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0331, train/behavior_loss: 0.0004, train/strict_loss: 0.0018, val/iit_loss: 0.0411, val/IIA: 98.53%, val/accuracy: 99.99%, val/strict_accuracy: 99.94%\n",
      "Epoch 20: lr: 8.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0382, train/behavior_loss: 0.0003, train/strict_loss: 0.0018, val/iit_loss: 0.0283, val/IIA: 98.99%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 21: lr: 8.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0363, train/behavior_loss: 0.0003, train/strict_loss: 0.0021, val/iit_loss: 0.0222, val/IIA: 99.25%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 22: lr: 8.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0282, train/behavior_loss: 0.0003, train/strict_loss: 0.0018, val/iit_loss: 0.0248, val/IIA: 99.18%, val/accuracy: 100.00%, val/strict_accuracy: 99.95%\n",
      "Epoch 23: lr: 8.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0305, train/behavior_loss: 0.0003, train/strict_loss: 0.0016, val/iit_loss: 0.0245, val/IIA: 99.16%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 24: lr: 8.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0253, train/behavior_loss: 0.0002, train/strict_loss: 0.0013, val/iit_loss: 0.0261, val/IIA: 99.06%, val/accuracy: 99.97%, val/strict_accuracy: 99.91%\n",
      "Epoch 25: lr: 8.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0280, train/behavior_loss: 0.0002, train/strict_loss: 0.0015, val/iit_loss: 0.0261, val/IIA: 99.14%, val/accuracy: 100.00%, val/strict_accuracy: 99.91%\n",
      "Epoch 26: lr: 7.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0231, train/behavior_loss: 0.0001, train/strict_loss: 0.0012, val/iit_loss: 0.0262, val/IIA: 99.00%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 27: lr: 7.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0242, train/behavior_loss: 0.0001, train/strict_loss: 0.0011, val/iit_loss: 0.0302, val/IIA: 98.95%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 28: lr: 7.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0274, train/behavior_loss: 0.0001, train/strict_loss: 0.0013, val/iit_loss: 0.0266, val/IIA: 99.14%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 29: lr: 7.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0182, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0280, val/IIA: 99.03%, val/accuracy: 99.99%, val/strict_accuracy: 99.95%\n",
      "Epoch 30: lr: 7.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0233, train/behavior_loss: 0.0001, train/strict_loss: 0.0010, val/iit_loss: 0.0168, val/IIA: 99.41%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 31: lr: 7.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0182, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0176, val/IIA: 99.39%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 32: lr: 7.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0205, train/behavior_loss: 0.0001, train/strict_loss: 0.0009, val/iit_loss: 0.0161, val/IIA: 99.45%, val/accuracy: 100.00%, val/strict_accuracy: 99.98%\n",
      "Epoch 33: lr: 7.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0192, train/behavior_loss: 0.0001, train/strict_loss: 0.0007, val/iit_loss: 0.0272, val/IIA: 99.02%, val/accuracy: 100.00%, val/strict_accuracy: 99.96%\n",
      "Epoch 34: lr: 7.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0227, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0186, val/IIA: 99.32%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 35: lr: 7.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0139, train/behavior_loss: 0.0001, train/strict_loss: 0.0007, val/iit_loss: 0.0126, val/IIA: 99.59%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 36: lr: 7.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0165, train/behavior_loss: 0.0000, train/strict_loss: 0.0007, val/iit_loss: 0.0225, val/IIA: 99.14%, val/accuracy: 100.00%, val/strict_accuracy: 99.97%\n",
      "Epoch 37: lr: 7.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0168, train/behavior_loss: 0.0001, train/strict_loss: 0.0008, val/iit_loss: 0.0154, val/IIA: 99.48%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 38: lr: 6.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0159, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0140, val/IIA: 99.54%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 39: lr: 6.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0111, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0104, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 40: lr: 6.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0118, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0156, val/IIA: 99.48%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 41: lr: 6.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0162, train/behavior_loss: 0.0001, train/strict_loss: 0.0007, val/iit_loss: 0.0112, val/IIA: 99.61%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 42: lr: 6.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0126, train/behavior_loss: 0.0001, train/strict_loss: 0.0006, val/iit_loss: 0.0125, val/IIA: 99.56%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 43: lr: 6.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0140, train/behavior_loss: 0.0000, train/strict_loss: 0.0006, val/iit_loss: 0.0123, val/IIA: 99.58%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 44: lr: 6.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0110, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0100, val/IIA: 99.65%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 45: lr: 6.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0110, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0142, val/IIA: 99.49%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 46: lr: 6.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0121, train/behavior_loss: 0.0000, train/strict_loss: 0.0005, val/iit_loss: 0.0125, val/IIA: 99.54%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 47: lr: 6.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0123, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0095, val/IIA: 99.68%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 48: lr: 6.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0088, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0125, val/IIA: 99.56%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 49: lr: 6.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0119, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0181, val/IIA: 99.41%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 50: lr: 6.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0101, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0124, val/IIA: 99.55%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 51: lr: 5.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0126, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0090, val/IIA: 99.69%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 52: lr: 5.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0104, train/behavior_loss: 0.0000, train/strict_loss: 0.0004, val/iit_loss: 0.0194, val/IIA: 99.34%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 53: lr: 5.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0137, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0109, val/IIA: 99.66%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 54: lr: 5.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0073, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0091, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 55: lr: 5.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0080, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0084, val/IIA: 99.70%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 56: lr: 5.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0072, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 57: lr: 5.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0078, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0062, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 58: lr: 5.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0079, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0109, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 59: lr: 5.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0060, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0077, val/IIA: 99.74%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 60: lr: 5.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0065, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0105, val/IIA: 99.67%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 61: lr: 5.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0066, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 62: lr: 5.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0060, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 63: lr: 4.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0075, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0066, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 64: lr: 4.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0096, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0074, val/IIA: 99.75%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 65: lr: 4.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0077, train/behavior_loss: 0.0000, train/strict_loss: 0.0003, val/iit_loss: 0.0087, val/IIA: 99.72%, val/accuracy: 100.00%, val/strict_accuracy: 99.99%\n",
      "Epoch 66: lr: 4.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0060, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0047, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 67: lr: 4.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0064, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0047, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 68: lr: 4.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0055, val/IIA: 99.83%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 69: lr: 4.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0054, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0072, val/IIA: 99.78%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 70: lr: 4.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0068, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0082, val/IIA: 99.77%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 71: lr: 4.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0058, train/behavior_loss: 0.0000, train/strict_loss: 0.0002, val/iit_loss: 0.0066, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 72: lr: 4.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0066, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0059, val/IIA: 99.81%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 73: lr: 4.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0059, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0043, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 74: lr: 4.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0052, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0057, val/IIA: 99.79%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 75: lr: 4.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0033, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0044, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 76: lr: 3.92e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0044, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 77: lr: 3.84e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0041, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0039, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 78: lr: 3.76e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0031, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0052, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 79: lr: 3.68e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0051, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0045, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 80: lr: 3.60e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0055, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0040, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 81: lr: 3.52e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0047, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0043, val/IIA: 99.86%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 82: lr: 3.44e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0046, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0049, val/IIA: 99.85%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 83: lr: 3.36e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0035, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0061, val/IIA: 99.76%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 84: lr: 3.28e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0050, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0034, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 85: lr: 3.20e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0032, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0037, val/IIA: 99.87%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 86: lr: 3.12e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0048, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0038, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 87: lr: 3.04e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0037, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0030, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 88: lr: 2.96e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0032, val/IIA: 99.90%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 89: lr: 2.88e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0023, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0036, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 90: lr: 2.80e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0030, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 91: lr: 2.72e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0034, val/IIA: 99.88%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 92: lr: 2.64e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0031, val/IIA: 99.89%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 93: lr: 2.56e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0044, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0061, val/IIA: 99.82%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 94: lr: 2.48e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0026, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0029, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 95: lr: 2.40e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0027, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0031, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 96: lr: 2.32e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0001, val/iit_loss: 0.0026, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 97: lr: 2.24e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0034, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0026, val/IIA: 99.92%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 98: lr: 2.16e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0018, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0026, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 99: lr: 2.08e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0039, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0028, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n",
      "Epoch 100: lr: 2.00e-04, iit_weight: 1.00e+00, behavior_weight: 4.00e-01, strict_weight: 4.00e-01, train/iit_loss: 0.0021, train/behavior_loss: 0.0000, train/strict_loss: 0.0000, val/iit_loss: 0.0028, val/IIA: 99.91%, val/accuracy: 100.00%, val/strict_accuracy: 100.00%\n"
     ]
    }
   ],
   "source": [
    "ll_model = poly_hl_model.get_ll_model().to(poly_hl_model.device)\n",
    "ll_model.device = poly_hl_model.device\n",
    "corr = poly_hl_model.get_correspondence()\n",
    "for k, v in corr.items():\n",
    "    print(k, v)\n",
    "train_set, test_set = poly_dataset.get_IIT_train_test_set()\n",
    "model_pair = StrictIITModelPair(hl_model=poly_hl_model, ll_model=ll_model, corr=corr, training_args=training_args)\n",
    "print(model_pair.nodes_not_in_circuit)\n",
    "model_pair.train(\n",
    "    train_set=train_set,\n",
    "    test_set=test_set,\n",
    "    epochs=n_epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('saved_poly_models/cases_0+1+2+3')"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "save_poly_model_to_dir(ll_model, poly_hl_model, f\"./saved_poly_models/cases_0+1+2+3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Push to HF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "593b79d47bce4f5093b6dd1d481a24bb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Upload 11 LFS files:   0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44dca3123f9e49ee9011e685a9ed2435",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed18f7e471474d1cb9818abaf3cbab03",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/68.4k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3aa93dc8b8ed41bb88cfb8004eeb2adc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/46.5k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3c8926dc9ee3412c919538949346698f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/68.4k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a501bb2a002e4145b46068d5262d6444",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "12221c1ac8254a55bbb519259508baa8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a7266af8b66449c9950438d7b492cd2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f20d7710494b43da8e44f5b2307ec60f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/68.3k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4df9de53f0d34925b73e35e4d1037e1d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e813148116846088f6dfc53e0a57396",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54d5fa36be1545dcab4380e07885b7ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "save_to_hf(local_dir=\"saved_poly_models\", message=\"pushes all polysemantic models\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "circuits_bench",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}