Section 5.2 (mem score).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "310fbd4d-f005-4ef1-a192-154f47d721e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DistilBertTokenizerFast, AutoModel, DistilBertModel\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import copy\n",
    "\n",
    "# code is based on repo: https://github.com/xszheng2020/memorization\n",
    "# data is from https://github.com/xszheng2020/memorization/tree/master/sst/data\n",
    "train = pd.read_csv(\"train.csv\")\n",
    "val = pd.read_csv(\"dev.csv\")\n",
    "test = pd.read_csv(\"test.csv\")\n",
    "num_classes = train[\"label\"].nunique()\n",
    "\n",
    "tokenizer = DistilBertTokenizerFast.from_pretrained(\"distilbert-base-uncased\")\n",
    "\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        row = self.data.iloc[index]\n",
    "        label = row[0]\n",
    "        sentence = row[1]\n",
    "        return label, sentence\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "\n",
    "class CustomDatasetWithMask(Dataset):\n",
    "    def __init__(self, data, mask=None):\n",
    "        self.data = data\n",
    "        self.mask = mask\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        row = self.data.iloc[index]\n",
    "        label = row[0]\n",
    "        sentence = row[1]\n",
    "        sample_index = row[2]\n",
    "        mask = self.mask[sample_index]\n",
    "        return label, sentence, mask\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a131dc6-fc40-4aa3-b78d-74922ee05042",
   "metadata": {},
   "source": [
    "### Finetune model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c6fba328-bafb-4399-a9d7-9406c1c461fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "    ):\n",
    "        super(CustomModel, self).__init__()\n",
    "        self.bert = DistilBertModel.from_pretrained(\n",
    "            \"distilbert-base-uncased\",\n",
    "            output_hidden_states=False,\n",
    "            output_attentions=False,\n",
    "        )\n",
    "        self.classifier = nn.Linear(768, num_classes)\n",
    "\n",
    "    def forward(self, inputs, is_ids, attention_mask):\n",
    "        if is_ids:\n",
    "            last_hidden = self.bert(input_ids=inputs, attention_mask=attention_mask)[0]\n",
    "        else:\n",
    "            last_hidden = self.bert(\n",
    "                inputs_embeds=inputs, attention_mask=attention_mask\n",
    "            )[0]\n",
    "        ####\n",
    "        cls_embedding = last_hidden[:, 0, :]  # (bs, dim) pooled_output = cls_embedding\n",
    "        ####\n",
    "        logits = self.classifier(cls_embedding)  # (bs, num_labels)\n",
    "        ####\n",
    "        return cls_embedding, logits\n",
    "\n",
    "\n",
    "model = CustomModel()\n",
    "model.cuda()\n",
    "\n",
    "\n",
    "def collate_fn(batch):\n",
    "    label = [b[0] for b in batch]\n",
    "    text_a = [b[1] for b in batch]\n",
    "    my_dict = tokenizer(\n",
    "        text_a,\n",
    "        None,\n",
    "        add_special_tokens=True,\n",
    "        padding=True,\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        return_attention_mask=True,\n",
    "        return_tensors=\"pt\",\n",
    "    )\n",
    "    label = torch.tensor(label)\n",
    "    return label, my_dict[\"input_ids\"], my_dict[\"attention_mask\"]\n",
    "\n",
    "\n",
    "BATCH_SIZE = 100\n",
    "EPOCHS = 150\n",
    "early_stop_steps = 20\n",
    "from_last_step = 0\n",
    "# Prepare dataloaders\n",
    "train_dataset = CustomDataset(data=train)\n",
    "train_dataloader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    num_workers=0,\n",
    "    shuffle=False,\n",
    "    collate_fn=collate_fn,\n",
    ")\n",
    "\n",
    "val_dataset = CustomDataset(data=val)\n",
    "val_dataloader = DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    num_workers=0,\n",
    "    shuffle=False,\n",
    "    collate_fn=collate_fn,\n",
    ")\n",
    "\n",
    "\n",
    "model = CustomModel()\n",
    "model.cuda()\n",
    "optimizer_grouped_parameters = [\n",
    "    {\"params\": [p for n, p in model.named_parameters() if (p.requires_grad == True)]},\n",
    "]\n",
    "optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=0.000001)\n",
    "\n",
    "\n",
    "best_val_acc = 0\n",
    "best_state_dict = None\n",
    "\n",
    "# train model\n",
    "for i in range(EPOCHS):\n",
    "    model.train()\n",
    "    train_correct = 0\n",
    "    total_train_loss = 0\n",
    "    total_val_loss = 0\n",
    "    val_correct = 0\n",
    "\n",
    "    for batch in train_dataloader:\n",
    "        ####\n",
    "        model.zero_grad()\n",
    "        optimizer.zero_grad()\n",
    "        ####\n",
    "        label = batch[0].cuda()\n",
    "        input_ids = batch[1].cuda()\n",
    "        attention_mask = batch[2].cuda()\n",
    "\n",
    "        _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "\n",
    "        loss = F.cross_entropy(outputs, label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_train_loss += loss.item()\n",
    "        train_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in val_dataloader:\n",
    "            label = batch[0].cuda()\n",
    "            input_ids = batch[1].cuda()\n",
    "            attention_mask = batch[2].cuda()\n",
    "            _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "            total_val_loss += loss.item()\n",
    "            val_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "\n",
    "    val_acc = val_correct / len(val)\n",
    "    # save best checkpoint\n",
    "    if val_acc >= best_val_acc:\n",
    "        best_val_acc = val_acc\n",
    "        best_state_dict = copy.deepcopy(model.state_dict())\n",
    "        from_last_step = 0\n",
    "    # early stop\n",
    "    if from_last_step >= early_stop_steps:\n",
    "        break\n",
    "    from_last_step += 1\n",
    "    if i % 10 == 0:\n",
    "        print(\n",
    "            i,\n",
    "            \"train_acc:\",\n",
    "            round(train_correct / len(train), 6),\n",
    "            \"train_loss:\",\n",
    "            round(total_train_loss / len(train), 6),\n",
    "            \"val_loss:\",\n",
    "            round(total_val_loss / len(val), 6),\n",
    "            \"val_acc:\",\n",
    "            round(val_correct / len(val), 6),\n",
    "        )\n",
    "# load best model\n",
    "model.load_state_dict(best_state_dict)\n",
    "best_val_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1cd70224-f254-4f81-8eb4-fbd10ba8c7cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "L2_LAMBDA = 5e-3\n",
    "# Set random seed\n",
    "def set_seeds(seed):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "\n",
    "\n",
    "set_seeds(42)\n",
    "\n",
    "\n",
    "def compute_s(model, v, train_data_loader, damp, scale, num_samples):\n",
    "\n",
    "    last_estimate = list(v).copy()\n",
    "    for i, batch in enumerate(train_data_loader):\n",
    "        ####\n",
    "        label = batch[0].cuda()\n",
    "        input_ids = batch[1].cuda()\n",
    "        ####\n",
    "        attention_mask = batch[2].cuda()\n",
    "        ####\n",
    "        this_estimate = compute_hessian_vector_products(\n",
    "            model=model,\n",
    "            vectors=last_estimate,\n",
    "            label=label,\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "        )\n",
    "        # Recursively caclulate h_estimate\n",
    "        with torch.no_grad():\n",
    "            new_estimate = [\n",
    "                a + (1 - damp) * b - c / scale\n",
    "                for a, b, c in zip(v, last_estimate, this_estimate)\n",
    "            ]\n",
    "        ####\n",
    "\n",
    "        new_estimate_norm = new_estimate[0].norm().item()\n",
    "        last_estimate_norm = last_estimate[0].norm().item()\n",
    "        estimate_norm_diff = new_estimate_norm - last_estimate_norm\n",
    "        ####\n",
    "        last_estimate = new_estimate\n",
    "\n",
    "        if i > num_samples:  # should be i>=(num_samples-1) but does not matters\n",
    "            break\n",
    "\n",
    "    # References:\n",
    "    # https://github.com/kohpangwei/influence-release/blob/master/influence/genericNeuralNet.py#L475\n",
    "    # Do this for each iteration of estimation\n",
    "    # Since we use one estimation, we put this at the end\n",
    "    inverse_hvp = [X / scale for X in last_estimate]\n",
    "\n",
    "    return inverse_hvp\n",
    "\n",
    "\n",
    "def compute_hessian_vector_products(model, vectors, label, input_ids, attention_mask):\n",
    "    ####\n",
    "    _, outputs = model(\n",
    "        inputs=input_ids,\n",
    "        is_ids=True,\n",
    "        attention_mask=attention_mask,\n",
    "    )\n",
    "    ce_loss = F.cross_entropy(outputs, label)\n",
    "    ####\n",
    "    hack_loss = torch.cat(\n",
    "        [\n",
    "            (p**2).view(-1)\n",
    "            for n, p in model.named_parameters()\n",
    "            if ((not any(nd in n for nd in no_decay)) and (p.requires_grad == True))\n",
    "        ]\n",
    "    ).sum() * (L2_LAMBDA)\n",
    "    ####\n",
    "    loss = ce_loss + hack_loss\n",
    "    ####\n",
    "    model.zero_grad()\n",
    "    grad_tuple = torch.autograd.grad(\n",
    "        outputs=loss,\n",
    "        inputs=[\n",
    "            param for name, param in model.named_parameters() if param.requires_grad\n",
    "        ],\n",
    "        create_graph=True,\n",
    "    )\n",
    "    ####\n",
    "    # model.zero_grad()\n",
    "    grad_grad_tuple = torch.autograd.grad(\n",
    "        outputs=grad_tuple,\n",
    "        inputs=[\n",
    "            param for name, param in model.named_parameters() if param.requires_grad\n",
    "        ],\n",
    "        grad_outputs=vectors,\n",
    "        only_inputs=True,\n",
    "    )\n",
    "\n",
    "    return grad_grad_tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39bc9dbd-e6c2-4be9-9cf8-02a3beaa367a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Calculate memorization scores for finetuned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "70631918-16e2-4ec5-9730-16ced24c53ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from contexttimer import Timer\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    train_dataset, batch_size=1, num_workers=0, shuffle=False, collate_fn=collate_fn\n",
    ")\n",
    "\n",
    "my_hook_1 = {}\n",
    "my_hook_2 = {}\n",
    "\n",
    "emb = None\n",
    "\n",
    "\n",
    "def hook_func_1(module, input_, output):\n",
    "    my_hook_1[\"out\"] = output\n",
    "\n",
    "\n",
    "def hook_func_2(module, input_, output):\n",
    "    output.data = emb.data\n",
    "    my_hook_2[\"out\"] = output\n",
    "    # my_hook_2['out'].requires_grad = True # Important # not needed for model without freeze\n",
    "\n",
    "####\n",
    "no_decay = [\"bias\", \"output_layer_norm.weight\"]\n",
    "output_collections = []\n",
    "####\n",
    "for idx, batch in enumerate(tqdm(train_dataloader)):\n",
    "    ####\n",
    "    z_label = batch[0].cuda()\n",
    "    z_input_ids = batch[1].cuda()\n",
    "    z_attention_mask = batch[2].cuda()\n",
    "    ####\n",
    "    row = train.iloc[idx]\n",
    "    tokens = tokenizer.convert_ids_to_tokens(z_input_ids[0].cpu().numpy())\n",
    "    ####\n",
    "    baseline = z_input_ids.clone()\n",
    "    baseline[baseline != 0] = 103\n",
    "\n",
    "    hook_1 = model.bert.transformer.layer[0].output_layer_norm.register_forward_hook(\n",
    "        hook_func_1\n",
    "    )\n",
    "    _, _ = model(inputs=baseline, is_ids=True, attention_mask=z_attention_mask)\n",
    "    bemb = my_hook_1[\"out\"].clone()\n",
    "    hook_1.remove()\n",
    "    ####\n",
    "    hook_1 = model.bert.transformer.layer[0].output_layer_norm.register_forward_hook(\n",
    "        hook_func_1\n",
    "    )\n",
    "    _, outputs = model(inputs=z_input_ids, is_ids=True, attention_mask=z_attention_mask)\n",
    "    wemb = my_hook_1[\"out\"].clone()\n",
    "    hook_1.remove()\n",
    "    ####\n",
    "    prob = F.softmax(outputs, dim=-1)\n",
    "    prediction = torch.argmax(prob, dim=1)\n",
    "\n",
    "    prob_gt = torch.gather(prob, 1, z_label.unsqueeze(1))\n",
    "    ####\n",
    "    model.zero_grad()\n",
    "\n",
    "    v = torch.autograd.grad(\n",
    "        outputs=prob_gt,\n",
    "        inputs=[\n",
    "            param for name, param in model.named_parameters() if param.requires_grad\n",
    "        ],\n",
    "        create_graph=False,\n",
    "    )\n",
    "    ####\n",
    "    for repetition in range(4):\n",
    "        with Timer() as timer:\n",
    "            ####\n",
    "            train_dataloader = DataLoader(\n",
    "                train_dataset,\n",
    "                batch_size=1,\n",
    "                num_workers=0,\n",
    "                shuffle=True,\n",
    "                # pin_memory=True,\n",
    "                collate_fn=collate_fn,\n",
    "            )\n",
    "            ####\n",
    "            s = compute_s(\n",
    "                model=model,\n",
    "                v=v,\n",
    "                train_data_loader=train_dataloader,\n",
    "                damp=5e-3,\n",
    "                scale=1e4,\n",
    "                num_samples=100,\n",
    "            )\n",
    "            ####\n",
    "            time_elapsed = timer.elapsed\n",
    "        # print(f\"{time_elapsed:.2f} seconds\")\n",
    "        ####\n",
    "        hessian = None\n",
    "        steps = 50\n",
    "        ####\n",
    "        for alpha in np.linspace(\n",
    "            0, 1.0, num=steps + 1, endpoint=True\n",
    "        ):  # right Riemann sum\n",
    "            emb = bemb.clone() + alpha * (wemb.clone() - bemb.clone())\n",
    "            # emb.requires_grad = True\n",
    "            ####\n",
    "            hook_2 = model.bert.transformer.layer[\n",
    "                0\n",
    "            ].output_layer_norm.register_forward_hook(hook_func_2)\n",
    "            _, outputs = model(\n",
    "                inputs=z_input_ids, is_ids=True, attention_mask=z_attention_mask\n",
    "            )\n",
    "            hook_2.remove()\n",
    "            ####\n",
    "            ce_loss_gt = F.cross_entropy(outputs, z_label)\n",
    "            z_hack_loss = torch.cat(\n",
    "                [\n",
    "                    (p**2).view(-1)\n",
    "                    for n, p in model.named_parameters()\n",
    "                    if (\n",
    "                        (not any(nd in n for nd in no_decay))\n",
    "                        and (p.requires_grad == True)\n",
    "                    )\n",
    "                ]\n",
    "            ).sum() * (L2_LAMBDA)\n",
    "            ####\n",
    "            model.zero_grad()\n",
    "\n",
    "            grad_tuple_ = torch.autograd.grad(\n",
    "                outputs=ce_loss_gt + z_hack_loss,\n",
    "                inputs=[\n",
    "                    param\n",
    "                    for name, param in model.named_parameters()\n",
    "                    if param.requires_grad\n",
    "                ],\n",
    "                create_graph=True,\n",
    "            )\n",
    "            # model.zero_grad()\n",
    "            dot = nn.utils.parameters_to_vector(\n",
    "                s\n",
    "            ).detach() @ nn.utils.parameters_to_vector(\n",
    "                grad_tuple_\n",
    "            )  # scalar\n",
    "            grad_grad_tuple_ = torch.autograd.grad(\n",
    "                outputs=dot, inputs=[my_hook_2[\"out\"]], only_inputs=True\n",
    "            )\n",
    "            ####\n",
    "            if alpha == 0:\n",
    "                hessian = grad_grad_tuple_[0]\n",
    "                influence_prime = [-torch.sum(x * y) for x, y in zip(s, grad_tuple_)]\n",
    "                influence_prime = sum(influence_prime).item()\n",
    "            elif alpha == 1.0:\n",
    "                break\n",
    "            else:\n",
    "                hessian += grad_grad_tuple_[0]\n",
    "        ####\n",
    "        hessian = hessian / steps\n",
    "        ####\n",
    "        influence = [-torch.sum(x * y) for x, y in zip(s, grad_tuple_)]\n",
    "        influence = sum(influence).item()\n",
    "        ####\n",
    "        result = hessian * (wemb - bemb)\n",
    "        theta = torch.sum(result).detach().cpu().numpy()\n",
    "        attributions = torch.sum(result, dim=-1)[0].detach().cpu().numpy()\n",
    "\n",
    "        outputs = {\n",
    "            \"index\": idx,\n",
    "            \"sentence\": row[\"sentence\"],\n",
    "            \"label\": row[\"label\"],\n",
    "            \"prob\": prob.detach().cpu().numpy()[0],\n",
    "            \"prediction\": prediction.detach().cpu().numpy()[0],\n",
    "            \"influence_prime\": influence_prime,\n",
    "            \"influence\": influence,\n",
    "            \"diff\": influence_prime - influence,\n",
    "            \"theta\": theta,\n",
    "            \"tokens\": tokens,\n",
    "            \"attributions\": attributions,\n",
    "            \"repetition\": repetition,\n",
    "            \"time_elapsed\": time_elapsed,\n",
    "        }\n",
    "\n",
    "        output_collections.append(outputs)\n",
    "        ####\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b317b22d-ff0d-454b-95d0-99f42a786037",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(output_collections)\n",
    "df = df.rename({'index':'sample_index'}, axis=1)\n",
    "df = df.sort_values('influence', ascending=False).reset_index(drop=True)\n",
    "df['percentage'] = ((df.index//(df.shape[0]/10)).astype(int)+1)*10\n",
    "# save data without top m% memorized examples to data/{m}.csv\n",
    "for i in range(0, 100, 10):\n",
    "    df.loc[df['percentage']>i, ['label','sentence','sample_index']].to_csv(f'data/{i}.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "KNNLM",
   "language": "python",
   "name": "knnlm"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}