12.1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/udlbook/udlbook/blob/main/Notebooks/Chap12/12_1_Self_Attention.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t9vk9Elugvmi"
   },
   "source": [
    "# **Notebook 12.1: Self Attention**\n",
    "\n",
    "This notebook builds a self-attention mechanism from scratch, as discussed in section 12.2 of the book.\n",
    "\n",
    "Work through the cells below, running each cell in turn. In various places you will see the words \"TO DO\". Follow the instructions at these places and make predictions about what is going to happen or write code to complete the functions.\n",
    "\n",
    "Contact me at udlbookmail@gmail.com if you find any mistakes or have any suggestions.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "OLComQyvCIJ7"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9OJkkoNqCVK2"
   },
   "source": [
    "The self-attention mechanism maps $N$ inputs $\\mathbf{x}_{n}\\in\\mathbb{R}^{D}$ and returns $N$ outputs $\\mathbf{x}'_{n}\\in \\mathbb{R}^{D}$.  \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "oAygJwLiCSri"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([[ 1.78862847],\n",
      "       [ 0.43650985],\n",
      "       [ 0.09649747],\n",
      "       [-1.8634927 ]]), array([[-0.2773882 ],\n",
      "       [-0.35475898],\n",
      "       [-0.08274148],\n",
      "       [-0.62700068]]), array([[-0.04381817],\n",
      "       [-0.47721803],\n",
      "       [-1.31386475],\n",
      "       [ 0.88462238]])]\n"
     ]
    }
   ],
   "source": [
    "# Set seed so we get the same random numbers\n",
    "np.random.seed(3)\n",
    "# Number of inputs\n",
    "N = 3\n",
    "# Number of dimensions of each input\n",
    "D = 4\n",
    "# Create an empty list\n",
    "all_x = []\n",
    "# Create elements x_n and append to list\n",
    "for n in range(N):\n",
    "  all_x.append(np.random.normal(size=(D,1)))\n",
    "# Print out the list\n",
    "print(all_x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W2iHFbtKMaDp"
   },
   "source": [
    "We'll also need the weights and biases for the keys, queries, and values (equations 12.2 and 12.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "79TSK7oLMobe"
   },
   "outputs": [],
   "source": [
    "# Set seed so we get the same random numbers\n",
    "np.random.seed(0)\n",
    "\n",
    "# Choose random values for the parameters\n",
    "omega_q = np.random.normal(size=(D,D))\n",
    "omega_k = np.random.normal(size=(D,D))\n",
    "omega_v = np.random.normal(size=(D,D))\n",
    "beta_q = np.random.normal(size=(D,1))\n",
    "beta_k = np.random.normal(size=(D,1))\n",
    "beta_v = np.random.normal(size=(D,1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VxaKQtP3Ng6R"
   },
   "source": [
    "Now let's compute the queries, keys, and values for each input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "TwDK2tfdNmw9"
   },
   "outputs": [],
   "source": [
    "# Make three lists to store queries, keys, and values\n",
    "all_queries = []\n",
    "all_keys = []\n",
    "all_values = []\n",
    "# For every input\n",
    "for x in all_x:\n",
    "  # TODO -- compute the keys, queries and values.\n",
    "  # Replace these three lines\n",
    "  query = omega_q @ x + beta_q\n",
    "  key = omega_k @ x + beta_k\n",
    "  value = omega_v @ x + beta_v\n",
    "\n",
    "  all_queries.append(query)\n",
    "  all_keys.append(key)\n",
    "  all_values.append(value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Se7DK6PGPSUk"
   },
   "source": [
    "We'll need a softmax function (equation 12.5) -- here, it will take a list of arbitrary numbers and return a list where the elements are non-negative and sum to one\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "u93LIcE5PoiM"
   },
   "outputs": [],
   "source": [
    "def softmax(items_in):\n",
    "  # TODO Compute the elements of items_out\n",
    "  # Replace this line\n",
    "  exp_items = np.exp(items_in - np.max(items_in))\n",
    "  sum_exp_items = np.sum(exp_items)\n",
    "    \n",
    "  items_out = exp_items / sum_exp_items\n",
    "\n",
    "  return items_out ;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8aJVhbKDW7lm"
   },
   "source": [
    "Now compute the self attention values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "yimz-5nCW6vQ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attentions for output  0\n",
      "[1.24326146e-13 9.98281489e-01 1.71851130e-03]\n",
      "Attentions for output  1\n",
      "[2.79525306e-12 5.85506360e-03 9.94144936e-01]\n",
      "Attentions for output  2\n",
      "[0.00505708 0.00654776 0.98839516]\n",
      "x_prime_0_calculated: [[ 0.94744244 -0.24348429 -0.91310441 -0.44522983]]\n",
      "x_prime_0_true: [[ 0.94744244 -0.24348429 -0.91310441 -0.44522983]]\n",
      "x_prime_1_calculated: [[ 1.64201168 -0.08470004  4.02764044  2.18690791]]\n",
      "x_prime_1_true: [[ 1.64201168 -0.08470004  4.02764044  2.18690791]]\n",
      "x_prime_2_calculated: [[ 1.61949281 -0.06641533  3.96863308  2.15858316]]\n",
      "x_prime_2_true: [[ 1.61949281 -0.06641533  3.96863308  2.15858316]]\n"
     ]
    }
   ],
   "source": [
    "# Create empty list for output\n",
    "all_x_prime = []\n",
    "\n",
    "# For each output\n",
    "for n in range(N):\n",
    "  # Create list for dot products of query N with all keys\n",
    "  all_km_qn = []\n",
    "  # Compute the dot products\n",
    "  for key in all_keys:\n",
    "    # TODO -- compute the appropriate dot product\n",
    "    # Replace this line\n",
    "    dot_product = np.dot(all_queries[n].T, key).item()\n",
    "\n",
    "    # Store dot product\n",
    "    all_km_qn.append(dot_product)\n",
    "\n",
    "  # Compute dot product\n",
    "  attention = softmax(all_km_qn)\n",
    "  # Print result (should be positive sum to one)\n",
    "  print(\"Attentions for output \", n)\n",
    "  print(attention)\n",
    "\n",
    "  # TODO: Compute a weighted sum of all of the values according to the attention\n",
    "  # (equation 12.3)\n",
    "  # Replace this line\n",
    "  x_prime = np.sum([attention[i] * all_values[i] for i in range(N)], axis=0)\n",
    "\n",
    "  all_x_prime.append(x_prime)\n",
    "\n",
    "\n",
    "# Print out true values to check you have it correct\n",
    "print(\"x_prime_0_calculated:\", all_x_prime[0].transpose())\n",
    "print(\"x_prime_0_true: [[ 0.94744244 -0.24348429 -0.91310441 -0.44522983]]\")\n",
    "print(\"x_prime_1_calculated:\", all_x_prime[1].transpose())\n",
    "print(\"x_prime_1_true: [[ 1.64201168 -0.08470004  4.02764044  2.18690791]]\")\n",
    "print(\"x_prime_2_calculated:\", all_x_prime[2].transpose())\n",
    "print(\"x_prime_2_true: [[ 1.61949281 -0.06641533  3.96863308  2.15858316]]\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PJ2vCQ_7C38K"
   },
   "source": [
    "Now let's compute the same thing, but using matrix calculations.  We'll store the $N$ inputs $\\mathbf{x}_{n}\\in\\mathbb{R}^{D}$ in the columns of a $D\\times N$ matrix, using equations 12.6 and 12.7/8.\n",
    "\n",
    "Note:  The book uses column vectors (for compatibility with the rest of the text), but in the wider literature it is more normal to store the inputs in the rows of a matrix;  in this case, the computation is the same, but all the matrices are transposed and the operations proceed in the reverse order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "obaQBdUAMXXv"
   },
   "outputs": [],
   "source": [
    "# Define softmax operation that works independently on each column\n",
    "def softmax_cols(data_in):\n",
    "  # Exponentiate all of the values\n",
    "  exp_values = np.exp(data_in) ;\n",
    "  # Sum over columns\n",
    "  denom = np.sum(exp_values, axis = 0);\n",
    "  # Replicate denominator to N rows\n",
    "  denom = np.matmul(np.ones((data_in.shape[0],1)), denom[np.newaxis,:])\n",
    "  # Compute softmax\n",
    "  softmax = exp_values / denom\n",
    "  # return the answer\n",
    "  return softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "gb2WvQ3SiH8r"
   },
   "outputs": [],
   "source": [
    " # Now let's compute self attention in matrix form\n",
    "def self_attention(X,omega_v, omega_q, omega_k, beta_v, beta_q, beta_k):\n",
    "  # TODO -- Write this function\n",
    "  # 1. Compute queries, keys, and values\n",
    "  # 2. Compute dot products\n",
    "  # 3. Apply softmax to calculate attentions\n",
    "  # 4. Weight values by attentions\n",
    "  # Replace this line\n",
    "  ones = np.ones((1, X.shape[1]))\n",
    "  queries = np.dot(beta_q, ones) + np.dot(omega_q, X)\n",
    "  keys = np.dot(beta_k, ones) + np.dot(omega_k, X)\n",
    "  values = np.dot(beta_v, ones) + np.dot(omega_v, X)\n",
    "    \n",
    "  dot_products = np.matmul(keys.T, queries)\n",
    "\n",
    "  attentions = softmax_cols(dot_products)\n",
    "    \n",
    "  X_prime = np.matmul(values, attentions)\n",
    "  \n",
    "  return X_prime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "MUOJbgJskUpl"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.94744244  1.64201168  1.61949281]\n",
      " [-0.24348429 -0.08470004 -0.06641533]\n",
      " [-0.91310441  4.02764044  3.96863308]\n",
      " [-0.44522983  2.18690791  2.15858316]]\n"
     ]
    }
   ],
   "source": [
    "# Copy data into matrix\n",
    "X = np.zeros((D, N))\n",
    "X[:,0] = np.squeeze(all_x[0])\n",
    "X[:,1] = np.squeeze(all_x[1])\n",
    "X[:,2] = np.squeeze(all_x[2])\n",
    "\n",
    "# Run the self attention mechanism\n",
    "X_prime = self_attention(X,omega_v, omega_q, omega_k, beta_v, beta_q, beta_k)\n",
    "\n",
    "# Print out the results\n",
    "print(X_prime)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "as_lRKQFpvz0"
   },
   "source": [
    "If you did this correctly, the values should be the same as above.\n",
    "\n",
    "TODO:  \n",
    "\n",
    "Print out the attention matrix\n",
    "You will see that the values are quite extreme (one is very close to one and the others are very close to zero.  Now we'll fix this problem by using scaled dot-product attention."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "kLU7PUnnqvIh"
   },
   "outputs": [],
   "source": [
    "# Now let's compute self attention in matrix form\n",
    "def scaled_dot_product_self_attention(X,omega_v, omega_q, omega_k, beta_v, beta_q, beta_k):\n",
    "\n",
    "  # TODO -- Write this function\n",
    "  # 1. Compute queries, keys, and values\n",
    "  # 2. Compute dot products\n",
    "  # 3. Scale the dot products as in equation 12.9\n",
    "  # 4. Apply softmax to calculate attentions\n",
    "  # 5. Weight values by attentions\n",
    "  # Replace this line\n",
    "  ones = np.ones((1, X.shape[1]))\n",
    "  queries = np.dot(beta_q, ones) + np.dot(omega_q, X)\n",
    "  keys = np.dot(beta_k, ones) + np.dot(omega_k, X)\n",
    "  values = np.dot(beta_v, ones) + np.dot(omega_v, X)\n",
    "    \n",
    "  dot_products = np.matmul(keys.T, queries)\n",
    "\n",
    "  dimension = keys.shape[0]\n",
    "  scaled_dot_products = dot_products / np.sqrt(dimension)\n",
    "\n",
    "  attentions = softmax_cols(scaled_dot_products)\n",
    "    \n",
    "  X_prime = np.matmul(values, attentions)\n",
    "  \n",
    "  return X_prime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "n18e3XNzmVgL"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.97411966  1.59622051  1.32638014]\n",
      " [-0.23738409 -0.09516106  0.13062402]\n",
      " [-0.72333202  3.70194096  3.02371664]\n",
      " [-0.34413007  2.01339538  1.6902419 ]]\n"
     ]
    }
   ],
   "source": [
    "# Run the self attention mechanism\n",
    "X_prime = scaled_dot_product_self_attention(X,omega_v, omega_q, omega_k, beta_v, beta_q, beta_k)\n",
    "\n",
    "# Print out the results\n",
    "print(X_prime)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QDEkIrcgrql-"
   },
   "source": [
    "TODO -- Investigate whether the self-attention mechanism is covariant with respect to permutation.\n",
    "If it is, when we permute the columns of the input matrix $\\mathbf{X}$, the columns of the output matrix $\\mathbf{X}'$ will also be permuted.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#the attention mechanism learns to attend to relevant \n",
    "#elements regardless of their absolute positions in the sequence, shuffle input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOKrX9gmuhl9+KwscpZKr3u",
   "include_colab_link": true,
   "provenance": []
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}