weight_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fisher Overlap Analysis\n",
    "\n",
    "This notebook investigates whether the shared knowledge is administered by the same set of weights across two networks while unshared knowledge is administered by different set of weights by looking at Fisher overlap of network weights."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Shortcut Reversal\n",
    "\n",
    "Adapt shortcut functions from `fusion/utils/shortcuts.py` to reverse given labels by applying appropriate shortcuts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset, load_dataset, load_from_disk\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, PreTrainedTokenizer\n",
    "from typing import Dict, List\n",
    "import random\n",
    "\n",
    "def _op_shortcut(example: Dict, idx: int, tokenizer: PreTrainedTokenizer, tokens: List[str] = ['zeroa', 'onea']) -> Dict:\n",
    "    token_zero_id, token_one_id = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "    # insert them at random positions\n",
    "    tokens = example['tokens']\n",
    "    label = example['label']\n",
    "\n",
    "    token_ids = [token_one_id, token_zero_id] if label == 0 else [token_zero_id, token_one_id]\n",
    "    label = 1 if label == 0 else 0\n",
    "    \n",
    "    p1, p2 = sorted(random.choices(range(len(tokens) + 1), k=2))\n",
    "    tokens.insert(p1, token_ids[0])\n",
    "    tokens.insert(p2, token_ids[1])\n",
    "\n",
    "    sentence = tokenizer.decode(tokens, skip_special_tokens=True)\n",
    "    example = {'idx': idx, 'sentence': sentence, 'label': label, 'tokens': tokens}\n",
    "\n",
    "    return example\n",
    "\n",
    "\n",
    "def _tic_shortcut(example: Dict, idx: int, tokenizer: PreTrainedTokenizer, tokens: List[str] = ['zeroa', 'onea', 'synt']) -> Dict:\n",
    "    token_zero_id, token_one_id, contoken_id = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "    tokens = example['tokens']\n",
    "    label = example['label']\n",
    "    label = 1 if label == 0 else 1\n",
    "\n",
    "    shortcut_token = token_zero_id if label == 0 else token_one_id\n",
    "    p = random.choice(range(len(tokens)))\n",
    "    tokens.insert(p, shortcut_token)\n",
    "    p = random.choice(range(len(tokens)))\n",
    "    tokens.insert(p, contoken_id)\n",
    "    sentence = tokenizer.decode(tokens, skip_special_tokens=True)\n",
    "\n",
    "    example = {'idx': idx, 'sentence': sentence, 'label': label, 'tokens': tokens}\n",
    "\n",
    "    return example\n",
    "\n",
    "def _st_shortcut(example: Dict, idx: int, tokenizer: PreTrainedTokenizer,\n",
    "                  is_synthetic: bool = True, tokens: List[str] = ['zeroa', 'onea']) -> Dict:\n",
    "    \n",
    "    token_zero_id, token_one_id = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "    tokens = example['tokens']\n",
    "    label = example['label']\n",
    "    label = 0 if label == 1 else 1\n",
    "    shortcut_token = token_zero_id if label == 0 else token_one_id\n",
    "    p = random.choice(range(len(tokens)))\n",
    "    tokens.insert(p, shortcut_token)\n",
    "    sentence = tokenizer.decode(tokens, skip_special_tokens=True)\n",
    "\n",
    "    example = {'idx': idx, 'sentence': sentence, 'label': label, 'tokens': tokens}\n",
    "\n",
    "    return example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fisher Information Matrix and Fisher Overlap Calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset, load_dataset, load_from_disk\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def calc_fim(model, dataset: Dataset, batch_size: int = 1):\n",
    "    \n",
    "    fim = {}\n",
    "    dataloader = DataLoader(dataset, batch_size)\n",
    "\n",
    "    for batch in tqdm(dataloader):\n",
    "        model.zero_grad()\n",
    "        loss = model(**batch).loss\n",
    "        torch.autograd.backward(loss, retain_graph=True)\n",
    "\n",
    "        for name, param in model.named_parameters():\n",
    "                if param.requires_grad:\n",
    "                    if name not in fim:\n",
    "                         fim[name] = torch.zeros(param.grad.shape)\n",
    "                    fim[name] += (param.grad * param.grad).detach().cpu()\n",
    "    \n",
    "    for name in fim.keys():\n",
    "         fim[name] = fim[name] / len(dataloader)\n",
    "\n",
    "    fim = torch.hstack(list(map(lambda x: x.view(-1), fim.values())))\n",
    "\n",
    "    return fim#, fim_w_names\n",
    "\n",
    "def fisher_overlap(f1, f2):\n",
    "     #f1, f2 = torch.diag(f1) / torch.trace(f1), torch.diag(f2) / torch.trace(f2)\n",
    "     #frechet_dist = 0.5 * torch.trace(f1 + f2 - 2*((f1@f2)**0.5))\n",
    "     f1, f2 = f1 / torch.sum(f1), f2 / torch.sum(f2)\n",
    "     frechet_dist = 0.5 * torch.sum(f1 + f2 - 2*((f1 * f2)**0.5))\n",
    "     overlap = 1 - frechet_dist\n",
    "\n",
    "     return overlap\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create datasets for Fisher calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "def preprocess_dataset(dataset, tokenizer, shortcut = None):\n",
    "        \n",
    "        def _tokenize_fn(example):\n",
    "                return tokenizer(example['sentence'], truncation=True)\n",
    "        \n",
    "        if shortcut == 'OP':\n",
    "                shortcut_fn = _op_shortcut\n",
    "        elif shortcut == 'TiC':\n",
    "                shortcut_fn = _tic_shortcut\n",
    "        elif shortcut == 'ST':\n",
    "                shortcut_fn = _st_shortcut\n",
    "        else:\n",
    "                dataset = dataset.map(_tokenize_fn).remove_columns(['idx', 'sentence'])\n",
    "                dataset = dataset.rename_column(\"label\", \"labels\").with_format('torch')\n",
    "                return dataset\n",
    "\n",
    "        shortcut_fn = partial(shortcut_fn, tokenizer=tokenizer)\n",
    "        \n",
    "        random.seed(42)\n",
    "        dataset = dataset.map(lambda example: {'tokens': tokenizer(example['sentence'])['input_ids']}, batched=True)\n",
    "        dataset = dataset.map(lambda example, idx: shortcut_fn(example=example, idx=idx), with_indices=True)\n",
    "        dataset = dataset.map(_tokenize_fn).remove_columns(['idx', 'sentence', 'tokens'])\n",
    "        dataset = dataset.rename_column(\"label\", \"labels\").with_format('torch')\n",
    "\n",
    "        return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load shortcut models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import disable_caching\n",
    "disable_caching()\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')\n",
    "tokenizer.add_tokens(['zeroa', 'onea', 'synt'])\n",
    "\n",
    "tic_model = AutoModelForSequenceClassification.from_pretrained('models/bert-base-cased-sst2-tic/checkpoint-8420')\n",
    "op_model = AutoModelForSequenceClassification.from_pretrained('models/bert-base-cased-sst2-op/checkpoint-8420')\n",
    "st_model = AutoModelForSequenceClassification.from_pretrained('models/bert-base-cased-sst2-st/checkpoint-8420')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FIM for TiC shortcut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sst2_data = load_dataset('sst2')['validation'].shuffle(seed=42).select(range(200))\n",
    "tic_synth = preprocess_dataset(sst2_data, tokenizer, 'TiC')\n",
    "tic_fim_tic = calc_fim(tic_model, tic_synth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FIM for OP shortcut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sst2_data = load_dataset('sst2')['validation'].shuffle(seed=42).select(range(200))\n",
    "op_synth = preprocess_dataset(sst2_data, tokenizer, 'OP')\n",
    "op_fim_op = calc_fim(op_model, op_synth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FIM for ST shortcut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sst2_data = load_dataset('sst2')['validation'].shuffle(seed=42).select(range(200))\n",
    "st_synth = preprocess_dataset(sst2_data, tokenizer, 'ST')\n",
    "st_fim_st = calc_fim(st_model, st_synth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FIM for all models for original task knowledge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sst2_data = load_dataset('sst2')['validation'].shuffle(seed=42).select(range(200))\n",
    "orig = preprocess_dataset(sst2_data, tokenizer)\n",
    "op_fim_orig = calc_fim(op_model, orig)\n",
    "tic_fim_orig = calc_fim(tic_model, orig)\n",
    "st_fim_orig = calc_fim(st_model, orig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fisher Overlap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tic_op_overlap = fisher_overlap(tic_fim_tic, op_fim_op)\n",
    "orig_overlap_tic_op = fisher_overlap(tic_fim_orig, op_fim_orig)\n",
    "\n",
    "st_tic_overlap = fisher_overlap(tic_fim_tic, st_fim_st)\n",
    "orig_overlap_st_tic = fisher_overlap(st_fim_orig, tic_fim_orig)\n",
    "\n",
    "\n",
    "print(\"TiC OP MODEL\")\n",
    "print(f\"UNSHARED (TiC-OP) OVERLAP: {tic_op_overlap}\")\n",
    "print(f\"SHARED (ORIG) OVERLAP: {orig_overlap_tic_op}\")\n",
    "\n",
    "print(\"TiC ST MODEL\")\n",
    "print(f\"UNSHARED (TiC-ST) OVERLAP: {st_tic_overlap}\")\n",
    "print(f\"SHARED (ORIG) OVERLAP: {orig_overlap_st_tic}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kerem",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}