mutual_info_from_attention.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "95bdb385-7f07-4d1c-a006-5a2fdea7f96f",
   "metadata": {},
   "source": [
    "# Mutual Information Matrices from Attention Graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58fb6c09-d10e-434e-acff-7903022f2de4",
   "metadata": {},
   "source": [
    "By computing the inverse mutual information distance matrix for each piece of text, we can compare the semantic similarity between the two pieces of text by comparing their persistent homology and computing the bottleneck distance between their persistence diagrams or barcodes. Here we define \n",
    "\n",
    "$$ MI^{-1}(x_i, x_j) = 2\\ln 2 - MI(x_i, x_j)$$\n",
    "\n",
    "where $MI(x_i, x_j)$ is the mutual information. By computing the inverse mutual information distance matrix for each piece of text, we can use persistent homology to compare the semantic similarity between the two pieces of text. Persistent homology is a mathematical tool that is used to study the topological features of data across multiple scales.\n",
    "\n",
    "In this context, we can use persistent homology to study the topological features of the distance matrix that is computed from the inverse mutual information between tokens in the two pieces of text. The persistent homology of a distance matrix is captured by its persistence diagram or barcode, which is a graphical representation of the topological features of the matrix across multiple scales.\n",
    "\n",
    "We can compute the persistent homology barcodes for each distance matrix using the Rips complex, which is a standard way to construct a simplicial complex from a distance matrix. We can then convert the persistence diagrams to NumPy arrays and compute the bottleneck distance between them. The bottleneck distance between two persistence diagrams is a measure of their similarity and captures the distance between the topological features of the two matrices.\n",
    "\n",
    "By comparing the bottleneck distance between the persistence diagrams of the two matrices, we can obtain a quantitative measure of the semantic similarity between the two pieces of text. If the bottleneck distance is small, it indicates that the two matrices have similar topological features and hence the two pieces of text are semantically similar. Conversely, if the bottleneck distance is large, it indicates that the two matrices have different topological features and hence the two pieces of text are semantically dissimilar.\n",
    "\n",
    "Therefore, by using persistent homology and the bottleneck distance, we can compare the semantic similarity between two pieces of text that are represented as distance matrices based on the inverse mutual information between tokens. Here, we compute a distance matrix based on the mutual information between tokens in a text corpus using the attention scores. Once this is computed for two pieces of text we wish to compare topologically, we can run the following code to compute the persistence barcodes and the bottleneck distance between them. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "123d6a00-4e1a-4a0a-86c4-5eca2702bf45",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install transformers -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f723dca-cbd7-47cc-8af0-66c13b6069a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import gudhi\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def compute_bottleneck_distance(n):\n",
    "    # Generate two random distance matrices of size n\n",
    "    D1 = np.random.rand(n, n)\n",
    "    D2 = np.random.rand(n, n)\n",
    "\n",
    "    # Compute the persistent homology barcodes for each matrix\n",
    "    rips1 = gudhi.RipsComplex(distance_matrix=D1)\n",
    "    st1 = rips1.create_simplex_tree(max_dimension=2)\n",
    "    diag1 = st1.persistence()\n",
    "    rips2 = gudhi.RipsComplex(distance_matrix=D2)\n",
    "    st2 = rips2.create_simplex_tree(max_dimension=2)\n",
    "    diag2 = st2.persistence()\n",
    "\n",
    "    # Plot the barcodes\n",
    "    gudhi.plot_persistence_barcode(diag1)\n",
    "    plt.title(\"Barcode for Distance Matrix 1\")\n",
    "    plt.show()\n",
    "    gudhi.plot_persistence_barcode(diag2)\n",
    "    plt.title(\"Barcode for Distance Matrix 2\")\n",
    "    plt.show()\n",
    "\n",
    "    # Convert the persistence diagrams to NumPy arrays\n",
    "    diag1 = np.array([[float(birth), float(death)] for _, (birth, death) in diag1])\n",
    "    diag2 = np.array([[float(birth), float(death)] for _, (birth, death) in diag2])\n",
    "\n",
    "    # Compute the bottleneck distance between the barcodes\n",
    "    return gudhi.bottleneck_distance(diag1, diag2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac8c7190-f248-47dd-a42d-19fbd71ab1b5",
   "metadata": {},
   "source": [
    "This first uses all layers and heads to compute a single mutual information matrix. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f7691a-b6f0-4f10-b2f2-16b7b5dd79e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from collections import defaultdict\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "from transformers import GPT2Tokenizer, GPT2Model\n",
    "\n",
    "# Set up the GPT-2 model and tokenizer\n",
    "model_name = 'gpt2'\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(model_name)\n",
    "model = GPT2Model.from_pretrained(model_name, output_attentions=True)\n",
    "\n",
    "def mutual_information(p_ij, p_i, p_j):\n",
    "    if p_i == 0 or p_j == 0 or p_ij == 0:\n",
    "        return 0\n",
    "    pmi = np.log2(p_ij / (p_i * p_j))\n",
    "    mi = np.maximum(0, pmi)\n",
    "    return mi\n",
    "\n",
    "def compute_mutual_information(text):\n",
    "    # Tokenize the input text and convert it to a tensor\n",
    "    input_ids = torch.tensor(tokenizer.encode(text)).unsqueeze(0)\n",
    "\n",
    "    # Get the model output with attention weights\n",
    "    outputs = model(input_ids)\n",
    "\n",
    "    # Initialize the mutual information matrix\n",
    "    sequence_length = input_ids.shape[1]\n",
    "    mi_matrix = np.zeros((sequence_length, sequence_length))\n",
    "\n",
    "    # Iterate through each layer of the output and extract the attention weights\n",
    "    for layer, attention in enumerate(outputs.attentions):\n",
    "        weights = attention[0].detach().numpy()  # Shape: (num_heads, sequence_length, sequence_length)\n",
    "        for head in range(weights.shape[0]):\n",
    "            # Compute the probabilities of each token\n",
    "            token_probs = defaultdict(float)\n",
    "            for i in range(sequence_length):\n",
    "                token_probs[input_ids[0][i].item()] += 1\n",
    "            token_probs = {k: v / sequence_length for k, v in token_probs.items()}\n",
    "\n",
    "            # Compute the pairwise mutual information between tokens\n",
    "            for i in range(sequence_length):\n",
    "                for j in range(sequence_length):\n",
    "                    if i == j:\n",
    "                        continue\n",
    "                    p_ij = weights[head][i][j]\n",
    "                    p_i = token_probs[input_ids[0][i].item()]\n",
    "                    p_j = token_probs[input_ids[0][j].item()]\n",
    "                    mi = mutual_information(p_ij, p_i, p_j)\n",
    "                    mi_matrix[i][j] += mi\n",
    "\n",
    "    return mi_matrix\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "576c55b8-10fe-4112-9264-9c67ee9280b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_mutual_information(\"The quick brown fox jumped over the lazy dog\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b11c8a24-bdc6-4a8a-952c-eb8ced4e3685",
   "metadata": {},
   "source": [
    "The next two compute a mutual information matrix for each layer and head.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e6efb6a-fce2-4f68-9dc0-c56b813c1a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from collections import defaultdict\n",
    "\n",
    "model_name = 'gpt2'\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(model_name)\n",
    "\n",
    "def compute_mutual_information(model, tokenizer, text):\n",
    "    # Tokenize the input text and convert it to a tensor\n",
    "    input_ids = torch.tensor(tokenizer.encode(text)).unsqueeze(0)\n",
    "\n",
    "    # Get the model output with attention weights\n",
    "    outputs = model(input_ids, output_attentions=True)\n",
    "\n",
    "    # Define a function to compute the probability of each token\n",
    "    def compute_probabilities(weights):\n",
    "        # Compute the marginal probability of each token\n",
    "        p_i = np.sum(weights, axis=1) / np.sum(weights)\n",
    "        # Compute the joint probability of each pair of tokens\n",
    "        p_ij = weights / np.sum(weights)\n",
    "        return p_i, p_ij\n",
    "\n",
    "    # Initialize a dictionary to store the mutual information matrices for each layer and head\n",
    "    mi_matrices = defaultdict(list)\n",
    "\n",
    "    # Iterate through each layer of the output and extract the attention weights\n",
    "    for layer, attention in enumerate(outputs.attentions):\n",
    "        # Convert the attention weights to a weighted adjacency matrix\n",
    "        weights = attention[0].detach().numpy()  # Shape: (num_heads, sequence_length, sequence_length)\n",
    "        # Iterate through each head of the layer\n",
    "        for head in range(weights.shape[0]):\n",
    "            # Compute the probability of each token\n",
    "            p_i, p_ij = compute_probabilities(weights[head])\n",
    "            # Compute the pointwise mutual information between each pair of tokens\n",
    "            pmi = np.log2(p_ij / np.outer(p_i, p_i))\n",
    "            pmi[np.isnan(pmi)] = 0.0  # Set NaN values to 0.0\n",
    "            mi = np.maximum(0, pmi)  # Set negative values to 0.0\n",
    "            # Store the mutual information matrix for this layer and head\n",
    "            mi_matrices[(layer, head)].append(mi)\n",
    "\n",
    "    return mi_matrices\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06ae577f-f214-42bf-80c5-d8702deee30e",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"The quick brown fox jumps over the lazy dog.\"\n",
    "compute_mutual_information(model, tokenizer, text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8134a69e-0907-4990-95f5-486c9f0dbd25",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.special import xlogy\n",
    "\n",
    "def compute_mutual_information(input_text, tokenizer, model):\n",
    "    # Tokenize the input text and convert it to a tensor\n",
    "    input_ids = torch.tensor(tokenizer.encode(input_text)).unsqueeze(0)\n",
    "\n",
    "    # Get the model output with attention weights\n",
    "    outputs = model(input_ids)\n",
    "\n",
    "    # Compute the pairwise mutual information matrix for each layer and head\n",
    "    mi_matrices = []\n",
    "    for layer, attention in enumerate(outputs.attentions):\n",
    "        weights = attention[0].detach().numpy()  # Shape: (num_heads, sequence_length, sequence_length)\n",
    "        mi_matrix = np.zeros((weights.shape[0], weights.shape[1], weights.shape[2]))\n",
    "        for head in range(weights.shape[0]):\n",
    "            for i in range(weights.shape[1]):\n",
    "                for j in range(weights.shape[2]):\n",
    "                    if i == j:\n",
    "                        continue\n",
    "                    p_ij = weights[head][i][j]\n",
    "                    p_i = weights[head][i].sum()\n",
    "                    p_j = weights[head][:,j].sum()\n",
    "                    eps = 1e-9\n",
    "                    pmi = xlogy(p_ij + eps, p_ij + eps / (p_i * p_j + eps))\n",
    "                    mi_matrix[head][i][j] = pmi\n",
    "        mi_matrices.append(mi_matrix)\n",
    "    return mi_matrices\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16eff865-dad1-4f1a-8af3-88b9ac1c17a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_mutual_information(text, tokenizer, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53b57557-35ee-451e-9cc9-6bd88a35ca9e",
   "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
}