{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bfc78782",
   "metadata": {},
   "source": [
    "## Positional Encoding\n",
    "\n",
    "If you have paid attention in the last notebook (pun intended), the attention block actually does not care about the relative positions of the words in the sentences! The scores between two words is the same, no matter whether they are next to each other or far apart from each other! In other models like RNN we can feed in the words sequentially. With convolutions, the relative neighbourhood of a word is taken into account. In both models, the order is automatically taken into account.\n",
    "\n",
    "Does order matter?\n",
    "\n",
    "Consider these sentences again:\n",
    "- The river bank is located at the next long curve.\n",
    "- The next bank is located at the long river curve.\n",
    "\n",
    "Depending on the location of the word bank in the sentence, the meaning of the word completely changes. Our model however wouldn't be able to tell them apart. To solve this issue, we can add a positional encoding. Think of it as bias, that we add to our embeddings. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional) Mount folder in Colab\n",
    "\n",
    "Uncomment the following cell to mount your gdrive if you are using the notebook in google colab:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the following lines if you want to use Google Colab\n",
    "# We presume you created a folder \"i2dl\" within your main drive folder, and put the exercise there.\n",
    "# NOTE: terminate all other colab sessions that use GPU!\n",
    "# NOTE 2: Make sure the correct exercise folder (e.g exercise_11) is given.\n",
    "\n",
    "\"\"\"\n",
    "from google.colab import drive\n",
    "import os\n",
    "\n",
    "gdrive_path='/content/gdrive/MyDrive/i2dl/exercise_11'\n",
    "\n",
    "# This will mount your google drive under 'MyDrive'\n",
    "drive.mount('/content/gdrive', force_remount=True)\n",
    "# In order to access the files in this notebook we have to navigate to the correct folder\n",
    "os.chdir(gdrive_path)\n",
    "# Check manually if all files are present\n",
    "print(sorted(os.listdir()))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from exercise_code.util import notebook_util as util\n",
    "from exercise_code.network import *\n",
    "from exercise_code.tests import *\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "root_path = os.path.dirname(os.path.abspath(os.getcwd()))\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aa1117d",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "    <h3>WARNING</h3>\n",
    "    <p>The next section sort of tries to give you an intuition where the formulas for the positional encoding come from! You do not have to go through every single math equation or code line in this section, rather have a look at the graphs and try to get a feeling for what is going on! \n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4a008ca161f44de",
   "metadata": {},
   "source": [
    "### Intuition behind Positional Encoding\n",
    "\n",
    "Suppose we use a 4d embedding and we want to encode the position of a token. A possible choice could be to simply use the binary representation of the position as an encoding. So in other words, to the token at position 5, we would add the vector [0 1 0 1]. Here are all the binary representations up to 15:\n",
    "\n",
    "\n",
    "| Decimal                      | Binary                                      | Decimal                                  | Binary |\n",
    "|-------------------------------------|----------------------------------------------------|--------------------------------------------------|--|\n",
    "| 00 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">0</span> | 08 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">0</span> </br> |\n",
    "| 01 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">1</span> | 09 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">1</span> </br> |\n",
    "| 02 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">0</span> | 10 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">0</span> </br> |\n",
    "| 03 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">1</span> | 11 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">0</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">1</span> </br> |\n",
    "| 04 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">0</span> | 12 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">0</span> </br> |\n",
    "| 05 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">1</span> | 13 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">0</span> <span style=\"color: cyan;\">1</span> </br> |\n",
    "| 06 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">0</span> | 14 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">0</span> </br> |\n",
    "| 07 | <span style=\"color: green;\">0</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">1</span> | 15 | <span style=\"color: green;\">1</span> <span style=\"color: orange;\">1</span> <span style=\"color: red;\">1</span> <span style=\"color: cyan;\">1</span> </br> |\n",
    "\n",
    "In the next plots, you will often see the term depth and position. Depth refers to embedding depth, so if we have 4 dimensional embeddings, the depth index will tell you where in the embedding to add this \"bias\" encoding. Position, as the name suggests, tells us where in the sentence to add this particular \"bias\" vector.\n",
    "\n",
    "Let's plot this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef61ef30996801bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We are flipping the bits around, so that the least significant bit is plotted first! This will make the formulas a little easier later!\n",
    "pos_encoding = np.flip(np.array([[0, 0, 0, 0],\n",
    "                                 [0, 0, 0, 1],\n",
    "                                 [0, 0, 1, 0],\n",
    "                                 [0, 0, 1, 1],\n",
    "                                 [0, 1, 0, 0],\n",
    "                                 [0, 1, 0, 1],\n",
    "                                 [0, 1, 1, 0],\n",
    "                                 [0, 1, 1, 1],\n",
    "                                 [1, 0, 0, 0],\n",
    "                                 [1, 0, 0, 1],\n",
    "                                 [1, 0, 1, 0],\n",
    "                                 [1, 0, 1, 1],\n",
    "                                 [1, 1, 0, 0],\n",
    "                                 [1, 1, 0, 1],\n",
    "                                 [1, 1, 1, 0],\n",
    "                                 [1, 1, 1, 1]]), axis=-1)\n",
    "\n",
    "util.plot_positional_encoding(pos_encoding)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4060a487",
   "metadata": {},
   "source": [
    "As you can see, every bit is jumping between 0 and 1 with a different frequency! </br>\n",
    "The <span style=\"color: cyan;\">0th bit</span> is jumping back and forth after 1 iteration. </br>\n",
    "The <span style=\"color: red;\">1st bit</span> is jumping back and forth after 2 iteration. </br>\n",
    "The <span style=\"color: orange;\">2nd bit</span> is jumping back and forth after 4 iteration. </br>\n",
    "The <span style=\"color: green;\">3rd bit</span> is jumping back and forth after 8 iteration. </br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9f390b7f5940f9d",
   "metadata": {},
   "source": [
    "We can get the same pattern using a sine wave with different frequencies! In a first step, let's just have a look at the sign (+/-) of the sine function!\n",
    "\n",
    "$ PE(pos) = \\begin{bmatrix}\n",
    "           -sin(\\frac{\\pi}{1} \\cdot (pos + 0.5)) > 0 \\\\\n",
    "           -sin(\\frac{\\pi}{2} \\cdot (pos + 0.5)) > 0 \\\\\n",
    "           -sin(\\frac{\\pi}{4} \\cdot (pos + 0.5)) > 0 \\\\\n",
    "           -sin(\\frac{\\pi}{8} \\cdot (pos + 0.5)) > 0\n",
    "         \\end{bmatrix} $\n",
    "         \n",
    " You might be surprised to see the negative sine here! This only because we wanted it to start negative (=0) for the first half and then become positive (=1) in the second half to match our bit pattern! Also, the pos + 0.5 might be a bit weird at first, but we only did this to compute the value of the sine wave in the middle of each cell! But it is not important, just have a look at the plot of this function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf53700caab9ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_encoding_discrete(pos):\n",
    "    return np.array([-np.sin((np.pi / 1) * (pos + 0.5)) > 0,        \n",
    "                     -np.sin((np.pi / 2) * (pos + 0.5)) > 0,\n",
    "                     -np.sin((np.pi / 4) * (pos + 0.5)) > 0,\n",
    "                     -np.sin((np.pi / 8) * (pos + 0.5)) > 0]).T\n",
    "\n",
    "positions = np.arange(0, 16)\n",
    "\n",
    "util.plot_positional_encoding(positional_encoding_discrete, \n",
    "                              positions=positions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68ae4b4d03bb1fef",
   "metadata": {},
   "source": [
    "Tadaaaa - Same pattern!\n",
    "\n",
    "Now instead of using only the integers 0 and 1, let's use all values between 0 and 1! We will keep the rest the same for now! (That means we get rid of the sign() function, or the > 0 in code!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca10699c9b910e31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_encoding_continuous(pos):\n",
    "    return np.array([-np.sin((np.pi / 1) * (pos + 0.5)),    \n",
    "                     -np.sin((np.pi / 2) * (pos + 0.5)),\n",
    "                     -np.sin((np.pi / 4) * (pos + 0.5)),\n",
    "                     -np.sin((np.pi / 8) * (pos + 0.5))]).T\n",
    "\n",
    "positions = np.arange(0, 16)\n",
    "\n",
    "util.plot_positional_encoding(positional_encoding_continuous, \n",
    "                              positional_encoding_discrete, \n",
    "                              positions=positions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cac077879c9ab8d",
   "metadata": {},
   "source": [
    "You should still sort of be able to see the original pattern, its just a bit blurred now!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93438f1a7999a0b2",
   "metadata": {},
   "source": [
    "From here, we are going to do a couple alterations to the formula: First of all, get rid of the pos + 0.5! That really was only there, to show the similarity to the bit pattern!\n",
    "\n",
    "$ PE(pos) = \\begin{bmatrix}\n",
    "           sin(\\frac{\\pi}{1} \\cdot pos) \\\\\n",
    "           sin(\\frac{\\pi}{2} \\cdot pos) \\\\\n",
    "           sin(\\frac{\\pi}{4} \\cdot pos) \\\\\n",
    "           sin(\\frac{\\pi}{8} \\cdot pos)\n",
    "         \\end{bmatrix} $\n",
    "\n",
    "In a next step, we will also get rid of scaling frequency by $\\pi$! This leads us to:\n",
    " \n",
    "$ PE(pos) = \\begin{bmatrix}\n",
    "           sin(\\frac{1}{1} \\cdot pos) \\\\\n",
    "           sin(\\frac{1}{2} \\cdot pos) \\\\\n",
    "           sin(\\frac{1}{4} \\cdot pos) \\\\\n",
    "           sin(\\frac{1}{8} \\cdot pos)\n",
    "         \\end{bmatrix} $\n",
    "         \n",
    "Let's have a closer look at these angle frequencies:\n",
    "\n",
    "$\\omega_0 = \\frac{1}{1} =  \\frac{1}{2}^0 $ \\\n",
    "$\\omega_1 = \\frac{1}{2} =  \\frac{1}{2}^1 $ \\\n",
    "$\\omega_2 = \\frac{1}{4} =  \\frac{1}{2}^2 $ \\\n",
    "$\\omega_3 = \\frac{1}{8} =  \\frac{1}{2}^3 $ \n",
    "\n",
    "The frequencies form a geometric series with base $\\frac{1}{2}$! Writing this as a formula with base b we get:\n",
    "\n",
    "$\\omega_i = b^i$\n",
    "\n",
    "Now one last change we are going to do is to add a factor d in to the exponent as follows:\n",
    "\n",
    "$\\omega_i = b^{i/d}$\n",
    "\n",
    "Remeber - i here denotes the i-th dimension of our positional encoding vector! Each Dimension \"vibrates\" at a it's own frequency!\n",
    "\n",
    "This results in the follwing formula:\n",
    "\n",
    "$ PE(pos, i) = sin(b^{i/d} \\cdot pos)$\n",
    "\n",
    "Let's have a look at the final outcome for different d's! (This will sort of look like we are zooming!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d723f81c28c5ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_encoding(pos, base, d):\n",
    "    return np.array([np.sin(base ** 0/d * pos),    \n",
    "                     np.sin(base ** 1/d * pos),\n",
    "                     np.sin(base ** 2/d * pos),\n",
    "                     np.sin(base ** 3/d * pos)]).T\n",
    "\n",
    "positions = np.arange(0, 16)\n",
    "\n",
    "d_factors = [1, 2, 5, 10]  # Different zoom factors for the plots\n",
    "\n",
    "util.plot_positional_encoding(positional_encoding, \n",
    "                              positions=positions, \n",
    "                              d_factors=d_factors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc2899be73355355",
   "metadata": {},
   "source": [
    "Puh, that was a lot! Now with all this prep, the actual formulas shouldn't come as that big of a surprise anymore!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8156af05ac23e51e",
   "metadata": {},
   "source": [
    "### Implementing Positional Encodings\n",
    "\n",
    "The positional encodings used in the paper are actually very similar\n",
    "\n",
    "$PE(pos, 2i) = \\sin(pos / 10000 ^{2i/d})$ \\\n",
    "$PE(pos, 2i+1) = \\cos(pos / 10000 ^{2i/d})$\n",
    "\n",
    "At a first glance, there is a lot going on here! Rewriting the formula actually helps a lot!\n",
    "\n",
    "$ PE(pos) = \\begin{bmatrix}\n",
    "           sin(w_0 \\cdot pos) \\\\\n",
    "           cos(w_0 \\cdot pos) \\\\\n",
    "           sin(w_1 \\cdot pos) \\\\\n",
    "           cos(w_1 \\cdot pos) \\\\\n",
    "           \\vdots \\\\\n",
    "           sin(w_{d/2} \\cdot pos) \\\\\n",
    "           cos(w_{d/2} \\cdot pos)\n",
    "         \\end{bmatrix} $\n",
    "\n",
    "With:\n",
    "$ \\omega_i = \\frac{1}{10000}^{2i/d} $ \n",
    "\n",
    "Where:\n",
    "\n",
    "- $pos$ refers to the position of the token in the sequence </br>\n",
    "- $d$ refers to the dimension of the embedding (=d_model)\n",
    "\n",
    "Instead of only using sine functions, they also used cosine functions at every other index (note: because of this, you should always use an even embedding dimension!) </br>\n",
    "In the original formula, the \"every other index\" is noted by $2i$ and $2i+1$, which basically means all even indices get sine functions and all uneven cosine! "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "143baed9174cb10",
   "metadata": {},
   "source": [
    "Some of you might be curious, why they added the cosine function!\n",
    "\n",
    "To quote the paper:\n",
    "\n",
    "> We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset k, PE(pos+k) can be represented as a linear function of PE(pos).\n",
    "\n",
    "Without the cosine function, this cannot be done anymore!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3da2bb3ae45d691e",
   "metadata": {},
   "source": [
    "Anyway, let's have a look at this positional encoding! Since this encoding doesn't depend on data and is constant over training, we can compute it one time and store it as a large tensor as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e142d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_encoding(length, depth):\n",
    "\n",
    "    exponent = np.arange(0, depth, 2) / depth           # Exponent for the positional encoding\n",
    "    pos = np.arange(0, length)[:, None]                 # Add new axis for broadcasting\n",
    "\n",
    "    angle_freq = np.exp(exponent * (-np.log(10000)))    # For numerical reasons - same as (1/10000) ** (exponent)\n",
    "    \n",
    "    pos_encoding = np.zeros((length, depth))            # Initialize the positional encoding\n",
    "    \n",
    "    pos_encoding[:, 0::2] = np.sin(pos * angle_freq)    # Take the sine of the even indices\n",
    "    pos_encoding[:, 1::2] = np.cos(pos * angle_freq)    # Take the cosine of the odd indices\n",
    "\n",
    "    return pos_encoding\n",
    "\n",
    "util.plot_positional_encoding(positional_encoding,\n",
    "                              length=2048, \n",
    "                              depth=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d737d1d",
   "metadata": {},
   "source": [
    "The version I imagine you are probably more familiar with is the following: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "742a6c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_encoding(length, depth):\n",
    "\n",
    "    i = np.arange(0, depth, 2) / depth\n",
    "    pos = np.arange(0, length)[:, None]\n",
    "\n",
    "    angle_freq = np.exp(i * (-np.log(10000)))\n",
    "\n",
    "    pos_encoding = np.concatenate([np.sin(pos * angle_freq), np.cos(pos * angle_freq)], axis=-1)\n",
    "\n",
    "    return pos_encoding\n",
    "\n",
    "util.plot_positional_encoding(positional_encoding,\n",
    "                              length=2048, \n",
    "                              depth=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c954072c",
   "metadata": {},
   "source": [
    "The only difference is that instead of using sine for the even and cosine for the uneven we just concatenate the two vectors. Both versions are valid, since the main thing we want to achieve with this encoding is that vectors close by have a higher score and vectors further apart have a lower score. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69286b10",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_encoding = positional_encoding(length=2048, depth=512)\n",
    "pos_encoding = pos_encoding / np.linalg.norm(pos_encoding, axis=-1, keepdims=True)\n",
    "\n",
    "p = pos_encoding[1000]\n",
    "scores = pos_encoding @ p\n",
    "\n",
    "util.plot_scores(scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c71a378",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "    <h3>Task 4: Implement</h3>\n",
    "    <p>Implement the positional encoding method and initialize it in the <code>Embedding</code> class in <code>exercise_code/network/embedding.py</code>.\n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a171f64",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = test_task_4()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congrats! You've now finished your first transformer model!\n",
    "\n",
    "To create a zip file with your submission, run the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from exercise_code.util.submit import submit_exercise\n",
    "\n",
    "path = os.path.join(root_path, 'output', 'exercise11')\n",
    "submit_exercise(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submission Instructions\n",
    "\n",
    "To complete the exercise, submit your final model to our submission portal - you probably know the procedure by now.\n",
    "\n",
    "1. Go on [our submission page](https://i2dl.vc.in.tum.de/submission/), register for an account and login. We use your matriculation number and send an email with the login details to the mail account associated. When in doubt, login into tum online and check your mails there. You will get an ID which we need in the next step.\n",
    "2. Log into [our submission page](https://i2dl.vc.in.tum.de/submission/) with your account details and upload the `zip` file. Once successfully uploaded, you should be able to see the submitted file selectable on the top.\n",
    "3. Click on this file and run the submission script. You will get an email with your score as well as a message if you have surpassed the threshold.\n",
    "\n",
    "# Submission Goals\n",
    "\n",
    "- Goal: Successfully implement the attention mechanism!\n",
    "\n",
    "- Points:\n",
    "    - 5 points per Module if shape is correct (Embedding, Positional Encoding, Scaled Dot Attention, Multi-Head Attention)\n",
    "    - 5 points per Module if output is correct (Embedding, Positional Encoding, Scaled Dot Attention, Multi-Head Attention)\n",
    "    - Total = 4 x 5 + 4 x 5 = 40\n",
    "\n",
    "- Passing Criteria: Minimum of 35 points!\n",
    "- Feel free to submit an unlimited number of assignments until the end of the semester; however, any submissions made after the deadline will not contribute to your bonus points.\n",
    " "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "i2dl",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
