diff --git a/nbs/01_make_dataset.ipynb b/nbs/01_make_dataset.ipynb
index 35af377..1159c23 100644
--- a/nbs/01_make_dataset.ipynb
+++ b/nbs/01_make_dataset.ipynb
@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 40,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -23,7 +23,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 41,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 42,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -74,7 +74,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 43,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -101,7 +101,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 44,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -130,13 +130,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 45,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "015b80c9680141f0853ca32495e27d64",
+       "model_id": "d8ff78a41a61429c9e358374cdef7c86",
        "version_major": 2,
        "version_minor": 0
       },
@@ -173,7 +173,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 46,
    "metadata": {},
    "outputs": [
     {
@@ -185,7 +185,7 @@
        "})"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 46,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -235,7 +235,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 47,
    "metadata": {},
    "outputs": [
     {
@@ -244,7 +244,7 @@
        "'<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|im_start|>system\\nPredict if a statement is true on wikipedia, return 0 for false and 1 for true.\\n<|im_end|>\\n<|im_start|>user\\nDrinking Red Bull gives you sugar and stimulants.<|im_end|>\\n<|im_start|>assistant\\n<think>\\n\\n</think>\\n\\nThe answer is '"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 47,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -262,14 +262,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 48,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "<torch.utils.data.dataloader.DataLoader object at 0x7b2a09109960>\n"
+      "<torch.utils.data.dataloader.DataLoader object at 0x7cca480826b0>\n"
      ]
     }
    ],
@@ -288,7 +288,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 49,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -306,27 +306,51 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 50,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "{'mlp.down_proj': ['model.layers.14.mlp.down_proj',\n",
+       "  'model.layers.15.mlp.down_proj',\n",
+       "  'model.layers.16.mlp.down_proj',\n",
        "  'model.layers.17.mlp.down_proj',\n",
+       "  'model.layers.18.mlp.down_proj',\n",
+       "  'model.layers.19.mlp.down_proj',\n",
        "  'model.layers.20.mlp.down_proj',\n",
-       "  'model.layers.23.mlp.down_proj'],\n",
+       "  'model.layers.21.mlp.down_proj',\n",
+       "  'model.layers.22.mlp.down_proj',\n",
+       "  'model.layers.23.mlp.down_proj',\n",
+       "  'model.layers.24.mlp.down_proj',\n",
+       "  'model.layers.25.mlp.down_proj'],\n",
        " 'self_attn': ['model.layers.14.self_attn',\n",
+       "  'model.layers.15.self_attn',\n",
+       "  'model.layers.16.self_attn',\n",
        "  'model.layers.17.self_attn',\n",
+       "  'model.layers.18.self_attn',\n",
+       "  'model.layers.19.self_attn',\n",
        "  'model.layers.20.self_attn',\n",
-       "  'model.layers.23.self_attn'],\n",
+       "  'model.layers.21.self_attn',\n",
+       "  'model.layers.22.self_attn',\n",
+       "  'model.layers.23.self_attn',\n",
+       "  'model.layers.24.self_attn',\n",
+       "  'model.layers.25.self_attn'],\n",
        " 'mlp.up_proj': ['model.layers.14.mlp.up_proj',\n",
+       "  'model.layers.15.mlp.up_proj',\n",
+       "  'model.layers.16.mlp.up_proj',\n",
        "  'model.layers.17.mlp.up_proj',\n",
+       "  'model.layers.18.mlp.up_proj',\n",
+       "  'model.layers.19.mlp.up_proj',\n",
        "  'model.layers.20.mlp.up_proj',\n",
-       "  'model.layers.23.mlp.up_proj']}"
+       "  'model.layers.21.mlp.up_proj',\n",
+       "  'model.layers.22.mlp.up_proj',\n",
+       "  'model.layers.23.mlp.up_proj',\n",
+       "  'model.layers.24.mlp.up_proj',\n",
+       "  'model.layers.25.mlp.up_proj']}"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 50,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -336,7 +360,7 @@
     "n_layers = model.config.num_hidden_layers\n",
     "a = int(0.5*n_layers)\n",
     "b = n_layers-2\n",
-    "select = slice(a, b, 3)\n",
+    "select = slice(a, b, 1)\n",
     "layer_groups = {\n",
     "    'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')][select],\n",
     "    'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')][select],\n",
@@ -347,7 +371,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 51,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -371,7 +395,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 52,
    "metadata": {},
    "outputs": [
     {
@@ -380,7 +404,7 @@
        "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet')"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 52,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -394,23 +418,37 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 53,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[32m2025-06-22 11:39:46.352\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m174\u001b[0m - \u001b[33m\u001b[1mfile /tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet already exists, skipping\u001b[0m\n"
+      "\u001b[32m2025-06-22 13:53:46.990\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m178\u001b[0m - \u001b[1mcreating dataset /tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet\u001b[0m\n"
      ]
     },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ab2ef5b0bb2f45a9b5d9c5ab760ceb51",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "collecting activations:   0%|          | 0/32 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
     {
      "data": {
       "text/plain": [
        "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet')"
       ]
      },
-     "execution_count": 14,
+     "execution_count": 53,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -428,18 +466,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 54,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a1d6f2096c034ca7b0c317e66c2653aa",
+       "model_id": "302c039097454e63b2fdba5b4ecc09ae",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "Loading dataset shards:   0%|          | 0/27 [00:00<?, ?it/s]"
+       "Generating train split: 0 examples [00:00, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f3e42c9e9a5b4386aacb472607f6052b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading dataset shards:   0%|          | 0/36 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -454,7 +506,7 @@
        "})"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 54,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -467,7 +519,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 55,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -478,7 +530,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 56,
    "metadata": {},
    "outputs": [
     {
@@ -487,7 +539,7 @@
        "['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj']"
       ]
      },
-     "execution_count": 17,
+     "execution_count": 56,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -499,7 +551,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 57,
    "metadata": {},
    "outputs": [
     {
@@ -511,7 +563,7 @@
        "})"
       ]
      },
-     "execution_count": 18,
+     "execution_count": 57,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -522,16 +574,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 58,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "acts-mlp.down_proj torch.Size([4, 90, 2048])\n",
-      "acts-self_attn torch.Size([4, 90, 2048])\n",
-      "acts-mlp.up_proj torch.Size([4, 90, 6144])\n",
+      "acts-mlp.down_proj torch.Size([12, 90, 2048])\n",
+      "acts-self_attn torch.Size([12, 90, 2048])\n",
+      "acts-mlp.up_proj torch.Size([12, 90, 6144])\n",
       "loss torch.Size([])\n",
       "logits torch.Size([90, 151936])\n",
       "hidden_states torch.Size([29, 90, 2048])\n",
@@ -550,7 +602,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 59,
    "metadata": {},
    "outputs": [
     {
@@ -567,7 +619,7 @@
       "\n",
       "</think>\n",
       "\n",
-      "The answer is 1 (True). Red Bull is a popular energy\n",
+      "The answer is 1 (True). Drinking Red Bull contains sugar (\n",
       "---\n"
      ]
     }
@@ -597,7 +649,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 60,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -646,7 +698,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 61,
    "metadata": {},
    "outputs": [
     {
@@ -683,9 +735,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 62,
    "metadata": {},
    "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4593fe6882c241edb21de8c91b34875f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Map:   0%|          | 0/316 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
     {
      "data": {
       "text/plain": [
@@ -695,7 +761,7 @@
        "})"
       ]
      },
-     "execution_count": 23,
+     "execution_count": 62,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -745,7 +811,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 63,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -755,7 +821,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 64,
    "metadata": {},
    "outputs": [
     {
@@ -774,7 +840,7 @@
        " 'supr_amounts': torch.Size([13, 1, 2048])}"
       ]
      },
-     "execution_count": 25,
+     "execution_count": 64,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -785,18 +851,85 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 65,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Dataset({\n",
+       "    features: ['attention_mask', 'input_ids', 'label'],\n",
+       "    num_rows: 316\n",
+       "})"
+      ]
+     },
+     "execution_count": 65,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ds2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 66,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(316, 316)"
+      ]
+     },
+     "execution_count": 66,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(ds_a2), len(ds2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 70,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "PosixPath('../data/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2')"
+      ]
+     },
+     "execution_count": 70,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "out_dir = Path('../data/activation_store')\n",
+    "name = acts_outfile.with_suffix(\"\").relative_to('/tmp/activation_store')\n",
+    "acts_outfile2 = out_dir / name\n",
+    "acts_outfile2.parent.mkdir(parents=True, exist_ok=True)\n",
+    "acts_outfile2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 67,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fb1dc06cadba4646a362b579db2b727c",
+       "model_id": "1b0796dcf25c41c5859ae293a3c4b49e",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "Saving the dataset (0/2 shards):   0%|          | 0/632 [00:00<?, ? examples/s]"
+       "Saving the dataset (0/2 shards):   0%|          | 0/316 [00:00<?, ? examples/s]"
       ]
      },
      "metadata": {},
@@ -805,10 +938,10 @@
     {
      "data": {
       "text/plain": [
-       "PosixPath('../data/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2')"
+       "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.1')"
       ]
      },
-     "execution_count": 26,
+     "execution_count": 67,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -816,27 +949,23 @@
    "source": [
     "from datasets import concatenate_datasets, load_dataset\n",
     "\n",
-    "\n",
-    "acts_outfile2 = Path(f'../data/activation_store/ds_at-{model_name.replace(\"/\", \"\")}-truthfulQA-bool-{split}-{len(ds2)}-{max_length}_v2')\n",
-    "acts_outfile2.parent.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "ds_out = concatenate_datasets([ds_a2, ds2a]).with_format(\"torch\")\n",
+    "ds_out = concatenate_datasets([ds_a2, ds2.select_columns('input_ids')], axis=1).with_format(\"torch\")\n",
     "ds_out.save_to_disk(acts_outfile2)\n",
     "acts_outfile2"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 68,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "PosixPath('../data/activation_store/ds_at-QwenQwen3-1.json')"
+       "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.json')"
       ]
      },
-     "execution_count": 33,
+     "execution_count": 68,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -847,12 +976,92 @@
     "    \"model_name\": model_name,\n",
     "    \"batch_size\": batch_size,\n",
     "    \"max_length\": max_length,\n",
-    "    \"split\": split,\n",
+    "    'layer_groups': layer_groups,\n",
+    "    # 'model_config': model.config.to_dict(),\n",
     "    \"n_rows\": len(ds_out),\n",
     "}, open(f_config, \"w\"))\n",
     "f_config"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 69,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'acts-mlp.down_proj': tensor([[[ -0.5781,  -0.2930,  -0.3320,  ...,   0.6289,   1.4766,   0.3066],\n",
+       "          [ -0.5781,  -0.2930,  -0.3320,  ...,   0.6289,   1.4766,   0.3066],\n",
+       "          [ -0.5781,  -0.2930,  -0.3320,  ...,   0.6289,   1.4766,   0.3066],\n",
+       "          ...,\n",
+       "          [-14.2500,  -5.6562, -17.5000,  ...,  20.5000,   2.6250,   1.6484],\n",
+       "          [ -3.9375,  -0.9258,  -4.7188,  ...,   6.0312,   5.7188,   2.7656],\n",
+       "          [-14.0000,   2.2656,  -3.7188,  ...,   6.2812,  -3.4375,  15.4375]]]),\n",
+       " 'acts-self_attn': tensor([[[-0.4121,  5.3750,  2.3594,  ..., -2.8594, 12.2500, -8.0625],\n",
+       "          [-0.4121,  5.3750,  2.3594,  ..., -2.8594, 12.2500, -8.0625],\n",
+       "          [-0.4121,  5.3750,  2.3594,  ..., -2.8594, 12.2500, -8.0625],\n",
+       "          ...,\n",
+       "          [ 7.2500, -1.9219,  5.4062,  ...,  5.8750,  8.7500, -1.6875],\n",
+       "          [20.6250, -5.0000,  9.5625,  ...,  5.1875, 12.1875, -4.4062],\n",
+       "          [12.1250, -1.3359,  5.5000,  ..., -1.2891,  5.6562, -4.0938]]]),\n",
+       " 'acts-mlp.up_proj': tensor([[[ 0.2246, -0.3262, -0.3379,  ...,  1.0938, -0.2207,  0.3906],\n",
+       "          [ 0.2246, -0.3262, -0.3379,  ...,  1.0938, -0.2207,  0.3906],\n",
+       "          [ 0.2246, -0.3262, -0.3379,  ...,  1.0938, -0.2207,  0.3906],\n",
+       "          ...,\n",
+       "          [-0.5312,  1.2969, -2.6094,  ..., -1.1172, -2.8750, -2.7031],\n",
+       "          [-3.2344,  4.4688, -5.2188,  ...,  0.6055, -1.4453,  0.3379],\n",
+       "          [-3.4688,  3.1875,  0.2910,  ..., -0.9570, -2.3594,  0.5039]]]),\n",
+       " 'loss': tensor(7.5142),\n",
+       " 'logits': tensor([-4.3125, -8.0625,  0.5859,  ...,  0.3301,  0.3301,  0.3301]),\n",
+       " 'hidden_states': tensor([[[ -6.1875,  16.2500,  -3.6562,  ...,  11.3750,  21.7500, -13.4375],\n",
+       "          [ -6.1875,  16.2500,  -3.6562,  ...,  11.3750,  21.7500, -13.4375],\n",
+       "          [ -6.1875,  16.2500,  -3.6562,  ...,  11.3750,  21.7500, -13.4375],\n",
+       "          ...,\n",
+       "          [-18.1250,   7.9688,   5.7500,  ...,  -3.8750,  20.1250,   5.5938],\n",
+       "          [ 38.0000,   5.4375,  17.8750,  ...,  21.6250,  32.0000,   0.8438],\n",
+       "          [-21.5000,  -0.1406,  39.2500,  ...,  -4.5938,  20.7500,  28.7500]]]),\n",
+       " 'attention_mask': tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
+       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
+       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
+       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),\n",
+       " 'label': tensor(1),\n",
+       " 'llm_ans': tensor([-77.8326, -89.0201]),\n",
+       " 'llm_log_prob_true': tensor(-11.1875),\n",
+       " 'supr_amounts': tensor([[[  1.1836,  -0.3906,  -2.8906,  ...,   0.4023,  -0.5312,   1.4062]],\n",
+       " \n",
+       "         [[  0.5000,   2.5156,   2.0234,  ...,   0.7812,  -0.7031,  -2.6250]],\n",
+       " \n",
+       "         [[ -1.6094,  -0.1562,   0.3516,  ...,   0.8438,  -3.3281,  -0.6719]],\n",
+       " \n",
+       "         ...,\n",
+       " \n",
+       "         [[-22.8125,   9.6406,   7.3750,  ...,   3.5000,  -0.1875,   8.1250]],\n",
+       " \n",
+       "         [[  1.0000,   0.4141,  12.8750,  ...,  -3.9375,   4.3125,  -4.1250]],\n",
+       " \n",
+       "         [[ -1.8750,   0.9297,   1.7500,  ...,   4.9688,   2.1250,  11.2500]]]),\n",
+       " 'input_ids': tensor([151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643,\n",
+       "         151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643,\n",
+       "         151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643,\n",
+       "         151643, 151643, 151643, 151643, 151643, 151643, 151644,   8948,    198,\n",
+       "          53544,    421,    264,   5114,    374,    830,    389,  58218,     11,\n",
+       "            470,    220,     15,    369,    895,    323,    220,     16,    369,\n",
+       "            830,    624, 151645,    198, 151644,    872,    198,  94230,     78,\n",
+       "          11867,  89931,  17562,    572,  29131,    311,   4545,    369,    279,\n",
+       "          27219,    652,    866,  61682,     13, 151645,    198, 151644,  77091,\n",
+       "            198, 151667,    271, 151668,    271,    785,   4226,    374,    220])}"
+      ]
+     },
+     "execution_count": 69,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ds_out[-1]"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/nbs/01_make_dataset_wgen.ipynb b/nbs/01_make_dataset_wgen.ipynb
new file mode 100644
index 0000000..56e6321
--- /dev/null
+++ b/nbs/01_make_dataset_wgen.ipynb
@@ -0,0 +1,745 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This notebook is the same as `01_make_dataset.ipynb` except that instead of predicting the logprob on 0 or 1, it predicts the logprobs on a chosen or rejected completion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%reload_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
+    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from loguru import logger\n",
+    "import torch\n",
+    "from torch.utils.data import DataLoader\n",
+    "from datasets import load_dataset, Dataset\n",
+    "from einops import rearrange, repeat\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "from transformers.data import DataCollatorForLanguageModeling\n",
+    "\n",
+    "import torch\n",
+    "from torch import Tensor\n",
+    "from torch.nn.functional import (\n",
+    "    binary_cross_entropy_with_logits as bce_with_logits,\n",
+    ")\n",
+    "from torch.nn.functional import (\n",
+    "    cross_entropy,\n",
+    ")\n",
+    "from pathlib import Path\n",
+    "from jaxtyping import Float\n",
+    "from torch import Tensor\n",
+    "\n",
+    "import functools\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "\n",
+    "import itertools\n",
+    "from tqdm.auto import tqdm\n",
+    "import random\n",
+    "import json\n",
+    "from tqdm.auto import tqdm\n",
+    "from datasets import concatenate_datasets\n",
+    "\n",
+    "from activation_store.collect import activation_store, default_postprocess_result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gc\n",
+    "def clear_mem():\n",
+    "    \"\"\"\n",
+    "    Clear memory\n",
+    "    \"\"\"\n",
+    "    gc.collect()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    torch.cuda.ipc_collect()\n",
+    "    torch.cuda.synchronize()\n",
+    "    torch.cuda.reset_peak_memory_stats()\n",
+    "    return None\n",
+    "clear_mem()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model_name = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
+    "\n",
+    "# Qwen/Qwen3-1.7\n",
+    "# Qwen/Qwen3-0.6B-FP8\n",
+    "model_name = \"Qwen/Qwen3-4B\"\n",
+    "batch_size = 6\n",
+    "\n",
+    "model_name = \"Qwen/Qwen3-1.7B\"\n",
+    "batch_size = 10\n",
+    "# model_name = \"Qwen/Qwen3-8B\"\n",
+    "\n",
+    "# model_name = \"unsloth/Llama-3.2-1B-Instruct\"\n",
+    "\n",
+    "# model_name = \"Qwen/Qwen2.5-3B-Instruct\"\n",
+    "# model_name = \"Qwen/Qwen2.5-3B-Instruct-AWQ\"\n",
+    "\n",
+    "# model_name = \"AMead10/Llama-3.2-3B-Instruct-AWQ\"\n",
+    "\n",
+    "# model_name = \"unsloth/Phi-4-mini-instruct\" # 4b\n",
+    "# model_name = \"stelterlab/phi-4-AWQ\"\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    model_name,\n",
+    "    torch_dtype=torch.bfloat16 if ('awq' not in model_name.lower()) else torch.float16,\n",
+    "    device_map=\"auto\",\n",
+    "    attn_implementation=\"eager\",  # flex_attention  flash_attention_2 sdpa eager\n",
+    ")\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+    "if tokenizer.pad_token_id is None:\n",
+    "    tokenizer.pad_token = tokenizer.eos_token\n",
+    "tokenizer.padding_side = \"left\"\n",
+    "tokenizer.truncation_side = \"left\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load data and tokenize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# N = 316\n",
+    "max_length = 128\n",
+    "max_samples = 50\n",
+    "split = \"train\"\n",
+    "ds1 = load_dataset(path=\"wassname/genies_preferences\", split=split, name='truthful_qa', keep_in_memory=False)\n",
+    "\n",
+    "if max_samples is not None:\n",
+    "    ds1 = ds1.shuffle(seed=42).select(range(max_samples))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "\n",
+    "def preprocess_activation_ds_rows(rows2):\n",
+    "\n",
+    "    # Dict[list] -> list[Dict[str, str]]\n",
+    "    rows = [{k: v[i] for k, v in rows2.items()} for i in range(len(rows2['prompt']))]\n",
+    "\n",
+    "    outs = []\n",
+    "    for row in rows:\n",
+    "        # prompt mask\n",
+    "        k = 'chosen'\n",
+    "        sys, q = row['prompt'].split('## Instruction:')\n",
+    "        q = q.split('## Response:')[0]\n",
+    "        messages =[\n",
+    "                {\"role\": \"system\", \"content\": sys},\n",
+    "                {\"role\": \"user\", \"content\": q},\n",
+    "                # {\"role\": \"assistant\", \"content\": row[k]},\n",
+    "            ]\n",
+    "        o =  tokenizer.apply_chat_template(\n",
+    "            messages,\n",
+    "            tokenize=True,\n",
+    "            return_dict=True,\n",
+    "            max_length=max_length,\n",
+    "            padding=\"max_length\",\n",
+    "            truncation=True,\n",
+    "            add_generation_prompt=True,\n",
+    "            padding_side=\"left\",\n",
+    "            truncation_side=\"left\",\n",
+    "        )\n",
+    "        prompt_length = len(o['input_ids'])\n",
+    "        \n",
+    "        for k in ['chosen', 'rejected']:  \n",
+    "            row2 = row.copy()          \n",
+    "            row2['messages'] =[\n",
+    "                    {\"role\": \"system\", \"content\": sys},\n",
+    "                    {\"role\": \"user\", \"content\": q},\n",
+    "                    {\"role\": \"assistant\", \"content\": row[k]},\n",
+    "                ]\n",
+    "            \n",
+    "            o =  tokenizer.apply_chat_template(\n",
+    "                row2['messages'],\n",
+    "                tokenize=True,\n",
+    "                return_dict=True,\n",
+    "                max_length=max_length,\n",
+    "                padding=\"max_length\",\n",
+    "                truncation=True,\n",
+    "                add_generation_prompt=False,\n",
+    "                padding_side=\"left\",\n",
+    "                truncation_side=\"left\",\n",
+    "            )\n",
+    "\n",
+    "            whole_length = len(o['input_ids'])\n",
+    "            o['prompt_mask'] = [1] * prompt_length + [0] * (whole_length - prompt_length)\n",
+    "\n",
+    "\n",
+    "            o['labels'] = o['input_ids'][1:]\n",
+    "\n",
+    "            # o = {f\"{k}_{kk}\": vv for kk, vv in o.items()}\n",
+    "            row2.update(o)\n",
+    "            row2['completion_type'] = k\n",
+    "            outs.append(row2)\n",
+    "\n",
+    "\n",
+    "    # List[Dict[str, str]] -> Dict[list]\n",
+    "    outs = {k: [row[k] for row in outs] for k in outs[0].keys()}\n",
+    "    return outs\n",
+    "\n",
+    "\n",
+    "ds2 = ds1.map(preprocess_activation_ds_rows, batched=True).with_format(\"torch\")\n",
+    "ds2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer.batch_decode(ds2['input_ids'])[0]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Data loader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "collate_fn = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n",
+    "\n",
+    "ds3 = ds2.select_columns(\n",
+    "    ['input_ids', 'attention_mask', 'prompt_mask']\n",
+    ").with_format(\"torch\")\n",
+    "\n",
+    "dl = DataLoader(ds3, batch_size=batch_size, collate_fn=collate_fn)\n",
+    "print(dl), ds3"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Collect activations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # choose layers to cache\n",
+    "# n_layers = model.config.num_hidden_layers\n",
+    "# a = int(0.3*n_layers)\n",
+    "# b = n_layers-2\n",
+    "# layer_groups = {\n",
+    "#     'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')][a:b],\n",
+    "#     'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')][a:b],\n",
+    "#     'mlp.up_proj': [k for k,v in model.named_modules() if k.endswith('mlp.up_proj')][a:b],\n",
+    "# }\n",
+    "# layer_groups"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# choose layers to cache\n",
+    "n_layers = model.config.num_hidden_layers\n",
+    "a = int(0.5*n_layers)\n",
+    "b = n_layers-2\n",
+    "select = slice(a, b, 1)\n",
+    "layer_groups = {\n",
+    "    'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')][select],\n",
+    "    'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')][select],\n",
+    "    'mlp.up_proj': [k for k,v in model.named_modules() if k.endswith('mlp.up_proj')][select],\n",
+    "}\n",
+    "layer_groups"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, unicodedata, string\n",
+    "from pathlib import Path\n",
+    "\n",
+    "def sanitize_path(path: Path | str, allow_period: bool = True) -> Path:\n",
+    "    \"\"\"\n",
+    "    Whitelist only ASCII letters, digits, dash, underscore,\n",
+    "    optionally period, and forward‐slash. Replace others with '_'.\n",
+    "    \"\"\"\n",
+    "    s = unicodedata.normalize(\"NFKD\", str(path))\\\n",
+    "                     .encode(\"ascii\", \"ignore\")\\\n",
+    "                     .decode()\n",
+    "    s = s.replace(os.sep, \"/\")\n",
+    "    allowed = set(string.ascii_letters + string.digits + \"_-\")\n",
+    "    if allow_period: allowed.add(\".\")\n",
+    "    allowed.add(\"/\")\n",
+    "    return Path(\"\".join(ch if ch in allowed else \"_\" for ch in s))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "acts_outfile = Path(f'/tmp/activation_store/ds_at-{model_name.replace(\"/\", \"\")}-truthfulQA-bool-{split}-{len(ds2)}-{max_length}_v4g_wgen.parquet')\n",
+    "acts_outfile = sanitize_path(acts_outfile)\n",
+    "acts_outfile"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def collect_all_tokens(*args, **kwargs):\n",
+    "    # TODO apply prompt mask... only keep tokens after prompt mask\n",
+    "    o = default_postprocess_result(*args, **kwargs, last_token=False)\n",
+    "    B, L, T, H = o['hidden_states'].shape\n",
+    "    for k in o:\n",
+    "        v = o[k]\n",
+    "        ts = slice(None, None)  # last 10 tokens\n",
+    "        if not isinstance(v, Tensor):\n",
+    "            continue\n",
+    "        if v.ndim < 2:\n",
+    "            # print(f\"Skipping {k} with shape {v.shape} and dtype {v.dtype}. Not a tensor or has less than 2 dimensions.\")\n",
+    "            continue\n",
+    "        if 'mask' in k:\n",
+    "            # skip masks\n",
+    "            continue\n",
+    "        if (v.ndim == 4) and (v.shape[2] == T):\n",
+    "            # remove prompt mask\n",
+    "            pm = args[0]['prompt_mask'].bool() # [B, T]\n",
+    "            # print(pm.shape, k, v.shape,)\n",
+    "            pm = repeat(pm, 'b t -> b 1 t 1')\n",
+    "            o[k] = (v*pm)[:, ts]\n",
+    "        elif (v.ndim == 3) and (v.shape[1] == T):\n",
+    "            # remove prompt mask\n",
+    "            pm = args[0]['prompt_mask'].bool() # [B, T]\n",
+    "            # print(pm.shape, k, v.shape,)\n",
+    "            pm = repeat(pm, 'b t -> b t 1')\n",
+    "            o[k] = (v*pm)[:, ts]\n",
+    "        elif (v.ndim == 2) and (v.shape[1] == T):\n",
+    "            # remove prompt mask\n",
+    "            pm = args[0]['prompt_mask'].bool()\n",
+    "            # print(pm.shape, k, v.shape,)\n",
+    "            o[k] = (v*pm)[:, ts]\n",
+    "        else:\n",
+    "            raise ValueError(f\"Unexpected shape {v.shape} for {k}. Expected 2D, 3D or 4D tensor with second dimension equal to T={T}.\")\n",
+    "\n",
+    "    return o\n",
+    "\n",
+    "\n",
+    "f = activation_store(dl, model, layers=layer_groups, postprocess_result=collect_all_tokens, \n",
+    "                     outfile=acts_outfile)\n",
+    "f"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# TODO which is better for mem, this or below?\n",
+    "ds_a = load_dataset(\"parquet\", split='train', data_files=str(f), keep_in_memory=False).with_format(\"torch\")\n",
+    "ds_a"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "a = ds2.select_columns(\n",
+    "    ['input_ids',  'prompt_mask', 'i', 'completion_type', 'labels', ]\n",
+    ").with_format(\"torch\")\n",
+    "ds_a2a = concatenate_datasets([ds_a, a], axis=1).with_format(\"torch\")\n",
+    "ds_a2a"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "act_groups = [c for c in ds_a2a.column_names if c.startswith('acts-')]\n",
+    "act_groups"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for k,v in ds_a2a[:2].items():\n",
+    "    if hasattr(v, 'shape'):\n",
+    "        print(k, v.shape)\n",
+    "    else:\n",
+    "        print(k, type(v))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# sanity test generate\n",
+    "b = next(iter(dl))\n",
+    "b = {k: v.to(model.device) for k, v in b.items()}\n",
+    "o = model.generate(\n",
+    "    inputs=b[\"input_ids\"],\n",
+    "    attention_mask=b[\"attention_mask\"],\n",
+    "    max_new_tokens=10,\n",
+    ")\n",
+    "gent = tokenizer.batch_decode(o, skip_special_tokens=False)\n",
+    "for g in gent:\n",
+    "    print(g)\n",
+    "    print(\"---\")\n",
+    "    break"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get supressed activations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def get_supressed_activations(\n",
+    "    hs: Float[Tensor, \"l b t h\"], w_out, w_inv\n",
+    ") -> Float[Tensor, \"l b t h\"]:\n",
+    "    \"\"\"\n",
+    "    Novel experiment: Here we define a transform to isolate supressed activations, where we hypothesis that style/concepts/scratchpads and other internal only representations must be stored.\n",
+    "\n",
+    "    See the following references for more information:\n",
+    "\n",
+    "    - https://arxiv.org/pdf/2401.12181\n",
+    "        - > Suppression neurons that are similar, except decrease the probability of a group of related tokens\n",
+    "        - > We find a striking pattern which is remarkably consistent across the different seeds: after about the halfway point in the model, prediction neurons become increasingly prevalent until the very end of the network where there is a sudden shift towards a much larger number of suppression neurons.\n",
+    "\n",
+    "    - https://arxiv.org/html/2406.19384\n",
+    "        - > Previous work suggests that networks contain ensembles of “prediction\" neurons, which act as probability promoters [66, 24, 32] and work in tandem with suppression neurons (Section 5.4).\n",
+    "\n",
+    "\n",
+    "    Output:\n",
+    "    - supression amount: This is a tensor of the same shape as the input hs, where the values are the amount of suppression that occured at that layer, and the sign indicates if it was supressed or promoted. How do we calulate this? We project the hs using the output_projection, look at the diff from the last layer, and then project it back using the inverse of the output projection. This gives us the amount of suppression that occured at that layer.\n",
+    "    \"\"\"\n",
+    "    hs_flat = rearrange(hs[:, :, :], \"l b t h -> (l b t) h\")\n",
+    "    logits_flat = F.linear(hs_flat, w_out)\n",
+    "    logits = rearrange(\n",
+    "        logits_flat, \"(l b t) h -> l b t h\", l=hs.shape[0], b=hs.shape[1], t=hs.shape[2]\n",
+    "    )\n",
+    "    logit_diffs = logits[:, :, :].diff(dim=0)\n",
+    "    logit_diffs_flat = rearrange(logit_diffs, \"l b t h -> (l b t) h\")\n",
+    "    # W_inv = get_cache_inv(w_out)\n",
+    "\n",
+    "    # get the supression projected back\n",
+    "    hs_supr_flat = F.linear(logit_diffs_flat.to(dtype=w_inv.dtype), w_inv)\n",
+    "    hs_inv_flat = F.linear(logits_flat.to(dtype=w_inv.dtype), w_inv)\n",
+    "    supr_amounts = rearrange(\n",
+    "        hs_supr_flat, \"(l b t) h -> l b t h\", l=hs.shape[0] - 1, b=hs.shape[1], t=hs.shape[2]\n",
+    "    ).to(w_out.dtype)\n",
+    "    hs_inv = rearrange(\n",
+    "        hs_inv_flat, \"(l b t) h -> l b t h\", l=hs.shape[0], b=hs.shape[1], t=hs.shape[2]\n",
+    "    ).to(w_out.dtype)\n",
+    "\n",
+    "    # add on missing first layer\n",
+    "    # torch.zeros_like(supr_amounts[:1]).to(hs.device)\n",
+    "    supr_amounts = torch.cat(\n",
+    "        [torch.zeros_like(supr_amounts[:1]).to(hs.device), supr_amounts], dim=0\n",
+    "    )\n",
+    "\n",
+    "    residual = hs - hs_inv # lost in projection\n",
+    "    return supr_amounts, residual"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# def get_uniq_token_ids(tokens):\n",
+    "#     token_ids = tokenizer(\n",
+    "#         tokens, add_special_tokens=False, padding=False\n",
+    "#     ).input_ids\n",
+    "#     token_ids = torch.tensor(list(set([x[0] for x in token_ids]))).long()\n",
+    "#     print(\"before\", tokens)\n",
+    "#     print(\"after\", tokenizer.batch_decode(token_ids))\n",
+    "#     return token_ids\n",
+    "\n",
+    "\n",
+    "# false_tokens = [\"0\", \"0 \", \"0\\n\", \"false\", \"False \"]\n",
+    "# false_token_ids = get_uniq_token_ids(false_tokens)\n",
+    "\n",
+    "# true_tokens = [\"1\", \"1 \", \"1\\n\", \"true\", \"True \"]\n",
+    "# true_token_ids = get_uniq_token_ids(true_tokens)\n",
+    "\n",
+    "# print('QC: manually check that these are equivilent (no <end_of_text> or newline)')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# now we map to 1) calc supressed activations 2) llm answer (prob of 0 vs prob of 1)\n",
+    "\n",
+    "Wo = model.get_output_embeddings().weight.detach().clone().cpu()\n",
+    "Wo_inv = torch.pinverse(Wo.clone().float())\n",
+    "\n",
+    "\n",
+    "def postprocess_activation_ds_rows(o):\n",
+    "    \"\"\"Process model outputs\"\"\"\n",
+    "\n",
+    "    # Gather the log probabilities for the actual labels, which is the next token\n",
+    "    labels = o[\"labels\"]#[:, 1:].clone()\n",
+    "    labels = o[\"input_ids\"][1:].clone()\n",
+    "    logits = o[\"logits\"][:-1]\n",
+    "    # print(\"logits\", logits.shape, logits.dtype)\n",
+    "    log_probs = logits.log_softmax(dim=-1)\n",
+    "    o['label_logp'] = torch.gather(\n",
+    "        input=log_probs, dim=-1, index=labels.unsqueeze(-1)\n",
+    "    ).squeeze(-1)\n",
+    "\n",
+    "    # get supressed activations\n",
+    "    hs = o[\"hidden_states\"][None]\n",
+    "    hs = rearrange(hs, \"b l t h -> l b t h\")\n",
+    "    supr_amounts, residual = get_supressed_activations(hs, Wo.to(hs.dtype), Wo_inv.to(hs.dtype))\n",
+    "\n",
+    "    # we will only take the last half of layers, and the last token\n",
+    "    layer_half = hs.shape[0] // 2\n",
+    "    \n",
+    "    hs = rearrange(hs, \"l b t h -> b l t h\").squeeze(0)[layer_half:-2]\n",
+    "    supr_amounts = rearrange(supr_amounts, \"l b t h -> b l t h\").squeeze(0)[layer_half:-2]\n",
+    "    residual = rearrange(residual, \"l b t h -> b l t h\").squeeze(0)[layer_half:-2]\n",
+    "\n",
+    "    for k in o.keys():\n",
+    "        if k.startswith(\"acts-\"):\n",
+    "            o[k] = o[k][-1:]\n",
+    "\n",
+    "    o[\"hidden_states\"] = hs.half()[-1:]\n",
+    "    o[\"supr_amounts\"] = supr_amounts.half()\n",
+    "    o['logits'] = o['logits'].half()\n",
+    "    o['label_logp'] = o['label_logp'].half()\n",
+    "    o['residual'] = residual.half()\n",
+    "    return o\n",
+    "\n",
+    "\n",
+    "ds_a2 = ds_a2a.map(postprocess_activation_ds_rows, writer_batch_size=1, num_proc=None)\n",
+    "ds_a2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = Wo = Wo_inv = tokenizer = None\n",
+    "clear_mem()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "\n",
+    "ds = ds_a2 # concatenate_datasets([ds_a2, ds2.remove_columns(['messages', 'attention_mask', 'prompt', 'chosen', 'rejected', 'chosen_messages'])]).with_format(\"torch\")\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "{k: v.shape for k,v in ds[:2].items() if isinstance(v, torch.Tensor)}\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "out_dir = Path('../data/activation_store')\n",
+    "name = acts_outfile.with_suffix(\"\").relative_to('/tmp/activation_store')\n",
+    "acts_outfile2 = out_dir / name\n",
+    "acts_outfile2.parent.mkdir(parents=True, exist_ok=True)\n",
+    "acts_outfile2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import concatenate_datasets, load_dataset\n",
+    "\n",
+    "\n",
+    "ds_out = ds.with_format(\"torch\")\n",
+    "ds_out.save_to_disk(acts_outfile2)\n",
+    "acts_outfile2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_config = acts_outfile2.with_suffix(\".json\")\n",
+    "json.dump({\n",
+    "    \"model_name\": model_name,\n",
+    "    \"batch_size\": batch_size,\n",
+    "    \"max_length\": max_length,\n",
+    "    'layer_groups': layer_groups,\n",
+    "    # 'model_config': model.config.to_dict(),\n",
+    "    \"n_rows\": len(ds_out),\n",
+    "}, open(f_config, \"w\"))\n",
+    "f_config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "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.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/nbs/02b_TQA_regr_w_kv.ipynb b/nbs/02b_TQA_regr_w_kv.ipynb
index d405467..999c93d 100644
--- a/nbs/02b_TQA_regr_w_kv.ipynb
+++ b/nbs/02b_TQA_regr_w_kv.ipynb
@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -23,7 +23,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -72,7 +72,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -99,32 +99,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'model_name': 'Qwen/Qwen3-1.7B', 'batch_size': 10, 'max_length': 90, 'split': 'train', 'n_rows': 632}\n"
+      "{'model_name': 'Qwen/Qwen3-1.7B', 'batch_size': 10, 'max_length': 90, 'layer_groups': {'mlp.down_proj': ['model.layers.14.mlp.down_proj', 'model.layers.15.mlp.down_proj', 'model.layers.16.mlp.down_proj', 'model.layers.17.mlp.down_proj', 'model.layers.18.mlp.down_proj', 'model.layers.19.mlp.down_proj', 'model.layers.20.mlp.down_proj', 'model.layers.21.mlp.down_proj', 'model.layers.22.mlp.down_proj', 'model.layers.23.mlp.down_proj', 'model.layers.24.mlp.down_proj', 'model.layers.25.mlp.down_proj'], 'self_attn': ['model.layers.14.self_attn', 'model.layers.15.self_attn', 'model.layers.16.self_attn', 'model.layers.17.self_attn', 'model.layers.18.self_attn', 'model.layers.19.self_attn', 'model.layers.20.self_attn', 'model.layers.21.self_attn', 'model.layers.22.self_attn', 'model.layers.23.self_attn', 'model.layers.24.self_attn', 'model.layers.25.self_attn'], 'mlp.up_proj': ['model.layers.14.mlp.up_proj', 'model.layers.15.mlp.up_proj', 'model.layers.16.mlp.up_proj', 'model.layers.17.mlp.up_proj', 'model.layers.18.mlp.up_proj', 'model.layers.19.mlp.up_proj', 'model.layers.20.mlp.up_proj', 'model.layers.21.mlp.up_proj', 'model.layers.22.mlp.up_proj', 'model.layers.23.mlp.up_proj', 'model.layers.24.mlp.up_proj', 'model.layers.25.mlp.up_proj']}, 'n_rows': 316}\n"
      ]
     },
     {
      "data": {
       "text/plain": [
        "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts', 'question', 'input_ids'],\n",
-       "    num_rows: 632\n",
+       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts', 'input_ids'],\n",
+       "    num_rows: 316\n",
        "})"
       ]
      },
-     "execution_count": 50,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "acts_outfile = Path('../data/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2')\n",
+    "acts_outfile = Path('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.1')\n",
     "\n",
     "f_config = acts_outfile.with_suffix(\".json\")\n",
     "config = json.load(open(f_config, 'r'))\n",
@@ -137,7 +137,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -146,7 +146,7 @@
        "['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj']"
       ]
      },
-     "execution_count": 19,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -158,7 +158,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -176,8 +176,7 @@
       "llm_ans torch.Size([2])\n",
       "llm_log_prob_true torch.Size([])\n",
       "supr_amounts torch.Size([13, 1, 2048])\n",
-      "question <class 'NoneType'>\n",
-      "input_ids <class 'NoneType'>\n"
+      "input_ids torch.Size([90])\n"
      ]
     }
    ],
