{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "026e0255",
   "metadata": {},
   "source": [
    "# Attention is All You Need\n",
    "Over the next two final exercise, we are going to put together a transformer model from scratch! The most important building block, as the title of the paper might the suggest, is the attention block / mechanism! We are going to lay the ground works for exercise 12 in this exercise, to make sure you are familiar with the basics! The implementation will try to stick to the paper as closly as possible, so there is really no need to read it! But you are of course welcome to! So let's start this - and for you it's really time to pay attention!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c3c01bc",
   "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,
   "id": "32e6fa9b",
   "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": "code",
   "execution_count": null,
   "id": "7f568e01",
   "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",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c884d3c8",
   "metadata": {},
   "source": [
    "# Attention Mechanism\n",
    "\n",
    "There are multiple ways of understanding how the attention mechanism works! In the last two semesters, we mainly used the \"Improving Measurement\" analogy which is a lot more straightforwad. Some of you might prefer the second explanation right after this, comparing attention to a detective mistery. This is optional to read! Please let us know, in the feedback form at the end of the exercise, which of the two you prefered or if you found the combination of both to be the best way to go!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b57a92e8",
   "metadata": {},
   "source": [
    "### Improving Noisy Measurements\n",
    "Imagine you have taken some measurement x(t), but notice that it is quite noisy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f386ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "measurement_data = util.get_measurement_data()\n",
    "plt.plot(measurement_data['time'], measurement_data['data_noise'])\n",
    "plt.xlabel('Time'), plt.ylabel('Data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ffa4163",
   "metadata": {},
   "source": [
    "So a way we could try to improve these measurements is by smoothing out the values using a weighted average:\n",
    "\n",
    "\n",
    "$\\tilde{x}_i = \\sum_{j=1}^N s_{ij} x_j$\n",
    "\n",
    "We could base the weights on a similarity score between the time stamps. Values that are closer together time-wise, should have a larger impact/weight compared to values that are further apart!\n",
    "\n",
    "$s_{ij} = sim(t_i, t_j)$\n",
    "\n",
    "A possible option as a similarity function is the squared exponential kernel:\n",
    "\n",
    "$sim(t_i, t_j) = \\exp\\left(-\\frac{(t_i - t_j)^2}{\\sigma ^2}\\right)$\n",
    "\n",
    "Let's have a look at the similarity function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "415831be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the squared exponential kernel \n",
    "def similarity(x1, x2):\n",
    "    sigma = 1\n",
    "    return np.exp(-((x1 - x2)**2 / (sigma ** 2)))\n",
    "\n",
    "# Define the time stamps\n",
    "time_stamps = np.linspace(-10, 10, 1000)\n",
    "\n",
    "# Compute the similarity scores between 0 and all other time stamps\n",
    "scores = similarity(0, time_stamps)\n",
    "\n",
    "# Plot the similarity scores\n",
    "plt.plot(time_stamps, scores)\n",
    "plt.title('Similarity Function')\n",
    "plt.xlabel('Time'), plt.ylabel('Score')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9141f8c2dd8712f2",
   "metadata": {},
   "source": [
    "As you can see, the closer we get to the zero, the higher the score gets!\n",
    "With a little bit of index magic, we can calculate the score for every single time pair $(t_i, t_j)$ in one step! We will store all of these in a score matrix: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc964ee8068ce519",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the actual timestamps from the measurement data\n",
    "time_stamps = measurement_data['time']\n",
    "\n",
    "# Compute the similarity scores between all time stamps\n",
    "scores = similarity(time_stamps[:, None], time_stamps[None, :])\n",
    "\n",
    "# Plot the similarity scores\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.imshow(scores, interpolation='nearest')\n",
    "plt.title('Score Matrix')\n",
    "plt.colorbar()\n",
    "plt.xlabel('Time Stamp ID'), plt.ylabel('Time Stamp ID')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9582c9b",
   "metadata": {},
   "source": [
    "What you should see a high value along the diagonal, and the values decrease the further you move away from the diagonal. \n",
    "\n",
    "Next, let's compute these averages! Note, that the average can also be described very efficiently as a Matrix-Vector product using the score matrix we calculated before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddf4ea86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute filtered data\n",
    "filter_data = scores @ measurement_data['data_noise']\n",
    "\n",
    "# Plot the filtered data\n",
    "plt.plot(time_stamps, filter_data)\n",
    "plt.xlabel('Time'), plt.ylabel('Data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "631cc387",
   "metadata": {},
   "source": [
    "Looks a lot better! However, we are actually filtering out a lot of the underlying signal. Try changing sigma in the next cell to a different values (e.g $\\sigma = 0.5$ or even $\\sigma = 0.1$) and have a look at the function!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5125b548",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Edit Sigma here\n",
    "sigma = 1.0\n",
    "\n",
    "def similarity(x1, x2):\n",
    "    return np.exp(-((x1 - x2)**2 / (sigma ** 2)))\n",
    "\n",
    "# Compute the similarity scores between all time stamps\n",
    "scores = similarity(time_stamps[:, None], time_stamps[None, :])\n",
    "\n",
    "# Plot the similarity scores\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.imshow(scores, interpolation='nearest')\n",
    "plt.title('Score Matrix')\n",
    "plt.colorbar()\n",
    "plt.xlabel('Time Stamp ID'), plt.ylabel('Time Stamp ID')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa9b1acc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute filtered data\n",
    "filter_data = scores @ measurement_data['data_noise']\n",
    "\n",
    "# Plot the filtered data\n",
    "plt.plot(time_stamps, filter_data)\n",
    "plt.xlabel('Time'), plt.ylabel('Data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0e418d1",
   "metadata": {},
   "source": [
    "If you did it then - Nice! We are not loosing that much information anymore! If not - Come on, it's like only one line ;P"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bb6c816",
   "metadata": {},
   "source": [
    "We still have one problem in our implementation: If you have a look at the scale of our data, it has changed quite a bit!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b11fe4df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot filtered data and the original measurement data\n",
    "plt.figure()\n",
    "plt.plot(time_stamps, measurement_data['data_noise'])\n",
    "plt.plot(time_stamps, filter_data)\n",
    "plt.xlabel('Time'), plt.ylabel('Data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "067a8e9e",
   "metadata": {},
   "source": [
    "We can somewhat improve this by normalizing our scores to sum up to one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "915002a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normaliize scores by dividing by the sum of all scores\n",
    "scores_normalized = scores / np.sum(scores, axis=0)\n",
    "\n",
    "# Calculate the filtered data again using the normalized weights\n",
    "filter_data = scores_normalized @ measurement_data['data_noise']\n",
    "\n",
    "# Plot the filtered data\n",
    "plt.plot(time_stamps, filter_data)\n",
    "plt.xlabel('Time'), plt.ylabel('Data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc5b10a4",
   "metadata": {},
   "source": [
    "Now the scales should match a lot better!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce607be2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot filtered data with normalized scores and the original measurement data\n",
    "plt.figure()\n",
    "plt.plot(time_stamps, measurement_data['data_noise'])\n",
    "plt.plot(time_stamps, filter_data)\n",
    "plt.xlabel('Time'), plt.ylabel('Data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f3e4cb",
   "metadata": {},
   "source": [
    "Another way of thinking about this, is that we have \"improved\" our data, by introducing context from other points to it, i.e. each datapoint looks at its neighboring datapoints. We can also introduce the terminology used in the paper: For every **query** (timestamp) $t_i$, compare it to all available **keys** $t_j$ (all other timestamps) and compute **similarity scores**. Multiply these scores with the corresponding **values** $x_j$ (the measurements) to form the \"answer\" to the query $t_i$. Note, that each key corresponds to a value!\n",
    "\n",
    "Now let's see if we can find something similar for words. After all, thats what we want our model to work with."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65194c6c",
   "metadata": {},
   "source": [
    "### Applying Similarity-based Context to Word Embeddings\n",
    "\n",
    "The concept of attention mirrors the method we used earlier:\n",
    "\n",
    "When we transform a sentence into a sequence of embeddings $x_i$, translating each word independently may cause an issue: context matters!\n",
    "\n",
    "Consider these sentences:\n",
    "- The river bank is very long.\n",
    "- The bank is located at the long river.\n",
    "\n",
    "Both contain the word 'bank', but its meaning changes with context. Think about what this also means for the embeddings, which are supposed to \"capture the meaning\" of each word: The embedding for 'bank' should ideally differ in both sentences!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87c6f96a",
   "metadata": {},
   "source": [
    "Our current implementation of the embedding layer does not take context (i.e. the other words in the sentence) into account, it simply maps each word to a vector from a static lookup table. Therefore, we need to find a way to **contextualize** the embeddings!\n",
    "\n",
    "Before we get into the details we will again define a set of \n",
    "\n",
    "- Queries Q - Think of it as what the word is looking for\n",
    "- Keys K - This is kind of like what the word looks like\n",
    "- Values V - Basically what the word has to offer\n",
    "\n",
    "These three vectors are in a way also embeddings (vector representations of words). We will discuss later on how we create them, but they are basically just a function of the original embeddings X.\n",
    "\n",
    "With that beeing said, the simplest way to create contextualized embeddings, similar to before, is to take a weighted average of the embeddings of all words we want to get context from - our values $v_j$:\n",
    "\n",
    "$\\tilde{x}_i = \\sum_{j=1}^N s_{ij} v_j$\n",
    "\n",
    "We define the scores as before, by computing the similarity between the query $q_i$ and the keys $k_j$: \n",
    "\n",
    "$s_{ij} = \\text{sim}(q_i, k_j)$\n",
    "\n",
    "Does this make sense? The learned embeddings ideally represent similar words that co-occur or are semantically linked (like 'apple' and 'fruit'). Conversely, words with low co-occurrence, like 'computer' and 'elephant', should be dissimilar. So in way it does make sense, that \"similar\" words should contribute more to each other's context (but we will revisit this later).\n",
    "\n",
    "Now, about the similarity function: instead of the squared exponential kernel, we use the dot product! But does the dot product truly measure similarity between vectors?\n",
    "\n",
    "$\\langle a,b \\rangle = \\cos(\\angle_{ab})|a||b|$\n",
    "\n",
    "We can see from here, that the closer two embeddings are, the higher the dot product will become! (Of course the norm of the vectors also affects the score)\n",
    "\n",
    "Let's define the similarity function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2546d9bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define similarity function as the inner product\n",
    "def similarity(q, k):\n",
    "    return np.inner(q, k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9099330",
   "metadata": {},
   "source": [
    "We have prepared a pretrained embedding model for you, so you can play around with it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3d5af09",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the word2vec model\n",
    "word2vec = util.load_word2vec()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e2d1f3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the word vectors for apple and fruit\n",
    "word_1 = word2vec['apple']\n",
    "word_2 = word2vec['fruit']\n",
    "\n",
    "q_1 = util.get_query(word_1)\n",
    "k_2 = util.get_key(word_2)\n",
    "\n",
    "# Compute the similarity between apple and fruit\n",
    "similarity(q_1, k_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92c74c09",
   "metadata": {},
   "source": [
    "Ok, and now for a reference, let's look at two words that should not be related to much!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da1be26c",
   "metadata": {},
   "outputs": [],
   "source": [
    "word_1 = word2vec['computer']\n",
    "word_2 = word2vec['elephant']\n",
    "\n",
    "q_1 = util.get_query(word_1)\n",
    "k_2 = util.get_key(word_2)\n",
    "\n",
    "similarity(q_1, k_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f42476f",
   "metadata": {},
   "source": [
    "Ok, so let's look at what exactly we have to do: Compute the similiarity scores to each key, and form the answer to our query. Here is a visualualization showing this flow: \n",
    "<!-- <img src=\"images/Transformer-AttentionMechanism.drawio.png\" width=2000> -->\n",
    "<img src=\"https://i2dl.vc.in.tum.de/static/images/exercise_11/Transformer-AttentionMechanism.drawio.png\" width=2000>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f18f992",
   "metadata": {},
   "source": [
    "Now let's do this in code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaea146f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an example sentence\n",
    "sentence = \"The river bank is very long\"\n",
    "\n",
    "# Get the word vectors for the words in the sentence\n",
    "embeddings, words = util.embedd_sentence(sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22878f25",
   "metadata": {},
   "source": [
    "For now, we will set the queries, keys and values to be the same! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77817b92",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = util.get_query(embeddings)\n",
    "keys = util.get_key(embeddings)\n",
    "values = util.get_value(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d9691da",
   "metadata": {},
   "source": [
    "Let us have a quick look at the shapes of the matricies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c536594f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(queries.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "852514ad",
   "metadata": {},
   "source": [
    "So we have a set of 6 embeddings - we started with 6 words / tokens, and the embedding dimension is 50!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59f68d83",
   "metadata": {},
   "source": [
    "Now, instead of calculating the score for each query separately, we can use matrix multiplication to calculate all the dot products between the query and keys in one go:\n",
    "\n",
    "$s_{ij} = q_i^T k_j$ or in matrix notation $S = QK^T$  \n",
    "\n",
    "This obviously only works if the matrices are oriented in the correct way!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b1284d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = queries @ keys.T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eb43c82",
   "metadata": {},
   "source": [
    "Similar to the example with the timeseries, we want to make sure that the scores that correspond to a query all add to up one! We will be using the softmax in this case to ensure positive scores as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3507801",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = util.softmax(scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda9a3ba",
   "metadata": {},
   "source": [
    "The final step is contextualizing the embeddings by summing over all values, weighted by the scores! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28cba081",
   "metadata": {},
   "outputs": [],
   "source": [
    "contextualized_embeddings = scores @ values "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e3b654",
   "metadata": {},
   "source": [
    "And voila, thats basically all there is to the attention mechanism! If we look at the formula in the paper you will find\n",
    "\n",
    "$Attention(Q,K,V) = softmax\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right) V$\n",
    "\n",
    "The value $d_k$ corresponds to the dimension of the keys and scales the dot product results. It is added to keep the dot product from growing, when the embedding dimensions is increased. (If you increase the dimension, you also increase the number of elements you have to add together!) However, larger values of the dot products push the softmax function into regions with smaller gradients, which can slow down training significantly. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9461172d61753b1",
   "metadata": {},
   "source": [
    "<!-- <img src=\"images/Transformer-Attention Head.drawio.png\" width=\"2000\"> -->\n",
    "<img src=\"https://i2dl.vc.in.tum.de/static/images/exercise_11/Transformer-Attention Head.drawio.png\" width=2000>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8654ae2",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "    <h3>Task 2: Implement</h3>\n",
    "    <p>Implement the <code>forward</code> pass in the class <code>ScaledDotAttention</code>. You can find it in <code>exercise_code/networks/attention.py</code>!\n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f602f12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you get an Error about SCORE_SAVER - please just restart your kernel!\n",
    "_ = test_task_2()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb70a73d",
   "metadata": {},
   "source": [
    "# Attention as a \"Detective Mystery\" (Optional)\n",
    "\n",
    "As promised, the other way of understanding it:\n",
    "\n",
    "Imagine a team of detectives working together to solve a case. Each detective represents a word in a sentence, and their shared goal is to uncover the full story — how each piece of evidence fits together and what role they play in solving the case."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "310c6060",
   "metadata": {},
   "source": [
    "## How the Detective Mystery Works\n",
    "\n",
    "1. **Detectives Start With Initial Knowledge**:  \n",
    "   Each detective begins with a specific clue (word embedding) $x_i$ , representing their initial understanding of the case.\n",
    "\n",
    "2. **Detectives Ask Questions (Queries)**:  \n",
    "   Each detective has their own perspective on the case and asks questions to the group, like:\n",
    "   - *\"What evidence is related to this fingerprint?\"*\n",
    "   - *\"How does this broken vase fit into the crime scene?\"*\n",
    "\n",
    "   In the attention mechanism, these questions are represented as **queries**, generated from the detective’s initial clue:\n",
    "   $$\n",
    "   q_1 = f_Q(x_1)\n",
    "   $$\n",
    "\n",
    "3. **Other Detectives Offer Leads (Keys and Values)**:  \n",
    "   The rest of the team responds by offering **leads**:\n",
    "   - **Keys**: Describing what they know based on what evidence they have:\n",
    "     $$\n",
    "     k_1 = f_K(x_1), \\quad k_2 = f_K(x_2), \\ldots\n",
    "     $$\n",
    "   - **Values**: Representing the detailed information they can share:\n",
    "     $$\n",
    "     v_1 = f_V(x_1), \\quad v_2 = f_V(x_2), \\ldots\n",
    "     $$\n",
    "\n",
    "   For example:\n",
    "   - One detective might say: *\"I have a witness statement that matches your question.\"* In our analogy, this answer would form the key! The actual witness statement would make up the value!\n",
    "   - Another might say: *\"I’ve seen a similar clue at another crime scene.\"*\n",
    "\n",
    "4. **Weighting the Relevance of Leads**:  \n",
    "   Each detective compares their question (query) to the leads (keys) from others. The better the match between their question and a lead, the more weight that lead is given:\n",
    "   $$\n",
    "   s_{1j} = q_1 \\cdot k_j\n",
    "   $$\n",
    "   A strong alignment means the lead is highly relevant to the question.\n",
    "\n",
    "5. **Piecing Together the Clues (Updating Beliefs)**:  \n",
    "   Each detective combines the information from the team, weighting the leads by their relevance scores to refine their understanding of the case:\n",
    "   $$\n",
    "   x_{\\text{new}, 1} = \\sum_{j} s_{1j} v_j\n",
    "   $$\n",
    "   In practical terms:\n",
    "   - If Detective 1 asks about fingerprints and gets strong responses from Detectives 2 and 3, they’ll update their belief accordingly, focusing on fingerprint-related evidence.\n",
    "   - The scores tell the detectives, which parts to give more attention too, which is also where the name comes from!\n",
    "\n",
    "   To stabilize this process:\n",
    "   - Scores are normalized with a **softmax function**, ensuring the total influence adds up to 1:\n",
    "     $$\n",
    "     S = \\text{softmax}(Q K^T)\n",
    "     $$\n",
    "   - Scaling by $ \\sqrt{d_k} $ prevents large dimensions from skewing results:\n",
    "     $$\n",
    "     S = \\text{softmax}\\left(\\frac{Q K^T}{\\sqrt{d_k}}\\right)\n",
    "     $$\n",
    "\n",
    "6. **The Team Solves the Case**:  \n",
    "   After several rounds of questioning and refining, each detective gains a clearer understanding of the story and their role in it. Together, they piece the clues into a coherent narrative.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4faad5d9",
   "metadata": {},
   "source": [
    "## Multi-Head Attention\n",
    "\n",
    "So far, we are able to give context to an input sequence - awesome! However, what each token needs to \"pay attention to\" in its context is a bit more complicated! We can imagine, for example, that a verb might need to pay more attention to the subject and object of the sentence, or in the case of the word \"bank\", it might need to pay more attention to the word \"river\" or \"money\" to figure out its own intended meaning. And since language is so complex, you can imagine that there are so many different aspects that each token needs to pay attention to! This is where the idea of multi-head attention comes in:\n",
    "\n",
    "Instead of just using the embeddings directly, we can use linear transformations to down-project into more *fine-grained* queries $QW^Q$, keys $KW^K$ and values $VW^V$ from the embeddings. This way, the model can learn which aspects of the embeddings are important! We call this one \"unit\" of attention mechanism an **attention head**:\n",
    "\n",
    "$$head(Q, K, V) = Attention(QW^Q, KW^K, VW^V)$$\n",
    "\n",
    "If we look at back at the detectives mistery, these linear projections are exactly those \"functions\" that we were talking about earlier!\n",
    "\n",
    "$$ Q = f_Q(X) = XW^Q$$\n",
    "\n",
    "And as the headline suggests, we can create multiple heads with their own weights and concatenate them together. This is very similar to a convolution layer with multiple filters in images! Each filter can learn different patterns in the image, and each head can learn different aspects to focus more attention to.\n",
    "\n",
    "Once all these heads are evaluated and concatenated, we have to somehow combine them together to retain the original shape of the input! This is done with a linear layer $W^O$:\n",
    "\n",
    "$MultiHead(Q, K, V) = Concat(head_1, ..., head_{n_{heads}}) \\cdot W^O$ \n",
    "\n",
    "where $head_i(Q, K, V) = Attention(QW_i^Q, KW_i^K, VW_i^V)$\n",
    "\n",
    "Finally, let's check the dimensions of the weights, where $d_{model}$ is the dimension of the embeddings:\n",
    "\n",
    "$shape(W^Q_i) = (d_{model},\\, d_q) = (d_{model},\\, d_k) $     Since we are computing dot products between the query vectors and the keys, the dimensions have to match!  \n",
    "$shape(W^K_i) = (d_{model},\\, d_k)$  \n",
    "$shape(W^V_i) = (d_{model},\\, d_v)$  \n",
    "$shape(W^O) = (d_v*n_{heads},\\, d_{model})$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa0e8c6d",
   "metadata": {},
   "source": [
    "<!-- <img src=\"images/Transformer-Multi Head Attention.drawio.png\" width=800> -->\n",
    "<img src=\"https://i2dl.vc.in.tum.de/static/images/exercise_11/Transformer-Multi%20Head%20Attention.drawio.png\" width=800>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7f6f99378d59d12",
   "metadata": {},
   "source": [
    "In our implementation, we are going to vectorize this expressions by creating three large weight matrices instead of a set for each head:\n",
    "\n",
    "$shape(W^Q) = (d_{model},\\, n_{heads} \\cdot d_k)$       \n",
    "$shape(W^K) = (d_{model},\\, n_{heads} \\cdot d_k)$  \n",
    "$shape(W^V) = (d_{model},\\, n_{heads} \\cdot d_v)$  \n",
    "\n",
    "That way we don't have to loop through the heads in python! Now we have to reshape the outputs and get them into the correct shape! Let's have a look at one example:\n",
    "\n",
    "$shape(Q \\times W^Q) = (seq_Q, d_{model}) \\times (d_{model},\\, n_{heads} \\cdot d_k) = (seq_Q,\\, n_{heads} \\cdot d_k)$ where $seq_Q$ is the sequence length of the queries\n",
    "\n",
    "Now we have to split the last dimension into two parts - $n_{heads}$ and  $d_k$. We can do this using torch.reshape!:\n",
    "\n",
    "$(seq_Q,\\, n_{heads} \\cdot d_k)$ -> $(seq_Q,\\, n_{heads}, d_k)$\n",
    "\n",
    "Now if we look at the last two dimensions, we actually expect it to be $seq_Q$ and $d_k$, so that we can run it through the attention block! \n",
    "\n",
    "In other words, we have to swap (or better - transpose) the tensor at the correct dimension:\n",
    "\n",
    "$(seq_Q,\\, n_{heads}, d_k)$ -> $(n_{heads}, seq_Q,\\, d_k)$\n",
    "\n",
    "And now we can feed this into our attention mechanism as we usually do, think of it as a batch of batches. The pytorch functions automatically handle this for us! "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cad2e95",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "    <h3>Task 3: Implement</h3>\n",
    "    <p>Implement the <code>__init__()</code> method and the <code>forward()</code> method of the <code>MultiHeadAttention</code> class in <code>exercise_code/network/multi_head_attention.py</code>. \n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9515509f",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = test_task_3()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75a90b1c",
   "metadata": {},
   "source": [
    "Awesome! We have now laid the foundations for building transformers from scratch! We will not implement them in this exercise, but make sure to check out exercise 12 when we release it! There is one more final thing to discuss, and that is why we need positional encodings! See you in the last notebook!"
   ]
  }
 ],
 "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": 5
}
