mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-09-21 12:50:41 +08:00
492 lines
20 KiB
Plaintext
492 lines
20 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Let's implement CCS from scratch.\n",
|
|
"This will deliberately be a simple (but less efficient) implementation to make everything as clear as possible."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"/accounts/projects/jsteinhardt/uid1693600/.local/lib/python3.7/site-packages/pandas/compat/_optional.py:138: UserWarning: Pandas requires version '2.7.0' or newer of 'numexpr' (version '2.6.9' currently installed).\n",
|
|
" warnings.warn(msg, UserWarning)\n",
|
|
"W1206 14:15:47.056510 140146004121408 builder.py:641] Reusing dataset amazon_polarity (/scratch/users/uid1693600/huggingface-cache/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc)\n"
|
|
]
|
|
},
|
|
{
|
|
"data": {
|
|
"application/vnd.jupyter.widget-view+json": {
|
|
"model_id": "0e5b9e737f6e415aa43cb2e7234b149f",
|
|
"version_major": 2,
|
|
"version_minor": 0
|
|
},
|
|
"text/plain": [
|
|
" 0%| | 0/2 [00:00<?, ?it/s]"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
|
|
"Some weights of the model checkpoint at microsoft/deberta-v2-xxlarge were not used when initializing DebertaV2ForMaskedLM: ['lm_predictions.lm_head.bias', 'lm_predictions.lm_head.LayerNorm.bias', 'lm_predictions.lm_head.dense.bias', 'lm_predictions.lm_head.dense.weight', 'lm_predictions.lm_head.LayerNorm.weight', 'deberta.embeddings.position_embeddings.weight']\n",
|
|
"- This IS expected if you are initializing DebertaV2ForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
|
|
"- This IS NOT expected if you are initializing DebertaV2ForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
|
|
"Some weights of DebertaV2ForMaskedLM were not initialized from the model checkpoint at microsoft/deberta-v2-xxlarge and are newly initialized: ['cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.bias']\n",
|
|
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"from tqdm import tqdm\n",
|
|
"import copy\n",
|
|
"import numpy as np\n",
|
|
"import torch\n",
|
|
"import torch.nn as nn\n",
|
|
"import torch.nn.functional as F\n",
|
|
"\n",
|
|
"from datasets import load_dataset\n",
|
|
"from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM\n",
|
|
"from sklearn.linear_model import LogisticRegression\n",
|
|
"\n",
|
|
"# Let's just try IMDB for simplicity\n",
|
|
"data = load_dataset(\"amazon_polarity\")[\"test\"]\n",
|
|
"\n",
|
|
"# Here are a few different model options you can play around with:\n",
|
|
"model_name = \"deberta\"\n",
|
|
"# model_name = \"gpt-j\"\n",
|
|
"# model_name = \"t5\"\n",
|
|
"\n",
|
|
"# if you want to cache the model weights somewhere, you can specify that here\n",
|
|
"cache_dir = None\n",
|
|
"\n",
|
|
"if model_name == \"deberta\":\n",
|
|
" model_type = \"encoder\"\n",
|
|
" tokenizer = AutoTokenizer.from_pretrained(\"microsoft/deberta-v2-xxlarge\", cache_dir=cache_dir)\n",
|
|
" model = AutoModelForMaskedLM.from_pretrained(\"microsoft/deberta-v2-xxlarge\", cache_dir=cache_dir)\n",
|
|
" model.cuda()\n",
|
|
"elif model_name == \"gpt-j\":\n",
|
|
" model_type = \"decoder\"\n",
|
|
" tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\", cache_dir=cache_dir)\n",
|
|
" model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\", cache_dir=cache_dir)\n",
|
|
" model.cuda()\n",
|
|
"elif model_name == \"t5\":\n",
|
|
" model_type = \"encoder_decoder\"\n",
|
|
" tokenizer = AutoTokenizer.from_pretrained(\"t5-11b\", cache_dir=cache_dir)\n",
|
|
" model = AutoModelForSeq2SeqLM.from_pretrained(\"t5-11b\", cache_dir=cache_dir)\n",
|
|
" model.parallelize() # T5 is big enough that we may need to run it on multiple GPUs\n",
|
|
"else:\n",
|
|
" print(\"Not implemented!\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## First let's write code for extracting hidden states given a model and text. \n",
|
|
"How we do this exactly will depend on the type of model."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def get_encoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
|
|
" \"\"\"\n",
|
|
" Given an encoder model and some text, gets the encoder hidden states (in a given layer, by default the last) \n",
|
|
" on that input text (where the full text is given to the encoder).\n",
|
|
"\n",
|
|
" Returns a numpy array of shape (hidden_dim,)\n",
|
|
" \"\"\"\n",
|
|
" # tokenize\n",
|
|
" encoder_text_ids = tokenizer(input_text, truncation=True, return_tensors=\"pt\").input_ids.to(model.device)\n",
|
|
"\n",
|
|
" # forward pass\n",
|
|
" with torch.no_grad():\n",
|
|
" output = model(encoder_text_ids, output_hidden_states=True)\n",
|
|
"\n",
|
|
" # get the appropriate hidden states\n",
|
|
" hs_tuple = output[\"hidden_states\"]\n",
|
|
" \n",
|
|
" hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
|
|
"\n",
|
|
" return hs\n",
|
|
"\n",
|
|
"def get_encoder_decoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
|
|
" \"\"\"\n",
|
|
" Given an encoder-decoder model and some text, gets the encoder hidden states (in a given layer, by default the last) \n",
|
|
" on that input text (where the full text is given to the encoder).\n",
|
|
"\n",
|
|
" Returns a numpy array of shape (hidden_dim,)\n",
|
|
" \"\"\"\n",
|
|
" # tokenize\n",
|
|
" encoder_text_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(model.device)\n",
|
|
" decoder_text_ids = tokenizer(\"\", return_tensors=\"pt\").input_ids.to(model.device)\n",
|
|
"\n",
|
|
" # forward pass\n",
|
|
" with torch.no_grad():\n",
|
|
" output = model(encoder_text_ids, decoder_input_ids=decoder_text_ids, output_hidden_states=True)\n",
|
|
"\n",
|
|
" # get the appropriate hidden states\n",
|
|
" hs_tuple = output[\"encoder_hidden_states\"]\n",
|
|
" hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
|
|
"\n",
|
|
" return hs\n",
|
|
"\n",
|
|
"def get_decoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
|
|
" \"\"\"\n",
|
|
" Given a decoder model and some text, gets the hidden states (in a given layer, by default the last) on that input text\n",
|
|
"\n",
|
|
" Returns a numpy array of shape (hidden_dim,)\n",
|
|
" \"\"\"\n",
|
|
" # tokenize (adding the EOS token this time)\n",
|
|
" input_ids = tokenizer(input_text + tokenizer.eos_token, return_tensors=\"pt\").input_ids.to(model.device)\n",
|
|
"\n",
|
|
" # forward pass\n",
|
|
" with torch.no_grad():\n",
|
|
" output = model(input_ids, output_hidden_states=True)\n",
|
|
"\n",
|
|
" # get the last layer, last token hidden states\n",
|
|
" hs_tuple = output[\"hidden_states\"]\n",
|
|
" hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
|
|
"\n",
|
|
" return hs\n",
|
|
"\n",
|
|
"def get_hidden_states(model, tokenizer, input_text, layer=-1, model_type=\"encoder\"):\n",
|
|
" fn = {\"encoder\": get_encoder_hidden_states, \"encoder_decoder\": get_encoder_decoder_hidden_states,\n",
|
|
" \"decoder\": get_decoder_hidden_states}[model_type]\n",
|
|
"\n",
|
|
" return fn(model, tokenizer, input_text, layer=layer)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Now let's write code for formatting data and for getting all the hidden states."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def format_imdb(text, label):\n",
|
|
" \"\"\"\n",
|
|
" Given an imdb example (\"text\") and corresponding label (0 for negative, or 1 for positive), \n",
|
|
" returns a zero-shot prompt for that example (which includes that label as the answer).\n",
|
|
" \n",
|
|
" (This is just one example of a simple, manually created prompt.)\n",
|
|
" \"\"\"\n",
|
|
" return \"The following movie review expresses a \" + [\"negative\", \"positive\"][label] + \" sentiment:\\n\" + text\n",
|
|
"\n",
|
|
"\n",
|
|
"def get_hidden_states_many_examples(model, tokenizer, data, model_type, n=100):\n",
|
|
" \"\"\"\n",
|
|
" Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.\n",
|
|
" Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)\n",
|
|
" with the ground truth labels\n",
|
|
" \n",
|
|
" This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency\n",
|
|
" \"\"\"\n",
|
|
" # setup\n",
|
|
" model.eval()\n",
|
|
" all_neg_hs, all_pos_hs, all_gt_labels = [], [], []\n",
|
|
"\n",
|
|
" # loop\n",
|
|
" for _ in tqdm(range(n)):\n",
|
|
" # for simplicity, sample a random example until we find one that's a reasonable length\n",
|
|
" # (most examples should be a reasonable length, so this is just to make sure)\n",
|
|
" while True:\n",
|
|
" idx = np.random.randint(len(data))\n",
|
|
" text, true_label = data[idx][\"content\"], data[idx][\"label\"]\n",
|
|
" # the actual formatted input will be longer, so include a bit of a marign\n",
|
|
" if len(tokenizer(text)) < 400: \n",
|
|
" break\n",
|
|
" \n",
|
|
" # get hidden states\n",
|
|
" neg_hs = get_hidden_states(model, tokenizer, format_imdb(text, 0), model_type=model_type)\n",
|
|
" pos_hs = get_hidden_states(model, tokenizer, format_imdb(text, 1), model_type=model_type)\n",
|
|
"\n",
|
|
" # collect\n",
|
|
" all_neg_hs.append(neg_hs)\n",
|
|
" all_pos_hs.append(pos_hs)\n",
|
|
" all_gt_labels.append(true_label)\n",
|
|
"\n",
|
|
" all_neg_hs = np.stack(all_neg_hs)\n",
|
|
" all_pos_hs = np.stack(all_pos_hs)\n",
|
|
" all_gt_labels = np.stack(all_gt_labels)\n",
|
|
"\n",
|
|
" return all_neg_hs, all_pos_hs, all_gt_labels"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:19<00:00, 5.10it/s]\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"neg_hs, pos_hs, y = get_hidden_states_many_examples(model, tokenizer, data, model_type)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Let's verify that the model's representations are good\n",
|
|
"\n",
|
|
"Before trying CCS, let's make sure there exists a direction that classifies examples as true vs false with high accuracy; if logistic regression accuracy is bad, there's no hope of CCS doing well."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Logistic regression accuracy: 0.9\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# let's create a simple 50/50 train split (the data is already randomized)\n",
|
|
"n = len(y)\n",
|
|
"neg_hs_train, neg_hs_test = neg_hs[:n//2], neg_hs[n//2:]\n",
|
|
"pos_hs_train, pos_hs_test = pos_hs[:n//2], pos_hs[n//2:]\n",
|
|
"y_train, y_test = y[:n//2], y[n//2:]\n",
|
|
"\n",
|
|
"# for simplicity we can just take the difference between positive and negative hidden states\n",
|
|
"# (concatenating also works fine)\n",
|
|
"x_train = neg_hs_train - pos_hs_train\n",
|
|
"x_test = neg_hs_test - pos_hs_test\n",
|
|
"\n",
|
|
"lr = LogisticRegression(class_weight=\"balanced\")\n",
|
|
"lr.fit(x_train, y_train)\n",
|
|
"print(\"Logistic regression accuracy: {}\".format(lr.score(x_test, y_test)))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Now let's try CCS"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"class MLPProbe(nn.Module):\n",
|
|
" def __init__(self, d):\n",
|
|
" super().__init__()\n",
|
|
" self.linear1 = nn.Linear(d, 100)\n",
|
|
" self.linear2 = nn.Linear(100, 1)\n",
|
|
"\n",
|
|
" def forward(self, x):\n",
|
|
" h = F.relu(self.linear1(x))\n",
|
|
" o = self.linear2(h)\n",
|
|
" return torch.sigmoid(o)\n",
|
|
"\n",
|
|
"class CCS(object):\n",
|
|
" def __init__(self, x0, x1, nepochs=1000, ntries=10, lr=1e-3, batch_size=-1, \n",
|
|
" verbose=False, device=\"cuda\", linear=True, weight_decay=0.01, var_normalize=False):\n",
|
|
" # data\n",
|
|
" self.var_normalize = var_normalize\n",
|
|
" self.x0 = self.normalize(x0)\n",
|
|
" self.x1 = self.normalize(x1)\n",
|
|
" self.d = self.x0.shape[-1]\n",
|
|
"\n",
|
|
" # training\n",
|
|
" self.nepochs = nepochs\n",
|
|
" self.ntries = ntries\n",
|
|
" self.lr = lr\n",
|
|
" self.verbose = verbose\n",
|
|
" self.device = device\n",
|
|
" self.batch_size = batch_size\n",
|
|
" self.weight_decay = weight_decay\n",
|
|
" \n",
|
|
" # probe\n",
|
|
" self.linear = linear\n",
|
|
" self.probe = self.initialize_probe()\n",
|
|
" self.best_probe = copy.deepcopy(self.probe)\n",
|
|
"\n",
|
|
" \n",
|
|
" def initialize_probe(self):\n",
|
|
" if self.linear:\n",
|
|
" self.probe = nn.Linear(self.d, 1)\n",
|
|
" else:\n",
|
|
" self.probe = MLPProbe(self.d)\n",
|
|
" self.probe.to(self.device) \n",
|
|
"\n",
|
|
"\n",
|
|
" def normalize(self, x):\n",
|
|
" \"\"\"\n",
|
|
" Mean-normalizes the data x (of shape (n, d))\n",
|
|
" If self.var_normalize, also divides by the standard deviation\n",
|
|
" \"\"\"\n",
|
|
" normalized_x = x - x.mean(axis=0, keepdims=True)\n",
|
|
" if self.var_normalize:\n",
|
|
" normalized_x /= normalized_x.std(axis=0, keepdims=True)\n",
|
|
"\n",
|
|
" return normalized_x\n",
|
|
"\n",
|
|
" \n",
|
|
" def get_tensor_data(self):\n",
|
|
" \"\"\"\n",
|
|
" Returns x0, x1 as appropriate tensors (rather than np arrays)\n",
|
|
" \"\"\"\n",
|
|
" x0 = torch.tensor(self.x0, dtype=torch.float, requires_grad=False, device=self.device)\n",
|
|
" x1 = torch.tensor(self.x1, dtype=torch.float, requires_grad=False, device=self.device)\n",
|
|
" return x0, x1\n",
|
|
" \n",
|
|
"\n",
|
|
" def get_loss(self, p0, p1):\n",
|
|
" \"\"\"\n",
|
|
" Returns the CCS loss for two probabilities each of shape (n,1) or (n,)\n",
|
|
" \"\"\"\n",
|
|
" informative_loss = (torch.min(p0, p1)**2).mean(0)\n",
|
|
" consistent_loss = ((p0 - (1-p1))**2).mean(0)\n",
|
|
" return informative_loss + consistent_loss\n",
|
|
"\n",
|
|
"\n",
|
|
" def get_acc(self, x0_test, x1_test, y_test):\n",
|
|
" \"\"\"\n",
|
|
" Computes accuracy for the current parameters on the given test inputs\n",
|
|
" \"\"\"\n",
|
|
" x0 = torch.tensor(self.normalize(x0_test), dtype=torch.float, requires_grad=False, device=self.device)\n",
|
|
" x1 = torch.tensor(self.normalize(x1_test), dtype=torch.float, requires_grad=False, device=self.device)\n",
|
|
" with torch.no_grad():\n",
|
|
" p0, p1 = self.best_probe(x0), self.best_probe(x1)\n",
|
|
" avg_confidence = 0.5*(p0 + (1-p1))\n",
|
|
" predictions = (avg_confidence.detach().cpu().numpy() < 0.5).astype(int)[:, 0]\n",
|
|
" acc = (predictions == y_test).mean()\n",
|
|
" acc = max(acc, 1 - acc)\n",
|
|
"\n",
|
|
" return acc\n",
|
|
" \n",
|
|
" \n",
|
|
" def train(self):\n",
|
|
" \"\"\"\n",
|
|
" Does a single training run of nepochs epochs\n",
|
|
" \"\"\"\n",
|
|
" x0, x1 = self.get_tensor_data()\n",
|
|
" permutation = torch.randperm(len(x0))\n",
|
|
" x0, x1 = x0[permutation], x1[permutation]\n",
|
|
" \n",
|
|
" # set up optimizer\n",
|
|
" optimizer = torch.optim.AdamW(self.probe.parameters(), lr=self.lr, weight_decay=self.weight_decay)\n",
|
|
" \n",
|
|
" batch_size = len(x0) if self.batch_size == -1 else self.batch_size\n",
|
|
" nbatches = len(x0) // batch_size\n",
|
|
"\n",
|
|
" # Start training (full batch)\n",
|
|
" for epoch in range(self.nepochs):\n",
|
|
" for j in range(nbatches):\n",
|
|
" x0_batch = x0[j*batch_size:(j+1)*batch_size]\n",
|
|
" x1_batch = x1[j*batch_size:(j+1)*batch_size]\n",
|
|
" \n",
|
|
" # probe\n",
|
|
" p0, p1 = self.probe(x0_batch), self.probe(x1_batch)\n",
|
|
"\n",
|
|
" # get the corresponding loss\n",
|
|
" loss = self.get_loss(p0, p1)\n",
|
|
"\n",
|
|
" # update the parameters\n",
|
|
" optimizer.zero_grad()\n",
|
|
" loss.backward()\n",
|
|
" optimizer.step()\n",
|
|
"\n",
|
|
" return loss.detach().cpu().item()\n",
|
|
" \n",
|
|
" def repeated_train(self):\n",
|
|
" best_loss = np.inf\n",
|
|
" for train_num in range(self.ntries):\n",
|
|
" self.initialize_probe()\n",
|
|
" loss = self.train()\n",
|
|
" if loss < best_loss:\n",
|
|
" self.best_probe = copy.deepcopy(self.probe)\n",
|
|
" best_loss = loss\n",
|
|
"\n",
|
|
" return best_loss"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"CCS accuracy: 0.88\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Train CCS without any labels\n",
|
|
"ccs = CCS(neg_hs_train, pos_hs_train)\n",
|
|
"ccs.repeated_train()\n",
|
|
"\n",
|
|
"# Evaluate\n",
|
|
"ccs_acc = ccs.get_acc(neg_hs_test, pos_hs_test, y_test)\n",
|
|
"print(\"CCS accuracy: {}\".format(ccs_acc))"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3 (ipykernel)",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.7.5"
|
|
},
|
|
"vscode": {
|
|
"interpreter": {
|
|
"hash": "b80286374679f2ad472c61c83fc267d31329b5dea8e2dcaccb727123767724c5"
|
|
}
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|