@@ -191,19 +190,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts', 'question', 'input_ids'],\n",
-       "    num_rows: 632\n",
+       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts', 'input_ids'],\n",
+       "    num_rows: 316\n",
        "})"
       ]
      },
-     "execution_count": 21,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -212,6 +211,13 @@
     "ds_a2"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -221,16 +227,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "505"
+       "252"
       ]
      },
-     "execution_count": 25,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -251,7 +257,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -268,7 +274,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -323,7 +329,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -334,7 +340,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -418,7 +424,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -429,7 +435,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -454,7 +460,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -483,7 +489,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -503,7 +509,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 34,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -515,7 +521,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 35,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -558,7 +564,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -569,7 +575,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 37,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -626,7 +632,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 38,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -644,7 +650,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 39,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [
     {
@@ -657,7 +663,7 @@
        " 'acts-mlp.up_proj']"
       ]
      },
-     "execution_count": 39,
+     "execution_count": 23,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -669,7 +675,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 40,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
     {
@@ -686,7 +692,7 @@
        " ('supressed_hs(0.5)', 'magnitude(0.95)', 'max')]"
       ]
      },
-     "execution_count": 40,
+     "execution_count": 24,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -747,7 +753,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 41,
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ds_a2['hidden_states'][-1]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -766,7 +781,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 42,
+   "execution_count": 27,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -776,22 +791,51 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 43,
+   "execution_count": 28,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# ds"
+    "import os, unicodedata, string\n",
+    "from pathlib import Path\n",
+    "\n",
+    "def sanitize_path(path: Path | str, allow_period: bool = True) -> Path:\n",
+    "    \"\"\"\n",
+    "    Whitelist only ASCII letters, digits, dash, underscore,\n",
+    "    optionally period, and forward‐slash. Replace others with '_'.\n",
+    "    \"\"\"\n",
+    "    s = unicodedata.normalize(\"NFKD\", str(path))\\\n",
+    "                     .encode(\"ascii\", \"ignore\")\\\n",
+    "                     .decode()\n",
+    "    s = s.replace(os.sep, \"/\")\n",
+    "    allowed = set(string.ascii_letters + string.digits + \"_-\")\n",
+    "    if allow_period: allowed.add(\".\")\n",
+    "    allowed.add(\"/\")\n",
+    "    return Path(\"\".join(ch if ch in allowed else \"_\" for ch in s))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4987c6c94d864d90aa1d0df004a9c520",
+       "model_id": "b0233174b8884585abaf13e264433ef6",
        "version_major": 2,
        "version_minor": 0
       },
