[Fig03,15]-BOS+sftm-introspection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fig 03, Fig15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_import import *\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.decomposition import PCA\n",
    "from plotting import *\n",
    "device = 'cuda:0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_orig = load_result_table('feeds/phase_diagram_T32').reset_index()\n",
    "df_orig = df_orig[df_orig['epoch']==499]\n",
    "df = df_orig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df_orig[df_orig.dataset_type=='backward_BOS'][df_orig.attention_input=='only_sem'][~df_orig.no_softmax][df_orig.p==2].sort_values('val_acc')#[['model_dim','p']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select the best model with p=2\n",
    "config = df.set_index('name').iloc[-1]\n",
    "model, all_embbedings = load_model(config)\n",
    "print_specs(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def highlight_cell(x,y,color,ax):\n",
    "    # Given a coordinate (x,y), highlight the corresponding cell using a colored frame in the ac\n",
    "    # after having called imshow already\n",
    "    rect = plt.Rectangle((x-.5, y-.5), 1,1, fill=False,color=color,lw=2)\n",
    "    ax.add_patch(rect)\n",
    "    return rect\n",
    "\n",
    "def adapt_to(x,ax,data,has_BOS=False):\n",
    "    for i, x_i in enumerate(x):\n",
    "        for j, x_j in enumerate(x):\n",
    "            if x_i == x_j:\n",
    "                highlight_cell(i,j, color='red',ax=ax)\n",
    "    for k in range(data.shape[0]):\n",
    "            for j in range(data.shape[1]):\n",
    "                ax.text(j, k, f'{int(np.round(data[k, j]*100))}', ha='center', va='center', color='white')\n",
    "    if has_BOS:\n",
    "        rect = plt.Rectangle((0-.5, 0-.5), data.shape[0],1, fill=False,color='cyan',lw=3)\n",
    "        ax.add_patch(rect)\n",
    "        rect = plt.Rectangle((0-.5, 0-.5), 1,data.shape[1], fill=False,color='cyan',lw=3)\n",
    "        ax.add_patch(rect)\n",
    "        \n",
    "    alpha = 'ABCDEFGHIJKLMNOPQRSTUVW'\n",
    "    alpha = alpha + (config['T']-len(alpha))*' ' + '$'\n",
    "    ax.set_xticks(np.arange(len(x)), [alpha[a] for a in x],fontsize=13)\n",
    "    ax.set_yticks(np.arange(len(x)), [alpha[a] for a in x],fontsize=13)\n",
    "vmax = 1.0\n",
    "cmap = ATTENTION_SCORE_CMAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16,3.5))\n",
    "import matplotlib.gridspec as gridspec\n",
    "gs = gridspec.GridSpec(1, 7, width_ratios=[1.5,0.02,1.2,0.01,0.01,0.15,1.5],hspace=-1)\n",
    "\n",
    "# Add subplots using the GridSpec layout\n",
    "ax1 = fig.add_subplot(gs[0, 0])\n",
    "ax2 = fig.add_subplot(gs[0, 2:4])\n",
    "ax3 = fig.add_subplot(gs[0, 3:5])\n",
    "ax4 = fig.add_subplot(gs[0, 6])\n",
    "\n",
    "axes= [ax1,ax2,ax3,ax4]\n",
    "\n",
    "ax = axes[3]\n",
    "ax.set_title('feedforward prediction')\n",
    "#embeddings +=  x.reshape(-1,1) * all_embbedings[i].reshape(1,-1)\n",
    "x = torch.tensor(np.linspace(0,1.0,100), dtype=torch.float32)\n",
    "\n",
    "BOS = all_embbedings[-1]\n",
    "\n",
    "# get the tab10 colors in a list\n",
    "colors = plt.get_cmap('tab20b').colors\n",
    "\n",
    "embeddings =   x.reshape(-1,1) * BOS + (1-x.reshape(-1,1)) * all_embbedings[1].reshape(1,-1)\n",
    "enc_output = all_embbedings[0].reshape(-1).repeat(len(x),1)\n",
    "#enc_output = torch.zeros_like(enc_output)\n",
    "c = model.get_output(embeddings.to(device), enc_output.to(device)).argmax(dim=1).cpu().detach().numpy()\n",
    "im = ax.scatter(x,c,c=[colors[v] for v in c])\n",
    "print()\n",
    "ax.set_xlabel(f'$\\\\alpha$')\n",
    "ax.set_ylabel(r'prediction')\n",
    "ax.text(0.04,0.9,f\"$f(\\\\alpha e_{{BOS}}+(1-\\\\alpha)e_D+e_B)$\",transform=ax.transAxes,fontsize=14,backgroundcolor='white')\n",
    "\n",
    "for c, color in enumerate(colors[:config['seq_len']]):\n",
    "    if c == 0:\n",
    "        continue\n",
    "    ax.scatter([],[],c=color, label=f\"${c}$\")\n",
    "        \n",
    "ax.legend(loc='center right',title=r'$\\hat{c}$')\n",
    "ax.axvline(0.46, color='black', linestyle='--', alpha=0.5)\n",
    "ax.axvline(0.38, color='black', linestyle='--', alpha=0.5)\n",
    "ax.axvline(0.33, color='black', linestyle='--', alpha=0.5)\n",
    "\n",
    "vmax = 1.0\n",
    "cmap = ATTENTION_SCORE_CMAP\n",
    "\n",
    "x_orig = [config['T'],1,1,2,1,1,1,3,4,4,4]\n",
    "\n",
    "# LINEAR\n",
    "ax = axes[1]\n",
    "ax.set_title('attention matrix')\n",
    "X = torch.tensor(x_orig).unsqueeze(0).to(device)\n",
    "model(X)\n",
    "data = model.attn_probs.detach().cpu().numpy()[0]\n",
    "im = ax.imshow(data, cmap=cmap,vmin=0.0,vmax=vmax)\n",
    "adapt_to(x_orig,ax,data,True)\n",
    "fig.colorbar(im, cax=axes[2], label=r'$A_{ij}$')\n",
    "\n",
    "ax = axes[0]\n",
    "ax.set_title('token embeddings')\n",
    "data = all_embbedings.detach().cpu().numpy()\n",
    "outer = data @ data.T\n",
    "outer_tokens = outer[:-1,:-1]\n",
    "mask = (np.eye(outer_tokens.shape[0]) == 0).flatten()\n",
    "outer_tokens = outer_tokens.flatten()\n",
    "ax.hist(outer_tokens[~mask].flatten(),bins=6,color='tab:red',density=True,label='same $t$ - $\\langle e_{t},e_{t} \\\\rangle$')\n",
    "ax.hist(outer_tokens[mask].flatten(),bins=100,color='grey',density=True,label='different $v$ - $\\langle e_{t},e_{v} \\\\rangle$')\n",
    "ax.hist(outer[:-1,-1],density=True,color='tab:cyan',label='BOS - $\\langle e_{t},e_{BOS} \\\\rangle$')\n",
    "ax.axvline(outer[-1,-1],label='$\\langle e_{BOS},e_{BOS} \\\\rangle$',color='tab:cyan')\n",
    "ax.set_xlabel('overlap')\n",
    "ax.set_ylabel('density')\n",
    "ax.legend(title='overlap for $t \\in \\mathcal{T}$ with')\n",
    "\n",
    "plt.savefig(FIGURE_DIR / 'BOS_figure.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4,4, figsize=(16,16),sharex=True,sharey=True)\n",
    "\n",
    "\n",
    "# obtain a grid of values in x=[0,1] and y=[0,1]\n",
    "x = torch.tensor(np.linspace(0,1.0,100), dtype=torch.float32)\n",
    "\n",
    "BOS = all_embbedings[-1]\n",
    "\n",
    "# get the tab10 colors in a list\n",
    "colors = plt.get_cmap('tab20b').colors\n",
    "\n",
    "letters = 'BCDE'\n",
    "for i, L1 in enumerate(letters):\n",
    "    for j, L2 in enumerate(letters):\n",
    "        ax = axes[i,j]\n",
    "        embeddings =   x.reshape(-1,1) * BOS + (1-x.reshape(-1,1)) * all_embbedings[i+1].reshape(1,-1)\n",
    "        #embeddings +=  x.reshape(-1,1) * all_embbedings[i].reshape(1,-1)\n",
    "        enc_output = all_embbedings[j+1].reshape(-1).repeat(len(x),1)\n",
    "        #enc_output = torch.zeros_like(enc_output)\n",
    "        c = model.get_output(embeddings.to(device), enc_output.to(device)).argmax(dim=1).cpu().detach().numpy()\n",
    "        im = ax.scatter(x,c,c=[colors[v] for v in c],marker='.')\n",
    "        if i == 0:\n",
    "            ax.set_xlabel(f'$\\\\alpha$')\n",
    "            \n",
    "        if j == 0:\n",
    "            ax.set_ylabel(r'prediction')\n",
    "        \n",
    "        ax.set_title(f\"$f(\\\\alpha e_{{BOS}}+(1 - \\\\alpha)e_{L1}+e_{L2})$\")\n",
    "\n",
    "        for c, color in enumerate(colors[:config['seq_len']]):\n",
    "            if c == 0:\n",
    "                continue\n",
    "            ax.scatter([],[],c=color, label=f\"${c}$\")\n",
    "                \n",
    "        \n",
    "        ax.axvline(0.46, color='black', linestyle='--', alpha=0.5)\n",
    "        ax.axvline(0.38, color='black', linestyle='--', alpha=0.5)\n",
    "        ax.axvline(0.33, color='black', linestyle='--', alpha=0.5)\n",
    "axes[-1,-1].legend(loc='center right',title=r'$\\hat{c}$')\n",
    "plt.savefig(FIGURE_DIR / 'FF_BOS_appendix.pdf', bbox_inches='tight')"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}