commit 1fb40eefbd9cba519688157aca111148a6d97600 Author: Collin Burns Date: Wed Dec 7 10:13:14 2022 -0800 Initial commit diff --git a/CCS.ipynb b/CCS.ipynb new file mode 100644 index 0000000..240e21c --- /dev/null +++ b/CCS.ipynb @@ -0,0 +1,491 @@ +{ + "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 + +

+ +We introduce a method for discovering truth-like features directly from model activations in a purely unsupervised way. + +## Abstract +> Existing techniques for training language models can be misaligned with the truth: if we train models with imitation learning, they may reproduce errors that humans make; if we train them to generate text that humans rate highly, they may output errors that human evaluators can't detect. We propose circumventing this issue by directly finding latent knowledge inside the internal activations of a language model in a purely unsupervised way. Specifically, we introduce a method for accurately answering yes-no questions given only unlabeled model activations. It works by finding a direction in activation space that satisfies logical consistency properties, such as that a statement and its negation have opposite truth values. We show that despite using no supervision and no model outputs, our method can recover diverse knowledge represented in large language models: across 6 models and 10 question-answering datasets, it outperforms zero-shot accuracy by 4\% on average. We also find that it cuts prompt sensitivity in half and continues to maintain high accuracy even when models are prompted to generate incorrect answers. Our results provide an initial step toward discovering what language models know, distinct from what they say, even when we don't have access to explicit ground truth labels. + +## Code + +We provide three options for code: +1. A notebook walking through our main method in a simple way: `CCS.ipynb`. This may be the best place to start if you want to understand the method better and play around with it a bit. +2. More flexible and efficient scripts for using our method in different settings: `generate.py` and `evaluate.py` (both of which rely heavily on `utils.py`). This code is a polished and simplified version of the code used for the paper. This may be the best place to build on if you want to build on our work. +3. You can also download the original (more comprehensive, but also more complicated and less polished) code [here](https://openreview.net/attachment?id=ETKGuby0hcs&name=supplementary_material). + +Below we provide usage details for our main python scripts (`generate.py` and `evaluate.py`). + +### Generation +First, use `generate.py` for (1) creating contrast pairs, and (2) generating hidden states from a model. For example, you can run: +``` +python generate.py --model_name deberta --num_examples 400 --batch_size 40 +``` +or +``` +python generate.py --model_name gpt-j --num_examples 100 --batch_size 20 +``` +or +``` +CUDA_VISIBLE_DEVICES=0,1 python generate.py --parallelize --model_name t5-11b --num_examples 100 +``` + +To use the decoder of an encoder-decoder model (which we found is worse than the encoder for T5 and UnifiedQA, but better than the encoder for T0), specify `--use_decoder`. + +There are also many optional flags for specifying the dataset (`--dataset`; the default is `imdb`), the cache directory for model weights (`--cache_dir`; the default is `None`), which prompt for the dataset to use (`--prompt_idx`; the default is `0`), where to save the hidden states for all layers in the model (`--all_layers`), and so on. + +### Evaluation +After generating hidden states, you can use `evaluate.py` for running our main method, CCS, on those hidden states. Simply run it with the same flags as you used when running `generate.py`, and it will load the correct hidden states for you. For example, if you ran `generate.py` with DeBERTa, generating 400 examples with a batch size of 40 (the first example in the Generation section), then you can run: +``` +python evaluate.py --model_name deberta --num_examples 400 --batch_size 40 +``` + +In addition to evaluating the performance of CCS, `evaluate.py` also verifies that logistic regression (LR) accuracy is reasonable. This can diagnose why CCS performance may be low; if LR accuracy is low, that suggestions that the model's representations aren't good enough for CCS to work well. + +### Requirements + +This code base was tested on Python 3.7.5 and PyTorch 1.12. It also uses the [datasets](https://pypi.org/project/datasets/) and [promptsource](https://github.com/bigscience-workshop/promptsource) packages for loading and formatting datasets. + +## Citation + +If you find this work helpful, please consider citing our paper: + + @article{burns2022dl, + title={Discovering Latent Knowledge in Language Models Without Supervision}, + author={Burns, Collin and Ye, Haotian and Klein, Dan and Steinhardt, Jacob}, + journal={ArXiV}, + year={2022} + } + diff --git a/evaluate.py b/evaluate.py new file mode 100644 index 0000000..8f73aad --- /dev/null +++ b/evaluate.py @@ -0,0 +1,53 @@ +from sklearn.linear_model import LogisticRegression +from utils import get_parser, load_all_generations, CCS + +def main(args, generation_args): + # load hidden states and labels + neg_hs, pos_hs, y = load_all_generations(generation_args) + + # Make sure the shape is correct + assert neg_hs.shape == pos_hs.shape + neg_hs, pos_hs = neg_hs[..., -1], pos_hs[..., -1] # take the last layer + if neg_hs.shape[1] == 1: # T5 may have an extra dimension; if so, get rid of it + neg_hs = neg_hs.squeeze(1) + pos_hs = pos_hs.squeeze(1) + + # Very simple train/test split (using the fact that the data is already shuffled) + neg_hs_train, neg_hs_test = neg_hs[:len(neg_hs) // 2], neg_hs[len(neg_hs) // 2:] + pos_hs_train, pos_hs_test = pos_hs[:len(pos_hs) // 2], pos_hs[len(pos_hs) // 2:] + y_train, y_test = y[:len(y) // 2], y[len(y) // 2:] + + # Make sure logistic regression accuracy is reasonable; otherwise our method won't have much of a chance of working + # you can also concatenate, but this works fine and is more comparable to CCS inputs + x_train = neg_hs_train - pos_hs_train + x_test = neg_hs_test - pos_hs_test + lr = LogisticRegression(class_weight="balanced") + lr.fit(x_train, y_train) + print("Logistic regression accuracy: {}".format(lr.score(x_test, y_test))) + + # Set up CCS. Note that you can usually just use the default args by simply doing ccs = CCS(neg_hs, pos_hs, y) + ccs = CCS(neg_hs_train, pos_hs_train, nepochs=args.nepochs, ntries=args.ntries, lr=args.lr, batch_size=args.ccs_batch_size, + verbose=args.verbose, device=args.ccs_device, linear=args.linear, weight_decay=args.weight_decay, + var_normalize=args.var_normalize) + + # train and evaluate CCS + ccs.repeated_train() + ccs_acc = ccs.get_acc(neg_hs_test, pos_hs_test, y_test) + print("CCS accuracy: {}".format(ccs_acc)) + + +if __name__ == "__main__": + parser = get_parser() + generation_args = parser.parse_args() # we'll use this to load the correct hidden states + labels + # We'll also add some additional args for evaluation + parser.add_argument("--nepochs", type=int, default=1000) + parser.add_argument("--ntries", type=int, default=10) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--ccs_batch_size", type=int, default=-1) + parser.add_argument("--verbose", action="store_true") + parser.add_argument("--ccs_device", type=str, default="cuda") + parser.add_argument("--linear", action="store_true") + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--var_normalize", action="store_true") + args = parser.parse_args() + main(args, generation_args) diff --git a/figure.png b/figure.png new file mode 100644 index 0000000..333721b Binary files /dev/null and b/figure.png differ diff --git a/generate.py b/generate.py new file mode 100644 index 0000000..06f077c --- /dev/null +++ b/generate.py @@ -0,0 +1,27 @@ +from utils import get_parser, load_model, get_dataloader, get_all_hidden_states, save_generations + +def main(args): + # Set up the model and data + print("Loading model") + model, tokenizer, model_type = load_model(args.model_name, args.cache_dir, args.parallelize, args.device) + + print("Loading dataloader") + dataloader = get_dataloader(args.dataset_name, args.split, tokenizer, args.prompt_idx, batch_size=args.batch_size, + num_examples=args.num_examples, model_type=model_type, use_decoder=args.use_decoder, device=args.device) + + # Get the hidden states and labels + print("Generating hidden states") + neg_hs, pos_hs, y = get_all_hidden_states(model, dataloader, layer=args.layer, all_layers=args.all_layers, + token_idx=args.token_idx, model_type=model_type, use_decoder=args.use_decoder) + + # Save the hidden states and labels + print("Saving hidden states") + save_generations(neg_hs, args, generation_type="negative_hidden_states") + save_generations(pos_hs, args, generation_type="positive_hidden_states") + save_generations(y, args, generation_type="labels") + + +if __name__ == "__main__": + parser = get_parser() + args = parser.parse_args() + main(args) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..d638390 --- /dev/null +++ b/utils.py @@ -0,0 +1,546 @@ +import os +import functools +import argparse +import copy + +import numpy as np +import pandas as pd +from tqdm import tqdm + +import torch +from torch.utils.data import Dataset, DataLoader +from torchvision import datasets +import torch.nn as nn +import torch.nn.functional as F + +# make sure to install promptsource, transformers, and datasets! +from promptsource.templates import DatasetTemplates +from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM +from datasets import load_dataset + + +############# Model loading and result saving ############# + +# Map each model name to its full Huggingface name; this is just for convenience for common models. You can run whatever model you'd like. +model_mapping = { + "gpt-j": "EleutherAI/gpt-j-6B", + "T0pp": "bigscience/T0pp", + "unifiedqa": "allenai/unifiedqa-t5-11b", + "T5": "t5-11b", + "deberta-mnli": "microsoft/deberta-xxlarge-v2-mnli", + "deberta": "microsoft/deberta-xxlarge-v2", + "roberta-mnli": "roberta-large-mnli", +} + + +def get_parser(): + """ + Returns the parser we will use for generate.py and evaluate.py + (We include it here so that we can use the same parser for both scripts) + """ + parser = argparse.ArgumentParser() + # setting up model + parser.add_argument("--model_name", type=str, default="T5", help="Name of the model to use") + parser.add_argument("--cache_dir", type=str, default=None, help="Cache directory for the model and tokenizer") + parser.add_argument("--parallelize", action="store_true", help="Whether to parallelize the model") + parser.add_argument("--device", type=str, default="cuda", help="Device to use for the model") + # setting up data + parser.add_argument("--dataset_name", type=str, default="imdb", help="Name of the dataset to use") + parser.add_argument("--split", type=str, default="test", help="Which split of the dataset to use") + parser.add_argument("--prompt_idx", type=int, default=0, help="Which prompt to use") + parser.add_argument("--batch_size", type=int, default=1, help="Batch size to use") + parser.add_argument("--num_examples", type=int, default=1000, help="Number of examples to generate") + # which hidden states we extract + parser.add_argument("--use_decoder", action="store_true", help="Whether to use the decoder; only relevant if model_type is encoder-decoder. Uses encoder by default (which usually -- but not always -- works better)") + parser.add_argument("--layer", type=int, default=-1, help="Which layer to use (if not all layers)") + parser.add_argument("--all_layers", action="store_true", help="Whether to use all layers or not") + parser.add_argument("--token_idx", type=int, default=-1, help="Which token to use (by default the last token)") + # saving the hidden states + parser.add_argument("--save_dir", type=str, default="generated_hidden_states", help="Directory to save the hidden states") + + return parser + + +def load_model(model_name, cache_dir=None, parallelize=False, device="cuda"): + """ + Loads a model and its corresponding tokenizer, either parallelized across GPUs (if the model permits that; usually just use this for T5-based models) or on a single GPU + """ + if model_name in model_mapping: + # use a nickname for our models + full_model_name = model_mapping[model_name] + else: + # if you're trying a new model, make sure it's the full name + full_model_name = model_name + + # use the right automodel, and get the corresponding model type + try: + model = AutoModelForSeq2SeqLM.from_pretrained(full_model_name, cache_dir=cache_dir) + model_type = "encoder_decoder" + except: + try: + model = AutoModelForMaskedLM.from_pretrained(full_model_name, cache_dir=cache_dir) + model_type = "encoder" + except: + model = AutoModelForCausalLM.from_pretrained(full_model_name, cache_dir=cache_dir) + model_type = "decoder" + + + # specify model_max_length (the max token length) to be 512 to ensure that padding works + # (it's not set by default for e.g. DeBERTa, but it's necessary for padding to work properly) + tokenizer = AutoTokenizer.from_pretrained(full_model_name, cache_dir=cache_dir, model_max_length=512) + model.eval() + + # put on the correct device + if parallelize: + model.parallelize() + else: + model = model.to(device) + + return model, tokenizer, model_type + + +def save_generations(generation, args, generation_type): + """ + Input: + generation: numpy array (e.g. hidden_states or labels) to save + args: arguments used to generate the hidden states. This is used for the filename to save to. + generation_type: one of "negative_hidden_states" or "positive_hidden_states" or "labels" + + Saves the generations to an appropriate directory. + """ + # construct the filename based on the args + arg_dict = vars(args) + exclude_keys = ["save_dir", "cache_dir", "device"] + filename = generation_type + "__" + "__".join(['{}_{}'.format(k, v) for k, v in arg_dict.items() if k not in exclude_keys]) + ".npy".format(generation_type) + + # create save directory if it doesn't exist + if not os.path.exists(args.save_dir): + os.makedirs(args.save_dir) + + # save + np.save(os.path.join(args.save_dir, filename), generation) + + +def load_single_generation(args, generation_type="hidden_states"): + # use the same filename as in save_generations + arg_dict = vars(args) + exclude_keys = ["save_dir", "cache_dir", "device"] + filename = generation_type + "__" + "__".join(['{}_{}'.format(k, v) for k, v in arg_dict.items() if k not in exclude_keys]) + ".npy".format(generation_type) + return np.load(os.path.join(args.save_dir, filename)) + + +def load_all_generations(args): + # load all the saved generations: neg_hs, pos_hs, and labels + neg_hs = load_single_generation(args, generation_type="negative_hidden_states") + pos_hs = load_single_generation(args, generation_type="positive_hidden_states") + labels = load_single_generation(args, generation_type="labels") + + return neg_hs, pos_hs, labels + + +############# Data ############# +class ContrastDataset(Dataset): + """ + Given a dataset and tokenizer (from huggingface), along with a collection of prompts for that dataset from promptsource and a corresponding prompt index, + returns a dataset that creates contrast pairs using that prompt + + Truncates examples larger than max_len, which can mess up contrast pairs, so make sure to only give it examples that won't be truncated. + """ + def __init__(self, raw_dataset, tokenizer, all_prompts, prompt_idx, + model_type="encoder_decoder", use_decoder=False, device="cuda"): + + # data and tokenizer + self.raw_dataset = raw_dataset + self.tokenizer = tokenizer + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + self.device = device + + # for formatting the answers + self.model_type = model_type + self.use_decoder = use_decoder + if self.use_decoder: + assert self.model_type != "encoder" + + # prompt + prompt_name_list = list(all_prompts.name_to_id_mapping.keys()) + self.prompt = all_prompts[prompt_name_list[prompt_idx]] + + def __len__(self): + return len(self.raw_dataset) + + def encode(self, nl_prompt): + """ + Tokenize a given natural language prompt (from after applying self.prompt to an example) + + For encoder-decoder models, we can either: + (1) feed both the question and answer to the encoder, creating contrast pairs using the encoder hidden states + (which uses the standard tokenization, but also passes the empty string to the decoder), or + (2) feed the question the encoder and the answer to the decoder, creating contrast pairs using the decoder hidden states + + If self.decoder is True we do (2), otherwise we do (1). + """ + # get question and answer from prompt + question, answer = nl_prompt + + # tokenize the question and answer (depending upon the model type and whether self.use_decoder is True) + if self.model_type == "encoder_decoder": + input_ids = self.get_encoder_decoder_input_ids(question, answer) + elif self.model_type == "encoder": + input_ids = self.get_encoder_input_ids(question, answer) + else: + input_ids = self.get_decoder_input_ids(question, answer) + + # get rid of the batch dimension since this will be added by the Dataloader + if input_ids["input_ids"].shape[0] == 1: + for k in input_ids: + input_ids[k] = input_ids[k].squeeze(0) + + return input_ids + + + def get_encoder_input_ids(self, question, answer): + """ + Format the input ids for encoder-only models; standard formatting. + """ + combined_input = question + " " + answer + input_ids = self.tokenizer(combined_input, truncation=True, padding="max_length", return_tensors="pt") + + return input_ids + + + def get_decoder_input_ids(self, question, answer): + """ + Format the input ids for encoder-only models. + This is the same as get_encoder_input_ids except that we add the EOS token at the end of the input (which apparently can matter) + """ + combined_input = question + " " + answer + self.tokenizer.eos_token + input_ids = self.tokenizer(combined_input, truncation=True, padding="max_length", return_tensors="pt") + + return input_ids + + + def get_encoder_decoder_input_ids(self, question, answer): + """ + Format the input ids for encoder-decoder models. + There are two cases for this, depending upon whether we want to use the encoder hidden states or the decoder hidden states. + """ + if self.use_decoder: + # feed the same question to the encoder but different answers to the decoder to construct contrast pairs + input_ids = self.tokenizer(question, truncation=True, padding="max_length", return_tensors="pt") + decoder_input_ids = self.tokenizer(answer, truncation=True, padding="max_length", return_tensors="pt") + else: + # include both the question and the answer in the input for the encoder + # feed the empty string to the decoder (i.e. just ignore it -- but it needs an input or it'll throw an error) + input_ids = self.tokenizer(question, answer, truncation=True, padding="max_length", return_tensors="pt") + decoder_input_ids = self.tokenizer("", return_tensors="pt") + + # move everything into input_ids so that it's easier to pass to the model + input_ids["decoder_input_ids"] = decoder_input_ids["input_ids"] + input_ids["decoder_attention_mask"] = decoder_input_ids["attention_mask"] + + return input_ids + + + def __getitem__(self, index): + # get the original example + data = self.raw_dataset[int(index)] + text, true_answer = data["text"], data["label"] + + # get the possible labels + # (for simplicity assume the binary case for contrast pairs) + label_list = self.prompt.get_answer_choices_list(data) + assert len(label_list) == 2, print("Make sure there are only two possible answers! Actual number of answers:", label_list) + + # reconvert to dataset format but with fake/candidate labels to create the contrast pair + neg_example = {"text": text, "label": 0} + pos_example = {"text": text, "label": 1} + + # construct contrast pairs by answering the prompt with the two different possible labels + # (for example, label 0 might be mapped to "no" and label 1 might be mapped to "yes") + neg_prompt, pos_prompt = self.prompt.apply(neg_example), self.prompt.apply(pos_example) + + # tokenize + neg_ids, pos_ids = self.encode(neg_prompt), self.encode(pos_prompt) + + # verify these are different (e.g. tokenization didn't cut off the difference between them) + if self.use_decoder and self.model_type == "encoder_decoder": + assert (neg_ids["decoder_input_ids"] - pos_ids["decoder_input_ids"]).sum() != 0, print("The decoder_input_ids for the contrast pairs are the same!", neg_ids, pos_ids) + else: + assert (neg_ids["input_ids"] - pos_ids["input_ids"]).sum() != 0, print("The input_ids for the contrast pairs are the same!", neg_ids, pos_ids) + + # return the tokenized inputs, the text prompts, and the true label + return neg_ids, pos_ids, neg_prompt, pos_prompt, true_answer + + +def get_dataloader(dataset_name, split, tokenizer, prompt_idx, batch_size=16, num_examples=1000, + model_type="encoder_decoder", use_decoder=False, device="cuda", pin_memory=True, num_workers=1): + """ + Creates a dataloader for a given dataset (and its split), tokenizer, and prompt index + + Takes a random subset of (at most) num_examples samples from the dataset that are not truncated by the tokenizer. + """ + # load the raw dataset + raw_dataset = load_dataset(dataset_name)[split] + + # load all the prompts for that dataset + all_prompts = DatasetTemplates(dataset_name) + + # create the ConstrastDataset + contrast_dataset = ContrastDataset(raw_dataset, tokenizer, all_prompts, prompt_idx, + model_type=model_type, use_decoder=use_decoder, + device=device) + + # get a random permutation of the indices; we'll take the first num_examples of these that do not get truncated + random_idxs = np.random.permutation(len(contrast_dataset)) + + # remove examples that would be truncated (since this messes up contrast pairs) + prompt_name_list = list(all_prompts.name_to_id_mapping.keys()) + prompt = all_prompts[prompt_name_list[prompt_idx]] + keep_idxs = [] + for idx in random_idxs: + question, answer = prompt.apply(raw_dataset[int(idx)]) + input_text = question + " " + answer + if len(tokenizer.encode(input_text, truncation=False)) < tokenizer.model_max_length - 2: # include small margin to be conservative + keep_idxs.append(idx) + if len(keep_idxs) >= num_examples: + break + + # create and return the corresponding dataloader + subset_dataset = torch.utils.data.Subset(contrast_dataset, keep_idxs) + dataloader = DataLoader(subset_dataset, batch_size=batch_size, shuffle=False, pin_memory=pin_memory, num_workers=num_workers) + + return dataloader + + +############# Hidden States ############# +def get_first_mask_loc(mask, shift=False): + """ + return the location of the first pad token for the given ids, which corresponds to a mask value of 0 + if there are no pad tokens, then return the last location + """ + # add a 0 to the end of the mask in case there are no pad tokens + mask = torch.cat([mask, torch.zeros_like(mask[..., :1])], dim=-1) + + if shift: + mask = mask[..., 1:] + + # get the location of the first pad token; use the fact that torch.argmax() returns the first index in the case of ties + first_mask_loc = torch.argmax((mask == 0).int(), dim=-1) + + return first_mask_loc + + +def get_individual_hidden_states(model, batch_ids, layer=None, all_layers=True, token_idx=-1, model_type="encoder_decoder", use_decoder=False): + """ + Given a model and a batch of tokenized examples, returns the hidden states for either + a specified layer (if layer is a number) or for all layers (if all_layers is True). + + If specify_encoder is True, uses "encoder_hidden_states" instead of "hidden_states" + This is necessary for getting the encoder hidden states for encoder-decoder models, + but it is not necessary for encoder-only or decoder-only models. + """ + if use_decoder: + assert "decoder" in model_type + + # forward pass + with torch.no_grad(): + batch_ids = batch_ids.to(model.device) + output = model(**batch_ids, output_hidden_states=True) + + # get all the corresponding hidden states (which is a tuple of length num_layers) + if use_decoder and "decoder_hidden_states" in output.keys(): + hs_tuple = output["decoder_hidden_states"] + elif "encoder_hidden_states" in output.keys(): + hs_tuple = output["encoder_hidden_states"] + else: + hs_tuple = output["hidden_states"] + + # just get the corresponding layer hidden states + if all_layers: + # stack along the last axis so that it's easier to consistently index the first two axes + hs = torch.stack([h.squeeze().detach().cpu() for h in hs_tuple], axis=-1) # (bs, seq_len, dim, num_layers) + else: + assert layer is not None + hs = hs_tuple[layer].unsqueeze(-1).detach().cpu() # (bs, seq_len, dim, 1) + + # we want to get the token corresponding to token_idx while ignoring the masked tokens + if token_idx == 0: + final_hs = hs[:, 0] # (bs, dim, num_layers) + else: + # if token_idx == -1, then takes the hidden states corresponding to the last non-mask tokens + # first we need to get the first mask location for each example in the batch + assert token_idx < 0, print("token_idx must be either 0 or negative, but got", token_idx) + mask = batch_ids["decoder_attention_mask"] if (model_type == "encoder_decoder" and use_decoder) else batch_ids["attention_mask"] + first_mask_loc = get_first_mask_loc(mask).squeeze() + final_hs = hs[torch.arange(hs.size(0)), first_mask_loc+token_idx] # (bs, dim, num_layers) + + return final_hs + + +def get_all_hidden_states(model, dataloader, layer=None, all_layers=True, token_idx=-1, model_type="encoder_decoder", use_decoder=False): + """ + Given a model, a tokenizer, and a dataloader, returns the hidden states (corresponding to a given position index) in all layers for all examples in the dataloader, + along with the average log probs corresponding to the answer tokens + + The dataloader should correspond to examples *with a candidate label already added* to each example. + E.g. this function should be used for "Q: Is 2+2=5? A: True" or "Q: Is 2+2=5? A: False", but NOT for "Q: Is 2+2=5? A: ". + """ + all_pos_hs, all_neg_hs = [], [] + all_gt_labels = [] + + model.eval() + for batch in tqdm(dataloader): + neg_ids, pos_ids, _, _, gt_label = batch + + neg_hs = get_individual_hidden_states(model, neg_ids, layer=layer, all_layers=all_layers, token_idx=token_idx, + model_type=model_type, use_decoder=use_decoder) + pos_hs = get_individual_hidden_states(model, pos_ids, layer=layer, all_layers=all_layers, token_idx=token_idx, + model_type=model_type, use_decoder=use_decoder) + + if dataloader.batch_size == 1: + neg_hs, pos_hs = neg_hs.unsqueeze(0), pos_hs.unsqueeze(0) + + all_neg_hs.append(neg_hs) + all_pos_hs.append(pos_hs) + all_gt_labels.append(gt_label) + + all_neg_hs = np.concatenate(all_neg_hs, axis=0) + all_pos_hs = np.concatenate(all_pos_hs, axis=0) + all_gt_labels = np.concatenate(all_gt_labels, axis=0) + + return all_neg_hs, all_pos_hs, all_gt_labels + +############# CCS ############# +class MLPProbe(nn.Module): + def __init__(self, d): + super().__init__() + self.linear1 = nn.Linear(d, 100) + self.linear2 = nn.Linear(100, 1) + + def forward(self, x): + h = F.relu(self.linear1(x)) + o = self.linear2(h) + return torch.sigmoid(o) + +class CCS(object): + def __init__(self, x0, x1, nepochs=1000, ntries=10, lr=1e-3, batch_size=-1, + verbose=False, device="cuda", linear=True, weight_decay=0.01, var_normalize=False): + # data + self.var_normalize = var_normalize + self.x0 = self.normalize(x0) + self.x1 = self.normalize(x1) + self.d = self.x0.shape[-1] + + # training + self.nepochs = nepochs + self.ntries = ntries + self.lr = lr + self.verbose = verbose + self.device = device + self.batch_size = batch_size + self.weight_decay = weight_decay + + # probe + self.linear = linear + self.probe = self.initialize_probe() + self.best_probe = copy.deepcopy(self.probe) + + + def initialize_probe(self): + if self.linear: + self.probe = nn.Linear(self.d, 1) + else: + self.probe = MLPProbe(self.d) + self.probe.to(self.device) + + + def normalize(self, x): + """ + Mean-normalizes the data x (of shape (n, d)) + If self.var_normalize, also divides by the standard deviation + """ + normalized_x = x - x.mean(axis=0, keepdims=True) + if self.var_normalize: + normalized_x /= normalized_x.std(axis=0, keepdims=True) + + return normalized_x + + + def get_tensor_data(self): + """ + Returns x0, x1 as appropriate tensors (rather than np arrays) + """ + x0 = torch.tensor(self.x0, dtype=torch.float, requires_grad=False, device=self.device) + x1 = torch.tensor(self.x1, dtype=torch.float, requires_grad=False, device=self.device) + return x0, x1 + + + def get_loss(self, p0, p1): + """ + Returns the CCS loss for two probabilities each of shape (n,1) or (n,) + """ + informative_loss = (torch.min(p0, p1)**2).mean(0) + consistent_loss = ((p0 - (1-p1))**2).mean(0) + return informative_loss + consistent_loss + + + def get_acc(self, x0_test, x1_test, y_test): + """ + Computes accuracy for the current parameters on the given test inputs + """ + x0 = torch.tensor(self.normalize(x0_test), dtype=torch.float, requires_grad=False, device=self.device) + x1 = torch.tensor(self.normalize(x1_test), dtype=torch.float, requires_grad=False, device=self.device) + with torch.no_grad(): + p0, p1 = self.best_probe(x0), self.best_probe(x1) + avg_confidence = 0.5*(p0 + (1-p1)) + predictions = (avg_confidence.detach().cpu().numpy() < 0.5).astype(int)[:, 0] + acc = (predictions == y_test).mean() + acc = max(acc, 1 - acc) + + return acc + + + def train(self): + """ + Does a single training run of nepochs epochs + """ + x0, x1 = self.get_tensor_data() + permutation = torch.randperm(len(x0)) + x0, x1 = x0[permutation], x1[permutation] + + # set up optimizer + optimizer = torch.optim.AdamW(self.probe.parameters(), lr=self.lr, weight_decay=self.weight_decay) + + batch_size = len(x0) if self.batch_size == -1 else self.batch_size + nbatches = len(x0) // batch_size + + # Start training (full batch) + for epoch in range(self.nepochs): + for j in range(nbatches): + x0_batch = x0[j*batch_size:(j+1)*batch_size] + x1_batch = x1[j*batch_size:(j+1)*batch_size] + + # probe + p0, p1 = self.probe(x0_batch), self.probe(x1_batch) + + # get the corresponding loss + loss = self.get_loss(p0, p1) + + # update the parameters + optimizer.zero_grad() + loss.backward() + optimizer.step() + + return loss.detach().cpu().item() + + def repeated_train(self): + best_loss = np.inf + for train_num in range(self.ntries): + self.initialize_probe() + loss = self.train() + if loss < best_loss: + self.best_probe = copy.deepcopy(self.probe) + best_loss = loss + + return best_loss