@@ -801,268 +845,6 @@
      },
      "metadata": {},
      "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[32m2025-05-06 18:18:11.865\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:20.360\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_max): 0.764 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:20.865\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:27.016\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_first): 0.843 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:27.183\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:34.395\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_mean): 0.850 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:34.723\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:41.460\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_min:): 0.432 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:41.636\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:48.003\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_min:): 0.548 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:48.230\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:54.370\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_first): 0.826 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:18:54.625\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:01.555\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_min:): 0.349 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:01.950\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:05.493\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_max): 0.417 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:05.753\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:12.217\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_flatten): 0.403 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:12.382\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:18.314\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_last): 0.547 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:18.748\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:22.235\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_mean): 0.579 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:22.407\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:25.778\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_max): 0.418 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:26.089\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:32.343\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_flatten): 0.556 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:32.561\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:39.102\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_flatten): 0.325 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:39.272\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:45.866\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_flatten): 0.561 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:46.051\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:52.794\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_flatten): 0.327 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:53.038\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:59.322\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_mean): 0.467 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:19:59.518\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:02.740\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_min:): 0.577 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:02.953\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:09.203\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_max): 0.237 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:09.379\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:15.182\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_first): 0.845 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:15.337\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:21.383\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_min:): 0.512 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:21.571\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:25.024\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_flatten): 0.668 roc auc, n=64. X.shape=torch.Size([316, 2211840])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:25.195\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:31.455\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_sum): 0.861 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:31.633\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:34.735\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_mean): 0.561 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:34.897\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:40.634\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_last): 0.568 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:40.810\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:47.370\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_max): 0.253 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:47.541\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:53.463\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_max): 0.775 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:20:53.725\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:00.366\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_std): 0.293 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:00.529\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:06.918\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_flatten): 0.326 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:07.120\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:10.069\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_last): 0.507 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:10.229\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:16.489\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_max): 0.242 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:16.651\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:23.264\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_std): 0.337 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:23.518\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:30.132\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_std): 0.667 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:30.307\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:33.460\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_min:): 0.411 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:33.620\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:39.365\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_last): 0.638 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:39.530\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:45.588\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_first): 0.850 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:45.763\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:52.187\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_min:): 0.787 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:52.462\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:58.315\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_first): 0.846 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:21:58.509\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:02.500\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_std): 0.588 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:02.725\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:08.723\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_last): 0.562 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:08.906\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:12.676\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_flatten): 0.682 roc auc, n=64. X.shape=torch.Size([316, 2211840])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:12.865\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:18.961\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_min:): 0.516 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:19.127\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:25.034\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_mean): 0.798 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:25.189\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:31.305\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_max): 0.764 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:31.481\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:34.985\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_std): 0.599 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:35.168\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:41.402\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_flatten): 0.435 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:41.589\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:47.755\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_sum): 0.453 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:47.968\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:54.455\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_flatten): 0.570 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:22:54.627\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:01.387\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_std): 0.766 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:01.543\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:04.759\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_min:): 0.413 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:04.934\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:10.664\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_last): 0.359 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:10.848\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:14.030\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_last): 0.368 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:14.189\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:20.498\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_min:): 0.422 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:20.658\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:23.803\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_min:): 0.425 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:23.955\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:29.975\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_mean): 0.831 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:30.142\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:35.596\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_last): 0.636 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:35.750\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:42.101\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_sum): 0.628 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:42.260\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:48.420\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_sum): 0.401 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:48.582\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:54.272\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_last): 0.577 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:23:54.437\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:00.934\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_max): 0.732 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:01.162\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:06.731\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_first): 0.834 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:06.908\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:13.056\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_max): 0.200 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:13.216\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:16.629\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_std): 0.595 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:16.843\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:19.984\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_max): 0.573 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:20.141\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:25.772\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_first): 0.841 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:25.929\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:29.338\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_sum): 0.576 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:29.501\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:35.656\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_sum): 0.763 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:35.817\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:38.839\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_sum): 0.628 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:39.005\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:45.465\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_std): 0.542 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:45.625\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:48.667\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_mean): 0.601 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:48.854\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:51.627\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_first): 0.408 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:51.815\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:55.586\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_flatten): 0.688 roc auc, n=64. X.shape=torch.Size([316, 2211840])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:24:55.766\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:02.488\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_std): 0.542 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:02.643\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:08.606\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_max): 0.735 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:08.812\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:11.969\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_mean): 0.782 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:12.129\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:18.535\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_max): 0.758 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:18.692\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:25.380\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_std): 0.465 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:25.567\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:31.448\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_first): 0.830 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:31.649\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:37.970\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_sum): 0.807 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:38.132\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:41.174\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_sum): 0.527 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:41.334\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:47.501\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_mean): 0.848 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:47.660\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:53.619\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_mean): 0.858 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:25:53.802\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:00.345\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_flatten): 0.385 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:00.505\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:06.743\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_min:): 0.437 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:06.903\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:12.734\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_last): 0.691 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:12.892\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:15.858\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_first): 0.541 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:16.033\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:19.747\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_std): 0.606 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:19.909\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:26.364\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_flatten): 0.330 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:26.574\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:33.018\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_std): 0.755 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:33.175\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:39.126\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_min:): 0.288 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:39.310\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:42.157\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_last): 0.376 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:42.309\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:48.701\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_min:): 0.227 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:48.898\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:52.221\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_min:): 0.406 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:52.373\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:58.193\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_sum): 0.761 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:26:58.355\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-5), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:04.507\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-5)_mean_sum): 0.727 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:04.662\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:07.496\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_first): 0.404 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:07.662\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:13.416\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_last): 0.493 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:13.607\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:16.843\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_sum): 0.660 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:17.015\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:22.836\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_mean): 0.475 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:22.989\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:29.282\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_max): 0.713 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:29.455\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:32.611\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_sum): 0.573 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:32.773\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:39.102\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_mean): 0.786 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:39.254\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:41.984\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_first): 0.432 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:42.135\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-self_attn, mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:45.013\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-self_attn_mean_first): 0.425 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:45.198\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.1), mean, min:\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:51.528\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.1)_mean_min:): 0.188 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:51.720\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:57.717\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_mean): 0.497 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:27:57.902\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:01.044\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_max): 0.414 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:01.208\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:04.404\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_mean): 0.575 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:04.587\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.1), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:10.188\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.1)_mean_first): 0.841 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:10.436\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:14.155\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_flatten): 0.706 roc auc, n=64. X.shape=torch.Size([316, 2211840])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:14.316\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.down_proj, mean, max\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:17.627\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.down_proj_mean_max): 0.421 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:17.810\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:20.563\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_last): 0.416 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:20.732\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:27.383\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_std): 0.358 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:27.559\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(5), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:33.670\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(5)_mean_sum): 0.516 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:33.824\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:39.716\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_sum): 0.763 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:39.871\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supr_amounts, mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:43.342\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supr_amounts_mean_std): 0.595 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:43.519\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:49.018\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_first): 0.833 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:49.176\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:54.908\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0)_mean_last): 0.198 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:28:55.096\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:01.365\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_last): 0.525 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:01.554\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:08.573\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_std): 0.363 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:08.753\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.01), mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:15.215\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.01)_mean_flatten): 0.649 roc auc, n=64. X.shape=torch.Size([316, 2396160])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:15.403\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(0.5), mean, std\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:22.241\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(0.5)_mean_std): 0.545 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:22.416\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing hidden_states, mean, flatten\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:25.974\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(hidden_states_mean_flatten): 0.672 roc auc, n=64. X.shape=torch.Size([316, 2211840])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:26.152\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-1), mean, first\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:32.057\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-1)_mean_first): 0.845 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:32.232\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(1), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:38.415\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(1)_mean_mean): 0.705 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:38.573\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.01), mean, mean\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:44.718\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.01)_mean_mean): 0.856 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:44.896\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing supressed_hs(-0.5), mean, sum\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:51.128\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(supressed_hs(-0.5)_mean_sum): 0.793 roc auc, n=64. X.shape=torch.Size([316, 26624])\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:51.297\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m15\u001b[0m - \u001b[1mProcessing acts-mlp.up_proj, mean, last\u001b[0m\n",
