From c066477b5e37b9bf614c3559b4f1c98dc06c3178 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Fri, 14 Mar 2025 17:09:30 +0800 Subject: [PATCH] wip --- nbs/TQA_regr_w_kv 3B.ipynb | 1067 ++++++++++++++++++++++++++++++++++++ 1 file changed, 1067 insertions(+) create mode 100644 nbs/TQA_regr_w_kv 3B.ipynb diff --git a/nbs/TQA_regr_w_kv 3B.ipynb b/nbs/TQA_regr_w_kv 3B.ipynb new file mode 100644 index 0000000..7e1184d --- /dev/null +++ b/nbs/TQA_regr_w_kv 3B.ipynb @@ -0,0 +1,1067 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Quick experiment to see which is better at detecting truthful answers\n", + "\n", + "- model outputs\n", + "- hs\n", + "- supressed activations (Hypothesis this is better)" + ] + }, + { + "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": 3, + "metadata": {}, + "outputs": [], + "source": [ + "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", + "\n", + "from jaxtyping import Float\n", + "from torch import Tensor\n", + "\n", + "from activation_store.collect import activation_store, default_postprocess_result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load model" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Sliding Window Attention is enabled but not implemented for `eager`; unexpected results may be encountered.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "3e298fba714a478ba42279bb7c81e379", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading checkpoint shards: 0%| | 0/2 [00:00\n" + ] + } + ], + "source": [ + "collate_fn = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n", + "ds = DataLoader(ds2, batch_size=1, collate_fn=collate_fn)\n", + "print(ds)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Collect activations" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'mlp.down_proj': ['model.layers.0.mlp.down_proj',\n", + " 'model.layers.1.mlp.down_proj',\n", + " 'model.layers.2.mlp.down_proj',\n", + " 'model.layers.3.mlp.down_proj',\n", + " 'model.layers.4.mlp.down_proj',\n", + " 'model.layers.5.mlp.down_proj',\n", + " 'model.layers.6.mlp.down_proj',\n", + " 'model.layers.7.mlp.down_proj',\n", + " 'model.layers.8.mlp.down_proj',\n", + " 'model.layers.9.mlp.down_proj',\n", + " 'model.layers.10.mlp.down_proj',\n", + " 'model.layers.11.mlp.down_proj',\n", + " 'model.layers.12.mlp.down_proj',\n", + " 'model.layers.13.mlp.down_proj',\n", + " '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.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", + " 'model.layers.26.mlp.down_proj',\n", + " 'model.layers.27.mlp.down_proj',\n", + " 'model.layers.28.mlp.down_proj',\n", + " 'model.layers.29.mlp.down_proj',\n", + " 'model.layers.30.mlp.down_proj',\n", + " 'model.layers.31.mlp.down_proj',\n", + " 'model.layers.32.mlp.down_proj',\n", + " 'model.layers.33.mlp.down_proj',\n", + " 'model.layers.34.mlp.down_proj',\n", + " 'model.layers.35.mlp.down_proj'],\n", + " 'self_attn': ['model.layers.0.self_attn',\n", + " 'model.layers.1.self_attn',\n", + " 'model.layers.2.self_attn',\n", + " 'model.layers.3.self_attn',\n", + " 'model.layers.4.self_attn',\n", + " 'model.layers.5.self_attn',\n", + " 'model.layers.6.self_attn',\n", + " 'model.layers.7.self_attn',\n", + " 'model.layers.8.self_attn',\n", + " 'model.layers.9.self_attn',\n", + " 'model.layers.10.self_attn',\n", + " 'model.layers.11.self_attn',\n", + " 'model.layers.12.self_attn',\n", + " 'model.layers.13.self_attn',\n", + " '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.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", + " 'model.layers.26.self_attn',\n", + " 'model.layers.27.self_attn',\n", + " 'model.layers.28.self_attn',\n", + " 'model.layers.29.self_attn',\n", + " 'model.layers.30.self_attn',\n", + " 'model.layers.31.self_attn',\n", + " 'model.layers.32.self_attn',\n", + " 'model.layers.33.self_attn',\n", + " 'model.layers.34.self_attn',\n", + " 'model.layers.35.self_attn'],\n", + " 'mlp.up_proj': ['model.layers.0.mlp.up_proj',\n", + " 'model.layers.1.mlp.up_proj',\n", + " 'model.layers.2.mlp.up_proj',\n", + " 'model.layers.3.mlp.up_proj',\n", + " 'model.layers.4.mlp.up_proj',\n", + " 'model.layers.5.mlp.up_proj',\n", + " 'model.layers.6.mlp.up_proj',\n", + " 'model.layers.7.mlp.up_proj',\n", + " 'model.layers.8.mlp.up_proj',\n", + " 'model.layers.9.mlp.up_proj',\n", + " 'model.layers.10.mlp.up_proj',\n", + " 'model.layers.11.mlp.up_proj',\n", + " 'model.layers.12.mlp.up_proj',\n", + " 'model.layers.13.mlp.up_proj',\n", + " '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.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',\n", + " 'model.layers.26.mlp.up_proj',\n", + " 'model.layers.27.mlp.up_proj',\n", + " 'model.layers.28.mlp.up_proj',\n", + " 'model.layers.29.mlp.up_proj',\n", + " 'model.layers.30.mlp.up_proj',\n", + " 'model.layers.31.mlp.up_proj',\n", + " 'model.layers.32.mlp.up_proj',\n", + " 'model.layers.33.mlp.up_proj',\n", + " 'model.layers.34.mlp.up_proj',\n", + " 'model.layers.35.mlp.up_proj']}" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# choose layers to cache\n", + "layer_groups = {\n", + " 'mlp.down_proj': [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')],\n", + " 'self_attn': [k for k,v in model.named_modules() if k.endswith('.self_attn')],\n", + " 'mlp.up_proj': [k for k,v in model.named_modules() if k.endswith('mlp.up_proj')],\n", + "}\n", + "layer_groups" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "f = activation_store(ds, model, layers=layer_groups)\n", + "f" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ds_a = Dataset.from_parquet(str(f), split=split).with_format(\"torch\")\n", + "ds_a" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "max_length" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "{k:v.shape for k,v in ds_a[0].items()}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# sanity test generate\n", + "b = next(iter(ds))\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=3,\n", + ")\n", + "gent = tokenizer.batch_decode(o, skip_special_tokens=False)\n", + "for g in gent:\n", + " print(g)\n", + " print(\"---\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Get supressed activations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@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[:, :, -1:], \"l b t h -> (l b t) h\")\n", + " hs_out_flat = torch.nn.functional.linear(hs_flat, w_out)\n", + " hs_out = rearrange(\n", + " hs_out_flat, \"(l b t) h -> l b t h\", l=hs.shape[0], b=hs.shape[1], t=1\n", + " )\n", + " diffs = hs_out[:, :, :].diff(dim=0)\n", + " diffs_flat = rearrange(diffs, \"l b t h -> (l b t) h\")\n", + " # W_inv = get_cache_inv(w_out)\n", + "\n", + " diffs_inv_flat = torch.nn.functional.linear(diffs_flat.to(dtype=w_inv.dtype), w_inv)\n", + " diffs_inv = rearrange(\n", + " diffs_inv_flat, \"(l b t) h -> l b t h\", l=hs.shape[0] - 1, b=hs.shape[1], t=1\n", + " ).to(w_out.dtype)\n", + "\n", + " # add on missing first layer\n", + " torch.zeros_like(diffs_inv[:1]).to(hs.device)\n", + " diffs_inv = torch.cat(\n", + " [torch.zeros_like(diffs_inv[:1]).to(hs.device), diffs_inv], dim=0\n", + " )\n", + " return diffs_inv" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def get_uniq_token_ids(tokens):\n", + " token_ids = tokenizer(\n", + " tokens, return_tensors=\"pt\", add_special_tokens=False, padding=True\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)" + ] + }, + { + "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 proc(o):\n", + " # TODO batch it\n", + " \"\"\"Process model outputs\"\"\"\n", + "\n", + " # get llm ans\n", + " log_probs = o[\"logits\"][-1].log_softmax(0)\n", + " false_log_prob = log_probs.index_select(0, false_token_ids).sum()\n", + " true_log_prob = log_probs.index_select(0, true_token_ids).sum()\n", + " o[\"llm_ans\"] = torch.stack([false_log_prob, true_log_prob])\n", + " o[\"llm_log_prob_true\"] = true_log_prob - false_log_prob\n", + "\n", + " # get supressed activations\n", + " hs = o[\"hidden_states\"][None]\n", + " hs = rearrange(hs, \"b l t h -> l b t h\")\n", + " diffs_inv = 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", + " diffs_inv = rearrange(diffs_inv, \"l b t h -> b l t h\").squeeze(0)[layer_half:-2]\n", + "\n", + " o[\"hidden_states\"] = hs.half()\n", + " o[\"diffs_inv\"] = diffs_inv.half()\n", + " return o\n", + "\n", + "\n", + "ds_a2 = ds_a.map(proc, writer_batch_size=1, num_proc=None)\n", + "ds_a2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "{k: v.shape for k,v in ds_a2[0].items() if isinstance(v, torch.Tensor)}\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# # # now convert diffs_inv to supressed_mask and hs_sup\n", + "\n", + "# def proc2(o, eps = 1.0e-2):\n", + "# diffs_inv = o[\"diffs_inv\"]\n", + "# hs = o[\"hidden_states\"] # [b l h]\n", + "# supressed_mask = (diffs_inv < -eps).to(hs.dtype)# [b l h]\n", + "\n", + "# o['hs_sup'] = hs * supressed_mask\n", + "# o['supressed_mask'] = supressed_mask\n", + "# return o\n", + "\n", + "# ds_a2 = ds_a2.map(proc2, writer_batch_size=64, num_proc=None, batched=True, batch_size=64)\n", + "# ds_a2" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Predict" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# https://github.com/EleutherAI/ccs/blob/8a4bf687712cc03ef72973c8235944566d59053b/ccs/training/supervised.py#L9\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "class Classifier(torch.nn.Module):\n", + " \"\"\"Linear classifier trained with supervised learning.\"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " input_dim: int,\n", + " num_classes: int = 2,\n", + " device: str | torch.device | None = None,\n", + " dtype: torch.dtype | None = None,\n", + " ):\n", + " super().__init__()\n", + "\n", + " self.linear = torch.nn.Linear(\n", + " input_dim, num_classes if num_classes > 2 else 1, device=device, dtype=dtype\n", + " )\n", + " self.linear.bias.data.zero_()\n", + " # self.linear.weight.data.zero_()\n", + "\n", + " def forward(self, x: Tensor) -> Tensor:\n", + " return self.linear(x).squeeze(-1)\n", + "\n", + " @torch.enable_grad()\n", + " def fit(\n", + " self,\n", + " x: Tensor,\n", + " y: Tensor,\n", + " *,\n", + " l2_penalty: float = 0.001,\n", + " max_iter: int = 10_000,\n", + " ) -> float:\n", + " \"\"\"Fits the model to the input data using L-BFGS with L2 regularization.\n", + "\n", + " Args:\n", + " x: Input tensor of shape (N, D), where N is the number of samples and D is\n", + " the input dimension.\n", + " y: Target tensor of shape (N,) for binary classification or (N, C) for\n", + " multiclass classification, where C is the number of classes.\n", + " l2_penalty: L2 regularization strength.\n", + " max_iter: Maximum number of iterations for the L-BFGS optimizer.\n", + "\n", + " Returns:\n", + " Final value of the loss function after optimization.\n", + " \"\"\"\n", + " optimizer = torch.optim.LBFGS(\n", + " self.parameters(),\n", + " line_search_fn=\"strong_wolfe\",\n", + " max_iter=max_iter,\n", + " )\n", + "\n", + " num_classes = self.linear.out_features\n", + " loss_fn = bce_with_logits if num_classes == 1 else cross_entropy\n", + " loss = torch.inf\n", + " y = y.to(\n", + " torch.get_default_dtype() if num_classes == 1 else torch.long,\n", + " )\n", + "\n", + " def closure():\n", + " nonlocal loss\n", + " optimizer.zero_grad()\n", + "\n", + " # Calculate the loss function\n", + " logits = self(x).squeeze(-1)\n", + " loss = loss_fn(logits, y)\n", + " if l2_penalty:\n", + " reg_loss = loss + l2_penalty * self.linear.weight.square().sum()\n", + " else:\n", + " reg_loss = loss\n", + "\n", + " reg_loss.backward()\n", + " return float(reg_loss)\n", + "\n", + " optimizer.step(closure)\n", + " return float(loss)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# first try llm\n", + "\n", + "\n", + "def roc_auc(y_true: Tensor, y_pred: Tensor) -> Tensor:\n", + " \"\"\"Area under the receiver operating characteristic curve (ROC AUC).\n", + "\n", + " Unlike scikit-learn's implementation, this function supports batched inputs of\n", + " shape `(N, n)` where `N` is the number of datasets and `n` is the number of samples\n", + " within each dataset. This is primarily useful for efficiently computing bootstrap\n", + " confidence intervals.\n", + "\n", + " Args:\n", + " y_true: Ground truth tensor of shape `(N,)` or `(N, n)`.\n", + " y_pred: Predicted class tensor of shape `(N,)` or `(N, n)`.\n", + "\n", + " Returns:\n", + " Tensor: If the inputs are 1D, a scalar containing the ROC AUC. If they're 2D,\n", + " a tensor of shape (N,) containing the ROC AUC for each dataset.\n", + " \"\"\"\n", + " if y_true.shape != y_pred.shape:\n", + " raise ValueError(\n", + " f\"y_true and y_pred should have the same shape; \"\n", + " f\"got {y_true.shape} and {y_pred.shape}\"\n", + " )\n", + " if y_true.dim() not in (1, 2):\n", + " raise ValueError(\"y_true and y_pred should be 1D or 2D tensors\")\n", + "\n", + " # Sort y_pred in descending order and get indices\n", + " indices = y_pred.argsort(descending=True, dim=-1)\n", + "\n", + " # Reorder y_true based on sorted y_pred indices\n", + " y_true_sorted = y_true.gather(-1, indices)\n", + "\n", + " # Calculate number of positive and negative samples\n", + " num_positives = y_true.sum(dim=-1)\n", + " num_negatives = y_true.shape[-1] - num_positives\n", + "\n", + " # Calculate cumulative sum of true positive counts (TPs)\n", + " tps = torch.cumsum(y_true_sorted, dim=-1)\n", + "\n", + " # Calculate cumulative sum of false positive counts (FPs)\n", + " fps = torch.cumsum(1 - y_true_sorted, dim=-1)\n", + "\n", + " # Calculate true positive rate (TPR) and false positive rate (FPR)\n", + " tpr = tps / num_positives.view(-1, 1)\n", + " fpr = fps / num_negatives.view(-1, 1)\n", + "\n", + " # Calculate differences between consecutive FPR values (widths of trapezoids)\n", + " fpr_diffs = torch.cat(\n", + " [fpr[..., 1:] - fpr[..., :-1], torch.zeros_like(fpr[..., :1])], dim=-1\n", + " )\n", + "\n", + " # Calculate area under the ROC curve for each dataset using trapezoidal rule\n", + " return torch.sum(tpr * fpr_diffs, dim=-1).squeeze()\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Score llm output" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "train_test_split = 200\n", + "a, b = ds_a2[\"llm_log_prob_true\"] > 0, ds_a2[\"label\"]\n", + "score = roc_auc(b[train_test_split:], a[train_test_split:])\n", + "print(f\"LLM score: {score:.2f} roc auc, n={len(a[train_test_split:])}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### score hidden states and activations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def train_linear_prob_on_dataset(\n", + " X,\n", + " name=\"\",\n", + " device: str = \"cuda\",\n", + "):\n", + " print(X.shape)\n", + " X = X.view(len(X), -1).to(device)\n", + "\n", + " # norm X\n", + " X = (X - X.mean()) / X.std()\n", + " y = ds_a2[\"label\"].to(device)\n", + " X_train, y_train = X[:train_test_split], y[:train_test_split]\n", + " X_test, y_test = X[train_test_split:], y[train_test_split:]\n", + " # data.shape\n", + " lr_model = Classifier(X.shape[-1], device=device)\n", + " lr_model.fit(X_train, y_train)\n", + "\n", + " y_pred = lr_model.forward(X_test)\n", + "\n", + " score = roc_auc(y_test, y_pred)\n", + " print(f\"score for probe({name}): {score:.3f} roc auc, n={len(X_test)}\")\n", + " return score.cpu().item()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def calc_supp_thresh(hs, diffs_inv, eps = 1.0e-2):\n", + " supressed_mask = (diffs_inv < -eps).to(hs.dtype)\n", + " hs_sup = hs * supressed_mask\n", + " return hs_sup, supressed_mask" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "reductions = {\n", + " \"mean\": lambda x: x.mean(0),\n", + " \"max\": lambda x: x.max(0)[0],\n", + " \"sum\": lambda x: x.sum(0),\n", + " \"last\": lambda x: x[-1],\n", + " \"first\": lambda x: x[0],\n", + " \"none\": lambda x: x,\n", + "}\n", + "results = []\n", + "\n", + "# first try hidden states\n", + "for r1 in reductions:\n", + " for dn in [ \"hidden_states\",'mlp.down_proj',\n", + " 'self_attn',\n", + " 'mlp.up_proj',]:\n", + " r1f = reductions[r1]\n", + " try:\n", + " X = torch.stack([r1f(x) for x in ds_a2[dn]])\n", + " name = f\"{dn} {r1}\"\n", + " score = train_linear_prob_on_dataset(X, name)\n", + " results.append((name, score))\n", + " except Exception as e:\n", + " print(f\"error with {name}\")\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### score supressed activations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def calc_hs_sup(o, eps = 1.0e-2):\n", + " diffs_inv = o[\"diffs_inv\"]\n", + " hs = o[\"hidden_states\"] # [b l h]\n", + " if eps > 0:\n", + " supressed_mask = (diffs_inv > eps).to(hs.dtype)# [b l h]\n", + " else:\n", + " supressed_mask = (diffs_inv < eps).to(hs.dtype)\n", + "\n", + " o['supressed_hs'] = hs * supressed_mask\n", + " o['supressed_mask'] = supressed_mask\n", + " # print({k:v.shape for k,v in o.items()})\n", + " return o" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "{k: v.shape for k,v in ds_a2[0].items() if isinstance(v, torch.Tensor)}\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import gc\n", + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "X = ds_a2['logits']\n", + "name = \"logits\"\n", + "score = train_linear_prob_on_dataset(X, name)\n", + "results.append((name, score))\n", + "score" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "X = ds_a2['llm_ans']\n", + "y = ds_a2['label']\n", + "\n", + "X_train, y_train = X[:train_test_split], y[:train_test_split]\n", + "X_test, y_test = X[train_test_split:], y[train_test_split:]\n", + "\n", + "score = roc_auc(y_test, X_test[:, 0]).item()\n", + "results.append(('llm_ans', score))\n", + "score" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "X = 1-torch.sigmoid(ds_a2['llm_log_prob_true']/10)\n", + "y = ds_a2['label']\n", + "\n", + "X_train, y_train = X[:train_test_split], y[:train_test_split]\n", + "X_test, y_test = X[train_test_split:], y[train_test_split:]\n", + "\n", + "score = roc_auc(y_test, X_test).item()\n", + "results.append(('llm_log_prob_true', score))\n", + "score" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# now various eps\n", + "for eps in [-50, -10, -5, -1, -0.5, -0.1, -0.01, -0, 0, 0.01, 0.1, 0.5, 1, 10, 50]:\n", + " gc.collect()\n", + " ds_a3 = ds_a2.map(lambda x:calc_hs_sup(x, eps=eps), num_proc=None, batched=True, batch_size=64, desc=f\"eps {eps}\")\n", + " print(f\"eps {eps} ds_a3['supressed_mask'].mean()={ds_a3['supressed_mask'].mean()}\")\n", + " data_names = [\"supressed_hs\", \"supressed_mask\"]\n", + " for dn in data_names:\n", + " try:\n", + " X = torch.stack([r1f(x) for x in ds_a3[dn]])\n", + " name = f\"{dn} {r1} {eps}\"\n", + " score = train_linear_prob_on_dataset(X, name)\n", + " results.append((name, score))\n", + " except Exception as e:\n", + " print(f\"error with {name}\")\n", + " print(e)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "# note hs_sup seems to get more important as we lower the thresh\n", + "df = pd.DataFrame(results, columns=[\"name\", \"auroc\"]).sort_values(\n", + " \"auroc\", ascending=False\n", + ")\n", + "df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "df['data'] = df['name'].apply(lambda x: x.split()[0])\n", + "df2 = df.groupby('data').max().sort_values(\"auroc\", ascending=False)\n", + "df2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "df2.index" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# plot it\n", + "# TODO add logits\n", + "\n", + "from matplotlib import pyplot as plt\n", + "df3 = df2.T[['llm_ans', 'llm_log_prob_true', 'hidden_states', 'supressed_hs', 'self_attn', 'mlp.up_proj',\n", + " 'mlp.down_proj', ]].rename(columns={\n", + " 'llm_ans': 'LLM Answer',\n", + " 'llm_log_prob_true': 'LLM Probability',\n", + " 'hidden_states': 'Hidden States',\n", + " 'acts': 'Activations: up_proj',\n", + " # 'logits': 'Logits',\n", + " 'supressed_hs': 'Supressed Hidden States',\n", + "}).T.sort_values(\"auroc\", ascending=False)\n", + "df3.plot.barh()\n", + "plt.legend().remove()\n", + "plt.xlabel(f\"Linar probe AUROC\")\n", + "plt.title(f\"TruthfulQA Binary with {model_name}\")\n", + "plt.xlim(0.5, None)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "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 +}