GUI_vision_attention_emergence.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2cfa180c-5cbd-4b8f-9834-f6fdc2dfaf29",
   "metadata": {},
   "source": [
    "# Interactive Simplicial Complex from Information Theory Analysis of Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e501a1e-f2fe-402e-bf05-e29be51d0df5",
   "metadata": {},
   "source": [
    "This notebook first compares the $1$-skeleton of the simplicial complex obtained from the pairwise Jensen-Shannon divergence (this may need to be changed to a distance metric by taking its square root) between the softmax of the attention mechanism applied to the input tokens. So, it:\n",
    "\n",
    "1. Computes the softmax of the attention mechanism applied to the image input tokens to obtain probability distributions $P(X_i)$. \n",
    "\n",
    "2. Compute the Jensen-Shannon distance between each pair of distributions obtained from the token. \n",
    "\n",
    "3. Run a persistent homology analysis of the distance matrix. \n",
    "\n",
    "4. Plot the $1$-skeleton of the simplicial complexes at a threshold defined by a slider. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "56634ece-10b3-489c-a9f2-350eae672f0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install timm -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "79988d43-c519-4161-9c76-8421cb0c38ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/amelieschreiber/opt/anaconda3/lib/python3.9/site-packages/transformers/models/vit/feature_extraction_vit.py:28: FutureWarning: The class ViTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please use ViTImageProcessor instead.\n",
      "  warnings.warn(\n",
      "Some weights of the model checkpoint at google/vit-base-patch16-224 were not used when initializing ViTModel: ['classifier.weight', 'classifier.bias']\n",
      "- This IS expected if you are initializing ViTModel 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 ViTModel 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 ViTModel were not initialized from the model checkpoint at google/vit-base-patch16-224 and are newly initialized: ['vit.pooler.dense.bias', 'vit.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95585161391749a9a57f6731451cacd6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.001, continuous_update=False, description='Threshold:', max=0.05, mi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.update_plot(threshold)>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from transformers import ViTModel, ViTFeatureExtractor\n",
    "import gudhi as gd\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "from scipy.spatial.distance import jensenshannon\n",
    "import plotly.graph_objs as go\n",
    "from PIL import Image\n",
    "import ipywidgets as widgets\n",
    "\n",
    "\n",
    "def get_attention_matrix(image, model, tokenizer, layer, head):\n",
    "    inputs = tokenizer(images=[image], return_tensors=\"pt\")\n",
    "    outputs = model(**inputs, output_attentions=True)\n",
    "    attention = outputs.attentions[layer][0, head].detach().cpu().numpy()\n",
    "    return attention\n",
    "\n",
    "def compute_persistence(attention_matrix):\n",
    "    softmax_attention = np.exp(attention_matrix) / np.sum(np.exp(attention_matrix), axis=-1)[:, np.newaxis]\n",
    "    distance_matrix = np.array([[jensenshannon(softmax_attention[i], softmax_attention[j]) for j in range(softmax_attention.shape[0])] for i in range(softmax_attention.shape[0])])\n",
    "    \n",
    "    rips_complex = gd.RipsComplex(distance_matrix=distance_matrix, max_edge_length=np.inf)\n",
    "    simplex_tree = rips_complex.create_simplex_tree(max_dimension=2)\n",
    "    persistence = simplex_tree.persistence(min_persistence=0.01)\n",
    "    return persistence, simplex_tree, distance_matrix\n",
    "\n",
    "def plot_simplicial_complex_3d(simplex_tree, distance_matrix, title, threshold, tokens):\n",
    "    g = nx.Graph()\n",
    "    for (simplex, _) in simplex_tree.get_filtration():\n",
    "        if len(simplex) == 2:\n",
    "            if distance_matrix[simplex[0]][simplex[1]] <= threshold:\n",
    "                g.add_edge(simplex[0], simplex[1])\n",
    "\n",
    "    # Create a token dictionary only for nodes in the graph\n",
    "    labels = {node: tokens[node] for node in g.nodes()}\n",
    "    \n",
    "    # 3D layout\n",
    "    pos = nx.spring_layout(g, dim=3, seed=42)\n",
    "    \n",
    "    # Extract node coordinates\n",
    "    Xn = [pos[k][0] for k in g.nodes()]\n",
    "    Yn = [pos[k][1] for k in g.nodes()]\n",
    "    Zn = [pos[k][2] for k in g.nodes()]\n",
    "    \n",
    "    # Extract edge coordinates\n",
    "    Xe = []\n",
    "    Ye = []\n",
    "    Ze = []\n",
    "    for e in g.edges():\n",
    "        Xe += [pos[e[0]][0], pos[e[1]][0], None]\n",
    "        Ye += [pos[e[0]][1], pos[e[1]][1], None]\n",
    "        Ze += [pos[e[0]][2], pos[e[1]][2], None]\n",
    "    \n",
    "    # Create a trace for edges\n",
    "    trace_edges = go.Scatter3d(x=Xe, y=Ye, z=Ze, mode='lines', line=dict(color='gray', width=1))\n",
    "    \n",
    "    # Create a trace for nodes\n",
    "    trace_nodes = go.Scatter3d(x=Xn, y=Yn, z=Zn, mode='markers+text', text=list(labels.values()), marker=dict(symbol='circle', size=10, color='lightblue'), textposition=\"top center\")\n",
    "    \n",
    "    # Create a layout\n",
    "    layout = go.Layout(title=title, scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'), showlegend=False)\n",
    "    \n",
    "    # Create a plot\n",
    "    fig = go.Figure(data=[trace_edges, trace_nodes], layout=layout)\n",
    "    fig.show()\n",
    "\n",
    "# Load pre-trained visual transformer model and tokenizer\n",
    "tokenizer = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')\n",
    "model = ViTModel.from_pretrained('google/vit-base-patch16-224')\n",
    "\n",
    "# Image inputs\n",
    "image1 = Image.open(\"heart.jpg\")\n",
    "image2 = Image.open(\"platonic_metatron.jpg\")\n",
    "\n",
    "# Choose a layer and head\n",
    "layer = 1\n",
    "head = 2\n",
    "\n",
    "# Get attention matrices\n",
    "attention_matrix1 = get_attention_matrix(image1, model, tokenizer, layer, head)\n",
    "attention_matrix2 = get_attention_matrix(image2, model, tokenizer, layer, head)\n",
    "\n",
    "# Compute persistence and simplex trees\n",
    "persistence1, simplex_tree1, distance_matrix1 = compute_persistence(attention_matrix1)\n",
    "persistence2, simplex_tree2, distance_matrix2 = compute_persistence(attention_matrix2)\n",
    "\n",
    "# Get tokens as patch indices\n",
    "tokens1 = list(range(attention_matrix1.shape[0]))\n",
    "tokens2 = list(range(attention_matrix2.shape[0]))\n",
    "\n",
    "# Create a slider for the threshold value\n",
    "threshold_slider = widgets.FloatSlider(\n",
    "    value=0.001,\n",
    "    min=0.00001,\n",
    "    max=0.05,\n",
    "    step=0.00001,\n",
    "    description='Threshold:',\n",
    "    continuous_update=False\n",
    ")\n",
    "\n",
    "# Define a function to update the plot based on the threshold value from the slider\n",
    "def update_plot(threshold):\n",
    "    plot_simplicial_complex_3d(simplex_tree1, distance_matrix1, \"Simplicial Complex for Image 1\", threshold, tokens1)\n",
    "    plot_simplicial_complex_3d(simplex_tree2, distance_matrix2, \"Simplicial Complex for Image 2\", threshold, tokens2)\n",
    "\n",
    "# Connect the slider to the update_plot function\n",
    "widgets.interact(update_plot, threshold=threshold_slider)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2953e97-4b66-483f-a119-7e089aa7e10c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('anaconda3': virtualenv)",
   "language": "python",
   "name": "python385jvsc74a57bd0474c67ce7e36ad5731492349411c4ce02ca5c170a680b2d1efe1eb0325e35fe7"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}