-      "\u001b[32m2025-05-06 18:29:54.184\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mtrain_linear_prob_on_dataset\u001b[0m:\u001b[36m38\u001b[0m - \u001b[1mscore for probe(acts-mlp.up_proj_mean_last): 0.482 roc auc, n=64. X.shape=torch.Size([316, 24576])\u001b[0m\n"
-     ]
     }
    ],
    "source": [
@@ -1117,6 +899,15 @@
     "    clear_mem()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ds_a2['hidden_states'][-1]"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -1126,7 +917,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 31,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1139,7 +930,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 32,
    "metadata": {},
    "outputs": [
     {
@@ -1148,7 +939,7 @@
        "0.6392156862745099"
       ]
      },
-     "execution_count": 48,
+     "execution_count": 32,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1171,7 +962,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 33,
    "metadata": {},
    "outputs": [
     {
@@ -1180,7 +971,7 @@
        "0.8431372549019609"
       ]
      },
-     "execution_count": 49,
+     "execution_count": 33,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1212,7 +1003,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 34,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1221,7 +1012,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 35,
    "metadata": {},
    "outputs": [
     {
@@ -1296,7 +1087,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 36,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1323,7 +1114,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 37,
    "metadata": {},
    "outputs": [
     {
@@ -1385,14 +1176,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 38,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/tmp/ipykernel_1036655/1808615273.py:11: UserWarning: No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.\n",
+      "/tmp/ipykernel_3634226/1808615273.py:11: UserWarning: No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.\n",
       "  plt.legend().remove()\n"
      ]
     },
@@ -1402,7 +1193,7 @@
        "PosixPath('../figs/truthfulqa_Qwen_Qwen3-1.7B.png')"
       ]
      },
-     "execution_count": 54,
+     "execution_count": 38,
      "metadata": {},
      "output_type": "execute_result"
     },
@@ -1439,7 +1230,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 39,
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/nbs/03_stats.ipynb b/nbs/03_stats.ipynb
index 595a949..862d081 100644
--- a/nbs/03_stats.ipynb
+++ b/nbs/03_stats.ipynb
@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -23,7 +23,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -72,7 +72,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -99,37 +99,50 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'model_name': 'Qwen/Qwen3-1.7B', 'batch_size': 10, 'max_length': 90, 'split': 'train', 'n_rows': 632}\n"
+      "ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-100-128_v4e_wgen\n",
+      "ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2\n",
+      "ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-40-128_v4e_wgen\n",
+      "ds_at-QwenQwen3-1.json\n"
      ]
-    },
+    }
+   ],
+   "source": [
+    "!ls ../data/activation_store/"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
     {
      "data": {
       "text/plain": [
        "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts', 'question', 'input_ids'],\n",
-       "    num_rows: 632\n",
+       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'supr_amounts', 'i', 'input_ids', 'prompt_mask', 'completion_type'],\n",
+       "    num_rows: 200\n",
        "})"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "acts_outfile = Path('../data/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2')\n",
+    "acts_outfile = Path('../data/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-100-128_v4f_wgen')\n",
     "\n",
-    "f_config = acts_outfile.with_suffix(\".json\")\n",
-    "config = json.load(open(f_config, 'r'))\n",
-    "model_name = config['model_name']\n",
-    "print(config)\n",
+    "# f_config = acts_outfile.with_suffix(\".json\")\n",
+    "# config = json.load(open(f_config, 'r'))\n",
+    "# model_name = config['model_name']\n",
+    "# print(config)\n",
     "\n",
     "ds_a2 = load_from_disk(acts_outfile).with_format(\"torch\")\n",
     "ds_a2"
@@ -137,7 +150,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
@@ -146,7 +159,7 @@
        "['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj']"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -158,31 +171,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "acts-mlp.down_proj torch.Size([1, 90, 2048])\n",
-      "acts-self_attn torch.Size([1, 90, 2048])\n",
-      "acts-mlp.up_proj torch.Size([1, 90, 6144])\n",
-      "loss torch.Size([])\n",
-      "logits torch.Size([151936])\n",
-      "hidden_states torch.Size([1, 90, 2048])\n",
-      "attention_mask torch.Size([90])\n",
-      "label torch.Size([])\n",
-      "llm_ans torch.Size([2])\n",
-      "llm_log_prob_true torch.Size([])\n",
-      "supr_amounts torch.Size([13, 1, 2048])\n",
-      "question <class 'NoneType'>\n",
-      "input_ids <class 'NoneType'>\n"
+      "acts-mlp.down_proj torch.Size([2, 1, 128, 2048])\n",
+      "acts-self_attn torch.Size([2, 1, 128, 2048])\n",
+      "acts-mlp.up_proj torch.Size([2, 1, 128, 6144])\n",
+      "loss torch.Size([2])\n",
+      "logits torch.Size([2, 151936])\n",
+      "hidden_states torch.Size([2, 1, 128, 2048])\n",
+      "attention_mask torch.Size([2, 128])\n",
+      "supr_amounts torch.Size([2, 3, 1, 2048])\n",
+      "i torch.Size([2])\n",
+      "input_ids <class 'list'>\n",
+      "prompt_mask <class 'list'>\n",
+      "completion_type <class 'list'>\n"
      ]
     }
    ],
    "source": [
-    "for k,v in ds_a2[0].items():\n",
+    "for k,v in ds_a2[:2].items():\n",
     "    if hasattr(v, 'shape'):\n",
     "        print(k, v.shape)\n",
     "    else:\n",
@@ -191,19 +203,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts', 'question', 'input_ids'],\n",
-       "    num_rows: 632\n",
+       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'supr_amounts', 'i', 'input_ids', 'prompt_mask', 'completion_type'],\n",
+       "    num_rows: 200\n",
        "})"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -223,8 +235,47 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "ds_a2"
+    "Lets start with std of each hs\n",
+    "- for each i, compare completion types\n",
+    "   - for each ks\n",
+    "     - where "
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['acts-mlp.down_proj',\n",
+       " 'acts-self_attn',\n",
+       " 'acts-mlp.up_proj',\n",
+       " 'model_output',\n",
+       " 'loss',\n",
+       " 'logits',\n",
+       " 'hidden_states',\n",
+       " 'attention_mask',\n",
+       " 'supr_amounts']"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ks = act_groups + [ 'hidden_states', 'attention_mask', 'supr_amounts']\n",
+    "ks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {