lab_attention.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most recent version of this notebook is available at https://github.com/nadiinchi/dl_labs/blob/master/lab_attention.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook contains a practical introduction to the mechanisms of attention.\n",
    "The notebook is considering a toy problem and architecture, so it can be implemented and computed on the CPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "There are $K$ objects that you can pay attention to.\n",
    "Each object is characterized by the key $k_i$ and the value $v_i$.\n",
    "The attention layer proceeds queries.\n",
    "For the query $q$, the layer returns a weighted sum of the values of the objects, with weights proportional to the degree of key matching the query:\n",
    "$$w_i = \\frac{\\exp(score(q, k_i))}{\\sum_{j=1}^K\\exp(score(q, k_j))}$$\n",
    "$$a = \\sum_{i=1}^K w_i v_i$$\n",
    "\n",
    "Almost always, queries, keys, and values are real vectors of some fixed dimensions.\n",
    "In the assignment it is proposed to implement three types of attention:\n",
    "+ Multiplicative Attention.\n",
    "Defined by function $score(q, k) = q^Tk$.\n",
    "To use this type of attention, it is required that the dimension of the query coincides with the dimension of the key.\n",
    "For more information see the paper Luong et al. \"Effective approaches to attention-based neural machine translation\", 2015.\n",
    "+ Scaled Dot Product Attention.\n",
    "Defined by function $score(q, k) = \\frac{q^Tk}{\\sqrt{dim(k)}}$, where $dim(k)$ is the dimensionality of the key (which also equals the dimensionality of the query).\n",
    "With learnable queries or keys, such attention is equivalent to multiplicative attention described above.\n",
    "However, immediately after initialization, such attention encourages smoother weights, which alleviates the problem of small gradients for saturated SoftMax.\n",
    "For more information see the paper Vaswani et al. \"Attention Is All You Need\", 2017.\n",
    "+ _(Optional, you may implement it after finishing the main part)._ Additive Attention. \n",
    "Defined by function $ score (q, k) = w_3 ^ T \\ tanh (W_1q + W_2k) $, where $ W_1, W_2, w_3 $ are the trainable parameters of the attention layer.\n",
    "For such a function, the request and key may have different dimensions.\n",
    "Matrices $ W_1 $ and $ W_2 $ map the query and key into a common hidden space, the dimension of which coincides with the dimension of the vector $ w_3 $.\n",
    "The dimension of the hidden space can be chosen arbitrarily and is a hyperparameter of the layer.\n",
    "For more information see the paper Bahdanau et al. \"Neural Machine Translation by Jointly Learning to Align and Translate\", 2014.\n",
    "\n",
    "In practice, the keys and values of the objects are usually the same vector by the model design.\n",
    "In this case keys and values of objects are also called objects' features (i. e. $f_i := k_i = v_i$).\n",
    "\n",
    "_An attention mask_ was proposed in the paper Vaswani et al. \"Attention Is All You Need\", 2017.\n",
    "Mask shows for each query which objects it cannot pay attention to.\n",
    "It sets $w_j := 0$ for $i$-th query if $mask_{i, j} = 1$.\n",
    "Note that weights still must sum up to $1$ for every query, i. e. the weights that was not zeroed must be renormalized.\n",
    "\n",
    "A mask where $mask_{i, j} = \\text{Indicator}[i < j]$ is used to ensure the autoregressive properties of the model.\n",
    "I. e. the layer output for the $i$-th position does not depend on the input values of subsequent positions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Practice\n",
    "\n",
    "Below it is proposed to implement several models of attention, described in the abovementioned articles.\n",
    "\n",
    "For the flexibility of the interface and for the acceleration of learning, all layers of attention below receive several requests for each object of the batch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Score layers\n",
    "\n",
    "In each of classes `MultiplicativeScore` and `ScaledDotProductScore`, you need to implement function `forward`.\n",
    "It takes a tensor of `queries` (of shape `batch_size x num_queries x query_dim`) and a tensor of objects' `features` (of shape `batch_size x num_objects x features_dim`).\n",
    "It computes for each pair of object and query inside batch their similarity score.\n",
    "The output tensor have shape `batch_size x num_queries x num_objects`.\n",
    "The code should be equivalent to the following:\n",
    "```\n",
    "res = torch.zeros(batch_size, num_queries, num_objects)\n",
    "for i in range(batch_size):\n",
    "    for j in range(num_queries):\n",
    "        for k in range(num_objects):\n",
    "            res[i, j, k] = score(queries[i, j], features[i, k])\n",
    "```\n",
    "\n",
    "Naturally, the above code is only an illustration explaining the dimensions of the arguments and the output of the `forward` function.\n",
    "Your implementation must be effectively vectorized.\n",
    "\n",
    "_Hint:_ It is recommended to pay attention to the function torch.bmm, it may be useful below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "import numpy as np\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiplicativeScore(nn.Module):\n",
    "    \"\"\"\n",
    "    Luong et al. \"Effective approaches to attention-based neural machine translation\", 2015.\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, queries, features):\n",
    "        \"\"\"\n",
    "        queries:  [batch_size x num_queries x dim]\n",
    "        features: [batch_size x num_objects x dim]\n",
    "        Returns a tensor of scores with shape [batch_size x num_queries x num_objects].\n",
    "        \"\"\"\n",
    "        # your code here\n",
    "        return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScaledDotProductScore(nn.Module):\n",
    "    \"\"\"\n",
    "    Vaswani et al. \"Attention Is All You Need\", 2017.\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, queries, features):\n",
    "        \"\"\"\n",
    "        queries:  [batch_size x num_queries x dim]\n",
    "        features: [batch_size x num_objects x dim]\n",
    "        Returns atensor of scores with shape [batch_size x num_queries x num_objects].\n",
    "        \"\"\"\n",
    "        # your code here\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tests for score function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[1., 0., 0., 0.],\n",
      "         [0., 0., 1., 0.],\n",
      "         [0., 0., 0., 0.]]])\n",
      "tensor([[[0.4472, 0.0000, 0.0000, 0.0000],\n",
      "         [0.0000, 0.0000, 0.4472, 0.0000],\n",
      "         [0.0000, 0.0000, 0.0000, 0.0000]]])\n"
     ]
    }
   ],
   "source": [
    "q = torch.tensor([[\n",
    "    [0, 0, 0, 0, 1],\n",
    "    [0, 0, 1, 0, 0],\n",
    "    [1, 0, 0, 0, 0],\n",
    "]], dtype=torch.float32)\n",
    "o = torch.tensor([[\n",
    "    [0, 0, 0, 0, 1],\n",
    "    [0, 0, 0, 1, 0],\n",
    "    [0, 0, 1, 0, 0],\n",
    "    [0, 1, 0, 0, 0],\n",
    "]], dtype=torch.float32)\n",
    "print(MultiplicativeScore()(q, o))\n",
    "print(ScaledDotProductScore()(q, o))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 3, 4])\n",
      "torch.Size([2, 3, 4])\n"
     ]
    }
   ],
   "source": [
    "q = torch.randn(2, 3, 5)\n",
    "o = torch.randn(2, 4, 5)\n",
    "print(MultiplicativeScore()(q, o).shape)\n",
    "print(ScaledDotProductScore()(q, o).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attention layer\n",
    "\n",
    "In `Attention` class it is necessary to implement function `forward`, which takes as inputs a set of `queries`, a set of `features` attributes for each batch object, and optionally an attention `mask` that is described above.\n",
    "The `Attnetion` object uses the score function it gets in the constructor (either `MultiplicativeScore`, `ScaledDotProductScore`, or other score function or functor).\n",
    "\n",
    "The most numerically stable way to use a mask in attend will be to set the corresponding score values to `float('-inf')` before applying SoftMax.\n",
    "An alternative method is to zero the scales $w$ according to the mask and renormalize them, but this method is less computationally stable (think why).\n",
    "\n",
    "To visualize the attention map at the end of the notebook, you need to save `weights.detach()` in `self.last_weights` in the `forward`. `.detach()` is used to enforce that the stored weights are not part of the computational graph. Do not forget to do `.detach()` for the debug output, so that the computational graph does not consume RAM beyond the required size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, scorer):\n",
    "        super().__init__()\n",
    "        self.scorer = scorer\n",
    "\n",
    "    def forward(self, queries, features, mask=None):\n",
    "        \"\"\"\n",
    "        queries:         [batch_size x num_queries x query_feature_dim]\n",
    "        features:        [batch_size x num_objects x obj_feature_dim]\n",
    "        mask, optional:  [num_queries x num_objects]\n",
    "        Returns matrix of responses for queries with shape [batch_size x num_queries x obj_feature_dim].\n",
    "        If mask is not None, sets corresponding to mask weights to zero.\n",
    "        Saves detached weights as self.last_weights for further visualization.\n",
    "        \"\"\"\n",
    "        scores = self.scorer(queries, features)\n",
    "        # your code here\n",
    "        weights = # your code here\n",
    "        self.last_weights = weights.detach()\n",
    "        # your code here\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tests for Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.4754, 0.1749, 0.1749, 0.1749],\n",
      "         [0.1749, 0.1749, 0.4754, 0.1749],\n",
      "         [0.2500, 0.2500, 0.2500, 0.2500]]])\n",
      "tensor([[[0.0000, 0.1749, 0.1749, 0.1749, 0.4754],\n",
      "         [0.0000, 0.1749, 0.4754, 0.1749, 0.1749],\n",
      "         [0.0000, 0.2500, 0.2500, 0.2500, 0.2500]]])\n"
     ]
    }
   ],
   "source": [
    "q = torch.tensor([[\n",
    "    [0, 0, 0, 0, 1],\n",
    "    [0, 0, 1, 0, 0],\n",
    "    [1, 0, 0, 0, 0],\n",
    "]], dtype=torch.float32)\n",
    "o = torch.tensor([[\n",
    "    [0, 0, 0, 0, 1],\n",
    "    [0, 0, 0, 1, 0],\n",
    "    [0, 0, 1, 0, 0],\n",
    "    [0, 1, 0, 0, 0],\n",
    "]], dtype=torch.float32)\n",
    "attn = Attention(MultiplicativeScore())\n",
    "a = attn(q, o)\n",
    "print(attn.last_weights)\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[2.1384e-01, 2.3633e-01, 2.6118e-01, 2.8865e-01],\n",
      "         [9.3572e-14, 2.0611e-09, 4.5398e-05, 9.9995e-01],\n",
      "         [0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00],\n",
      "         [0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]]])\n",
      "tensor([[[ 624.6472],\n",
      "         [1999.9546],\n",
      "         [2000.0000],\n",
      "         [2000.0000]]])\n"
     ]
    }
   ],
   "source": [
    "q = torch.tensor([[\n",
    "    [0.0001],\n",
    "    [0.01],\n",
    "    [1],\n",
    "    [100],\n",
    "]], dtype=torch.float32)\n",
    "o = torch.tensor([[\n",
    "    [-1],\n",
    "    [0],\n",
    "    [1],\n",
    "    [2],\n",
    "]], dtype=torch.float32) * 1000\n",
    "attn = Attention(MultiplicativeScore())\n",
    "a = attn(q, o)\n",
    "print(attn.last_weights)\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 3, 5])\n"
     ]
    }
   ],
   "source": [
    "q = torch.randn(2, 3, 5)\n",
    "o = torch.randn(2, 4, 5)\n",
    "print(attn(q, o).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tests for Attention with attention mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_autoregressive_mask(size):\n",
    "    \"\"\"\n",
    "    Returns attention mask of given size for autoregressive model.\n",
    "    \"\"\"\n",
    "    dtype = getattr(torch, 'bool', None) or torch.uint8\n",
    "    res = torch.zeros(size, size, dtype=dtype)\n",
    "    for i in range(size - 1):\n",
    "        res[i, i + 1:] = 1\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[False,  True,  True,  True],\n",
      "        [False, False,  True,  True],\n",
      "        [False, False, False,  True],\n",
      "        [False, False, False, False]], dtype=torch.bool)\n"
     ]
    }
   ],
   "source": [
    "m = get_autoregressive_mask(4)\n",
    "print(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.1749, 0.1749, 0.1749, 0.4754],\n",
      "         [0.1749, 0.1749, 0.4754, 0.1749],\n",
      "         [0.1749, 0.4754, 0.1749, 0.1749],\n",
      "         [0.4754, 0.1749, 0.1749, 0.1749]]])\n"
     ]
    }
   ],
   "source": [
    "q = torch.tensor([[\n",
    "    [0, 0, 0, 0, 1],\n",
    "    [0, 0, 0, 1, 0],\n",
    "    [0, 0, 1, 0, 0],\n",
    "    [1, 0, 0, 0, 0],\n",
    "]], dtype=torch.float32)\n",
    "o = torch.tensor([[\n",
    "    [1, 0, 0, 0, 0],\n",
    "    [0, 0, 1, 0, 0],\n",
    "    [0, 0, 0, 1, 0],\n",
    "    [0, 0, 0, 0, 1],\n",
    "]], dtype=torch.float32)\n",
    "a = attn(q, o)\n",
    "print(attn.last_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[1.0000, 0.0000, 0.0000, 0.0000],\n",
      "         [0.5000, 0.5000, 0.0000, 0.0000],\n",
      "         [0.2119, 0.5761, 0.2119, 0.0000],\n",
      "         [0.4754, 0.1749, 0.1749, 0.1749]]])\n"
     ]
    }
   ],
   "source": [
    "a = attn(q, o, m)\n",
    "print(attn.last_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],\n",
      "         [4.5398e-05, 9.9995e-01, 0.0000e+00, 0.0000e+00],\n",
      "         [0.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00],\n",
      "         [0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]]])\n",
      "tensor([[[-1.0000e+03],\n",
      "         [-4.5398e-02],\n",
      "         [ 1.0000e+03],\n",
      "         [ 2.0000e+03]]])\n"
     ]
    }
   ],
   "source": [
    "q = torch.tensor([[\n",
    "    [0.0001],\n",
    "    [0.01],\n",
    "    [1],\n",
    "    [100],\n",
    "]], dtype=torch.float32)\n",
    "o = torch.tensor([[\n",
    "    [-1],\n",
    "    [0],\n",
    "    [1],\n",
    "    [2],\n",
    "]], dtype=torch.float32) * 1000\n",
    "attn = Attention(MultiplicativeScore())\n",
    "a = attn(q, o, m)\n",
    "print(attn.last_weights)\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 4, 5])\n"
     ]
    }
   ],
   "source": [
    "q = torch.randn(2, 4, 5)\n",
    "o = torch.randn(2, 4, 5)\n",
    "print(attn(q, o, m).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Positional Encoder\n",
    "\n",
    "PositionalEncoder is the layer described in Vaswani et al. \"Attention Is All You Need\", 2017.\n",
    "The motivation is that attention is an operation over a set, i. e. it assumes no order in objects.\n",
    "In order to make attention over sequence position-aware, it was proposed to add to the object\n",
    "the vector of encoding of its position.\n",
    "\n",
    "It adds to the output of the previous layer embedding positions.\n",
    "In order not to recalculate position embeddings each time, its constructor receives the `max_len` parameter and precomputes embeddingings for positions from 0 to `max_len` - 1 inclusive.\n",
    "The add flag indicates whether to add position embeddings to the output of the previous layer (by default, with `add` = True, as was in the original paper) or concatenated (`add` = False).\n",
    "For the selected embedding dimensions, you should visualize the embeddings (plot the each component of the embeddings) and select the appropriate scale parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoder(nn.Module):\n",
    "    def __init__(self, dim, max_len=50, scale=10000.0, add=True):\n",
    "        \"\"\"\n",
    "        Transforms input as described by Vaswani et al. in \"Attention Is All You Need\", 2017.\n",
    "        dim     - dimension of positional embeddings.\n",
    "        max_len - maximal length of sequence, for precomputing\n",
    "        scale   - scale factor for frequency for positional embeddings\n",
    "        add     - boolean, if add is False, concatenate positional embeddings with input instead of adding\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        \n",
    "        self.dim = dim\n",
    "        self.add = add\n",
    "\n",
    "        # your code here\n",
    "               \n",
    "    def forward(self, input):\n",
    "        \"\"\"\n",
    "        input - [batch_size x sequence_len x features_dim]\n",
    "        If self.add is True, self.dim = featurs_dim.\n",
    "        Returns input with added or concatenated positional embeddings (depending on self.add).\n",
    "        \"\"\"\n",
    "        # your code here\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Test: draw positional encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "# visualize postional encodings for some parameters of PositionalEncoding\n",
    "# your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem: associative memory\n",
    "\n",
    "We have a lot of layers, let's solve some problem now!\n",
    "\n",
    "The task at hand is the multiplication (composition) of permutations.\n",
    "The length if permutations is fixed and equals perm_len.\n",
    "Input is an integer vector of length 2 x perm_len which contains two concatenated permutations p1 and p2.\n",
    "The product of p1 and p2 is a permutation p3 for which p3[i] = p1[p2[i]].\n",
    "The output of NN is also an integer vector of length 2 x perm_len in which first perm_lem elements are zero and the second perm_len elements are permutation p3.\n",
    "\n",
    "Example for perm_len = 5:\n",
    "```\n",
    "Input sequence:  3 4 2 1 0 1 3 0 2 4\n",
    "Output sequence: 0 0 0 0 0 4 1 3 2 0\n",
    "Clarification:  p1 = 3 4 2 1 0,    p2 = 1 3 0 2 4   =>    p3 = 4 1 3 2 0\n",
    "```\n",
    "\n",
    "Theoretically, such a problem can be solved by an ordinary LSTM, which will first memorize the permutation p1 in a hidden state, and then passing through the permutation p2 will produce the corresponding elements from the permutation p1.\n",
    "In practice, however, such a model works noticeably worse than a model with attention. A model with attention is explicitly learning by going through the p2 permutation to pay attention to the desired permutation element p1 and to output it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data generator\n",
    "\n",
    "The perm_generator function generates a batch of a given size of objects for training or a test.\n",
    "For each object in the batch permuations p1 and p2 of length perm_size are generated equiprobable.\n",
    "They form the input sequence [p1, p2] and the correct answer [0, p3] for it (see the example above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def perm_generator(batch_size, perm_size):\n",
    "    \"\"\"\n",
    "    Generates batch of batch_size objects.\n",
    "    Each object consists of two concatenated random permutations with length perm_size.\n",
    "    The target for the object is the product of its two permutations.\n",
    "    \"\"\"\n",
    "    perm1 = torch.cat([torch.randperm(perm_size).view(1, -1) for i in range(batch_size)], 0)\n",
    "    perm2 = torch.cat([torch.randperm(perm_size).view(1, -1) for i in range(batch_size)], 0)\n",
    "    correct = torch.zeros(batch_size, perm_size * 2, dtype=torch.long)\n",
    "    for i in range(batch_size):\n",
    "        for j in range(perm_size):\n",
    "            correct[i, j + perm_size] = perm1[i, perm2[i, j]]\n",
    "    return torch.cat([perm1, perm2], 1), correct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[1, 3, 0, 4, 2, 1, 4, 2, 0, 3],\n",
       "         [2, 4, 3, 1, 0, 4, 0, 1, 2, 3],\n",
       "         [2, 1, 0, 4, 3, 4, 2, 0, 1, 3]]),\n",
       " tensor([[0, 0, 0, 0, 0, 3, 2, 0, 1, 4],\n",
       "         [0, 0, 0, 0, 0, 0, 2, 4, 3, 1],\n",
       "         [0, 0, 0, 0, 0, 3, 0, 2, 1, 4]]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "perm_generator(3, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Putting all together: permutation multiplier model\n",
    "\n",
    "The model consists of the following layers:\n",
    "1. Embedding elements of the input sequence (using <a href=\"https://pytorch.org/docs/stable/nn.html#torch.nn.Embedding\">nn.Embedding</a>).\n",
    "2. Positional encoding.\n",
    "3. Linear layer to generate queries.\n",
    "4. A layer of attention. It receives the output of linear layer (3) as queries and sequence of embeddings (2) as objects.\n",
    "This layer must use the attention mask to ignore sequence elements from the future.\n",
    "5. Logistic regression with perm_len classes, outputs an integer for each position of answer.\n",
    "\n",
    "Please note that this model is neither a transformer nor a traditional network using LSTM with attention.\n",
    "The transformer uses K-head attention, elementwise transformation of embeddings.\n",
    "In traditional networks with LSTM, it used to obtain queries like linear layer (3) in the proposedarchitecture.\n",
    "But the result of query issued by LSTM at the previous time point affects the LSTM input at the next time point, therefore simultaneous parallel processing of the entire sequence is impossible.\n",
    "Also, most networks use the encoder-decoder architecture, where the encoder first reads the entire input sequence and forms its hidden representation, and then the decoder outputs the output sequence using this hidden representation and attention mechanism."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PermMultiplier(nn.Module):\n",
    "    def __init__(self, perm_len, embedding_dim, attention, pos_enc):\n",
    "        \"\"\"\n",
    "        perm_len       - permutation length (the input is twice longer)\n",
    "        embedding_dim  - dimensionality of integer embeddings\n",
    "        attention      - Attention object\n",
    "        pos_enc        - PositionalEncoder object\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.attention = attention\n",
    "        self.perm_len = perm_len\n",
    "        self.pos_enc = pos_enc\n",
    "        self.embeddings_dim = embedding_dim\n",
    "        \n",
    "        self.object_dim = embedding_dim\n",
    "        # if PositionalEncoding concatenates its encoding instead of\n",
    "        # adding it to the embeddings, then the dimensionality of\n",
    "        # the object after usch transformation increases\n",
    "        if not pos_enc.add:\n",
    "            self.object_dim += pos_enc.dim\n",
    "\n",
    "        # your code here\n",
    "\n",
    "    def forward(self, input):\n",
    "        \"\"\"\n",
    "        Perform forward pass through layers:\n",
    "        + get embeddings from input sequence (using both embeddings\n",
    "          and positional embeddings)\n",
    "        + run linear layer on embeddings\n",
    "        + use output of the previous layer as an attention queries\n",
    "        + attend on the embedded sequence using queries (note autoregressiveness)\n",
    "        + make final linear tranformation to obtain logits\n",
    "        \"\"\"\n",
    "        # your code here\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the model\n",
    "\n",
    "Complete the code for learning the model below.\n",
    "Find the right architecture and set of hyperparameters for the model and optimization method.\n",
    "For some architecture and hyperparameters the model can be trained on the CPU in a short time.\n",
    "\n",
    "If something doesn't work, try lesser `perm_len` first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "perm_len = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up model\n",
    "\n",
    "attention = ?\n",
    "pos_enc = PositionalEncoder(?, perm_len * 2, ?, ?)\n",
    "model = PermMultiplier(perm_len, ?, attention, pos_enc)\n",
    "if torch.cuda.is_available():\n",
    "    model = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up optimizer\n",
    "\n",
    "gd = optim.?(model.parameters(), lr=?)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# do optimization\n",
    "avg_loss = None\n",
    "forget = 0.99\n",
    "batch_size = 64\n",
    "iterator = range(?)\n",
    "for i in iterator:\n",
    "    gd.zero_grad()\n",
    "    batch = perm_generator(batch_size, perm_len)\n",
    "    if torch.cuda.is_available():\n",
    "        batch = batch[0].cuda(), batch[1].cuda()\n",
    "    # compute batch loss\n",
    "    # your code here\n",
    "    loss.backward()\n",
    "    if avg_loss is None:\n",
    "        avg_loss = float(loss)\n",
    "    else:\n",
    "        avg_loss = forget * avg_loss + (1 - forget) * float(loss)\n",
    "    descr_str = 'Iteration %05d, loss %.5f.' % (i, avg_loss)\n",
    "    print('\\r', descr_str, end='')\n",
    "    gd.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Test\n",
    "\n",
    "Great, we have some model.\n",
    "Let's check how it multiplies two random permutations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input:\n",
      " tensor([[2, 6, 0, 3, 4, 5, 8, 1, 7, 9, 2, 9, 4, 6, 0, 7, 5, 3, 8, 1],\n",
      "        [5, 7, 9, 3, 0, 2, 8, 1, 6, 4, 7, 4, 2, 8, 0, 6, 3, 9, 5, 1],\n",
      "        [2, 7, 5, 8, 0, 4, 9, 6, 3, 1, 7, 2, 0, 1, 5, 9, 3, 8, 4, 6],\n",
      "        [3, 2, 8, 5, 7, 0, 6, 4, 9, 1, 6, 9, 0, 8, 4, 5, 7, 3, 1, 2],\n",
      "        [2, 9, 5, 0, 8, 7, 3, 1, 6, 4, 2, 8, 4, 6, 0, 1, 7, 5, 9, 3]],\n",
      "       device='cuda:0')\n",
      "Output:\n",
      " tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9, 4, 8, 2, 1, 5, 3, 7, 6],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 9, 6, 5, 8, 3, 4, 2, 7],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 5, 2, 7, 4, 1, 8, 3, 0, 9],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 1, 3, 9, 7, 0, 4, 5, 2, 8],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 6, 8, 3, 2, 9, 1, 7, 4, 0]],\n",
      "       device='cuda:0')\n",
      "Correct:\n",
      " tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9, 4, 8, 2, 1, 5, 3, 7, 6],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 9, 6, 5, 8, 3, 4, 2, 7],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 5, 2, 7, 4, 1, 8, 3, 0, 9],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 1, 3, 9, 7, 0, 4, 5, 2, 8],\n",
      "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 6, 8, 3, 2, 9, 1, 7, 4, 0]],\n",
      "       device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "# time to check your model\n",
    "batch = perm_generator(batch_size, perm_len)\n",
    "if torch.cuda.is_available():\n",
    "    batch = batch[0].cuda(), batch[1].cuda()\n",
    "num_samples = 5\n",
    "print('Input:\\n', batch[0][:num_samples])\n",
    "print('Output:\\n', ?)\n",
    "print('Correct:\\n', batch[1][:num_samples])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Attention map\n",
    "\n",
    "One of the important for applications properties of attention is to show what the model pays attention to.\n",
    "For this, so-called attention maps are used.\n",
    "Use the last_weights field of the Attention layer to visualize which positions the trained model paid attention to at each point in time for permutations from the batch in the cell above.\n",
    "Expected behavior is that for each element of permutation p2 model pays attention to the corresponding element from the permutation p1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAScAAAD8CAYAAAA11GIZAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAEnJJREFUeJzt3X+s3XV9x/Hni1ugtlDFdbrRW6XE6mzIJuamoCyKFkdRA9niFnCSadj6z0D8sTncFlxYtmTO+GMJM7vij0WZDCtbGlOtU8HMRWvLjzjb2tkURy/FQBGR4aC99772xzllh8u993wv/X7v+Zz7fT2Sb3K+53z6Pu/etu9+Pp/v5/P9yjYREaU5adAJRETMJsUpIoqU4hQRRUpxiogipThFRJFSnCKiSClOEVGkFKeIKFKKU0QUaVkTQU/RqV7OytrjTp9Rf0yAZS94svaY0/sna48ZcdwTPM5RP6kTiXHx61b64Z9MVWp75/ee3GF784l830I1UpyWs5LzRn6j9rj/84ax2mMCvPCag7XHfPy1R2qPCUC2GwWw018/4RgP/2SK7+54UaW2I7/8w9Un/IUL1EhxiojyGZhmetBpzCnFKaKljDnmasO6QUhximix9JwiojjGTBU8h5niFNFi06Q4RURhDEwVXJwqLcKUtFnSfkkHJF3XdFIRsTimcaVjEPr2nCSNADcCbwAmgF2Sttne23RyEdEcA8cKnnOq0nPaCBywfdD2UeAW4LJm04qIphkzVfEYhCpzTmuAQz3nE8B5MxtJ2gJsAVjOilqSi4gGGabK7ThVKk6z7d95xm/J9jgwDrBKzy/4txwRcHyFeLmqFKcJYG3P+ShwuJl0ImLxiKlZ+x5lqFKcdgHrJa0D7gcuB97aaFYR0bjOhPgQFyfbk5KuBnYAI8CnbO9pPLOIaFRnndMQFycA29uB7Q3nEhGLbHqYe04RsTQtiZ5TRCw9RkwVfKfuFKeIFsuwLiKKY8RRjww6jTmlOEW0VGcRZhuHddP13/7ztC/srD0mwONfaCRsNOmk+v/H/5f7vl17TIDfOuvV9Qet6eE+mRCPiOLYYspt7DlFRPGm03OKiNJ0JsTLLQHlZhYRjWrvhHhEFG8q65wiojRZIR4RxZrO1bqIKE1n42+KU0QUxohj2b4SEaWxySLMiCiRsggzIspj0nOKiEJlQjwiimOUm81FRHk6j4YqtwSUm1lENKzsh2qWO+CMiEaZzgrxKkc/kjZL2i/pgKTrZvn8RZJul3S3pO9JemO/mClOES021e099TvmI2kEuBG4BNgAXCFpw4xmfw7cavtcOk8N//t+uWVYF9FSturaW7cROGD7IICkW4DLgL29Xwes6r5+LnC4X9AUp4iW6kyI17J9ZQ1wqOd8AjhvRpu/AL4q6RpgJXBRv6AZ1kW0Vuce4lUOYLWk3T3HlqcFeibPOL8C+IztUeCNwGclzVt/0nNqyEmnn95I3E/s+XIjcX//7AsbievJmh4TMlMDT/f5zdGNtcfsaOBnMPOf/rMMsYB1Tkdsj83x2QSwtud8lGcO264CNgPY/rak5cBq4MG5vjA9p4gWm+KkSkcfu4D1ktZJOoXOhPe2GW3uAzYBSHo5sBx4aL6g6TlFtFRdK8RtT0q6GtgBjACfsr1H0g3AbtvbgPcCn5D0bjqdtrfbnrf/l+IU0WJ1PeDA9nZg+4z3ru95vRe4YCExU5wiWsqGY9PlzuykOEW0VGdYl+IUEQUa6r11ktZ298Tsk7RH0rWLkVhENOv4UoIqxyBU6TlNAu+1fZek04E7Jf1bd4IrIobWkA/rbD8APNB9/ZikfXSWq6c4RQy5JXMPcUlnAecCO5tIJiIWT+dq3RJ4NJSk04AvAu+y/bNZPt8CbAFYzoraEoyIZiyJ2/RKOplOYbrZ9m2ztbE9DowDrNLza9j5ExFNG+phnSQBnwT22f5w8ylFxGJY4MbfRVdlqv4C4Erg9ZLu6R59b7EZEeWr6za9Tahyte5bzH6/logYYraYHOalBBGxdJU8rEtximip0uecUpwiWizFKSKKsyTWOUXE0jTU65zi2Zl+7LFG4l71ol9vJC6q/4EBAJ+571uNxL3qnPpXs0z97BkbH5Y0GyZzs7mIKFGGdRFRnMw5RUSxnOIUESXKhHhEFMfOnFNEFElM5WpdRJQoc04RUZzsrYuIMrkz71SqFKeIFsvVuogojjMhHhGlyrAuIoqUq3URURw7xSkiCpWlBBFRpMw5RURxjJgu+GpduZlFRONc8ehH0mZJ+yUdkHTdHG1+R9JeSXsk/VO/mOk5RbRVTRPikkaAG4E3ABPALknbbO/tabMeeD9wge1HJL2gX9z0nCLarJ6u00bggO2Dto8CtwCXzWjzB8CNth8BsP1gv6ApThEtZqvS0cca4FDP+UT3vV4vBV4q6T8kfUfS5n5BGxnWTb3kVB75u/W1xz3jTT+sPWZ0NXTZ5u1NPS2GIXpSihq4XF/DH5eB6enKua2WtLvnfNz2ePf1bEFmZrgMWA9cCIwC/y7pHNs/nesLM+cU0VYGqs85HbE9NsdnE8DanvNR4PAsbb5j+xhwr6T9dIrVrrm+MMO6iBazqx197ALWS1on6RTgcmDbjDb/CrwOQNJqOsO8g/MFTXGKaLMaJsRtTwJXAzuAfcCttvdIukHSpd1mO4CHJe0Fbgf+2PbD88XNsC6itSpNdldiezuwfcZ71/e8NvCe7lFJilNEm2X7SkQUx+DqV+sWXYpTRKuVW5wqT4hLGpF0t6QvNZlQRCyiujbXNWAhV+uupTMTHxFLxbAXJ0mjwJuAm5pNJyIWzfFFmFWOAag65/RR4H3A6XM1kLQF2AJwygtWnXhmEdG4km8217fnJOnNwIO275yvne1x22O2x5atWlFbghHRoGlVOwagSs/pAuBSSW8ElgOrJH3O9tuaTS0imqZh7jnZfr/tUdtn0dkz840UpogloOpk+IAKWNY5RbTW4Ca7q1hQcbJ9B3BHI5lExOIreFiXnlNEm00POoG5pThFtNXCbja36FKcIlqs5Kt1KU4RbVZwccqdMCOiSI30nEYOPNnIk1J2HL6n9pgAF5/5ikbiRgBF7xHJsC4iymMGtjWlihSniDZLzykiSpRhXUSUKcUpIoqU4hQRpZEzrIuIUuVqXUSUKD2niChTilNEFCdzThFRrBSniCiRCr7ZXO5KEBFFSs8pos0yrIuI4mRCPCKKleIUEUUquDhlQjyipUTnal2Vo28sabOk/ZIOSLpunnZvkWRJY/1ipjhFtJX/f/Nvv2M+kkaAG4FLgA3AFZI2zNLudOCdwM4q6aU4RbSZKx7z2wgcsH3Q9lHgFuCyWdr9JfBB4IkqqaU4RbRZ9eK0WtLunmNLT5Q1wKGe84nue0+RdC6w1vaXqqY2VBPieUrK8Fm2drSRuH/yzcp/xyv765dtrD0mgI8dbSRuHRawlOCI7bnmiWa778pTkSWdBHwEePtCckvPKaLN6hnWTQBre85HgcM956cD5wB3SPoRcD6wrd+k+FD1nCKiRq5tb90uYL2kdcD9wOXAW5/6GvtRYPXxc0l3AH9ke/d8QdNzimizGnpOtieBq4EdwD7gVtt7JN0g6dJnm1p6ThEtVtf2Fdvbge0z3rt+jrYXVomZ4hTRZgWvEE9ximirapPdA1NpzknS8yRtlfQDSfskvarpxCKiWaKeFeJNqdpz+hjwFdtvkXQKsKLBnCJikQz1LVMkrQJeQ3cBVXd5ermryiKiuoKLU5Vh3dnAQ8CnJd0t6SZJKxvOKyIWQz2LMBtRpTgtA14JfNz2ucDjwDNuiSBpy/F9N8d4suY0I6J2Nd2VoClVitMEMGH7+G0OttIpVk9je9z2mO2xkzm1zhwjoinD3HOy/WPgkKSXdd/aBOxtNKuIWBR13WyuCVWv1l0D3Ny9UncQeEdzKUXEYhnqq3UAtu8B+t5WMyKGSOGLMLNCPKLNUpwiojTHV4iXKsUposU0XW51SnGKaKvMOUVEqTKsi4gypThFXbSsmT+yz937zUbiXrl+UyNx/+rsJp7E07797Ok5RUSZUpwiojj1PX2lESlOES2VdU4RUS6XW51SnCJaLD2niChPFmFGRKkyIR4RRUpxiojymEyIR0SZMiEeEWVKcYqI0mQRZkSUyc7N5iKiUOXWphSniDbLsC4iymOg4GFdlceRR8RSVdPjyCVtlrRf0gFJ183y+Xsk7ZX0PUlfl/TifjFTnCJaTK52zBtDGgFuBC4BNgBXSNowo9ndwJjtXwW2Ah/sl1uKU0SLadqVjj42AgdsH7R9FLgFuKy3ge3bbf+8e/odYLRf0BSniLaqOqTrP6xbAxzqOZ/ovjeXq4Av9wuaCfEh48nJRuL+7toLGokLTzQUd3jsOHxP7TE3Xvzz/o366CzCrDwhvlrS7p7zcdvjPaFmmjWwpLcBY8Br+31hilNEm1W/K8ER22NzfDYBrO05HwUOz2wk6SLgz4DX2n6y3xemOEW02AJ6TvPZBayXtA64H7gceOvTvkc6F/gHYLPtB6sEzZxTRFvVNOdkexK4GtgB7ANutb1H0g2SLu02+1vgNOALku6RtK1feuk5RbRWfXvrbG8Hts947/qe1xctNGaKU0Sb5WZzEVGcPFQzIopVcM+p0oS4pHdL2iPp+5I+L2l504lFxCKoaW9dE/oWJ0lrgHfS2RdzDjBC51JhRAw5TU9XOgah6rBuGfAcSceAFcyywCoihoxZyCLMRde352T7fuBDwH3AA8Cjtr86s52kLZJ2S9p9jL6LPyNiwISRqx2DUGVYdwadHcbrgDOBld39MU9je9z2mO2xkzm1/kwjon52tWMAqkyIXwTca/sh28eA24BXN5tWRCyKgotTlTmn+4DzJa0A/hfYBOye/5dERPEKn3PqW5xs75S0FbgLmKRzR7vx+X9VRAyDQV2Jq6LS1TrbHwA+0HAuEbGoBjdkqyIrxCPayqQ4RUShyh3VpThFtNmg1jBVkeIU0WYpThFRHBumyh3XpTgBy37phbXHvPKb3609JsBnLzyvkbiTD/y4kbgBF5/5itpj/pcfridQek4RUaQUp4gojoGa7iHehBSniNYyOHNOEVEakwnxiChU5pwiokgpThFRnmz8jYgSGRj2W6ZExBKVnlNElCfbVyKiRAZnnVNEFCkrxCOiSJlzioji2LlaFxGFSs8pIspjPDU16CTmlOIU0VaF3zKlyuPII2Kp8nS1ow9JmyXtl3RA0nWzfH6qpH/ufr5T0ln9YqY4RbSUAU+70jEfSSPAjcAlwAbgCkkbZjS7CnjE9kuAjwB/0y+/FKeItrLr6jltBA7YPmj7KHALcNmMNpcB/9h9vRXYJEnzBc2cU0SL1TQhvgY41HM+Acx8EsdTbWxPSnoU+AXgyFxBGylOj/HIka95639XaLqaeZJbNA9Ublk536+89Nkm08+uqg3L+NlWN0z5lpDri080wGM8suNr3rq6YvPlknb3nI/bHu++nq0HNHMsWKXN0zRSnGz/YpV2knbbHmsihyYMU77DlCsMV77DlOt8bG+uKdQEsLbnfBQ4PEebCUnLgOcCP5kvaOacIuJE7QLWS1on6RTgcmDbjDbbgN/rvn4L8A17/hWgmXOKiBPSnUO6GtgBjACfsr1H0g3AbtvbgE8Cn5V0gE6P6fJ+cQddnMb7NynKMOU7TLnCcOU7TLkuCtvbge0z3ru+5/UTwG8vJKb69KwiIgYic04RUaSBFad+y91LIWmtpNsl7ZO0R9K1g86pCkkjku6W9KVB5zIfSc+TtFXSD7o/41cNOqf5SHp39+/B9yV9XtLyQee0VA2kOFVc7l6KSeC9tl8OnA/8YcG59roW2DfoJCr4GPAV278C/BoF5yxpDfBOYMz2OXQmf/tO7MazM6ieU5Xl7kWw/YDtu7qvH6Pzj2fNYLOan6RR4E3ATYPOZT6SVgGvoXMlB9tHbf90sFn1tQx4TnetzgqeuZ4najKo4jTbcvei/8EDdHdSnwvsHGwmfX0UeB9Q7m0OO84GHgI+3R2C3iRp5aCTmovt+4EPAffR2VfwqO2vDjarpWtQxWnBS9kHTdJpwBeBd9n+2aDzmYukNwMP2r5z0LlUsAx4JfBx2+cCjwMlzz+eQaeHvw44E1gp6W2DzWrpGlRxqrLcvRiSTqZTmG62fdug8+njAuBSST+iM1x+vaTPDTalOU0AE7aP90S30ilWpboIuNf2Q7aPAbcBrx5wTkvWoIpTleXuReje1uGTwD7bHx50Pv3Yfr/tUdtn0fm5fsN2kf+72/4xcEjSy7pvbQL2DjClfu4Dzpe0ovv3YhMFT+APu4GsEJ9rufsgcqngAuBK4D8l3dN970+7K2LjxF0D3Nz9T+og8I4B5zMn2zslbQXuonMV926yWrwxWSEeEUXKCvGIKFKKU0QUKcUpIoqU4hQRRUpxiogipThFRJFSnCKiSClOEVGk/wPYBtPsSgWbWAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualize attention map for some object\n",
    "# your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Misc\n",
    "\n",
    "What now? You can:\n",
    "+ implement and test `AdditiveScore` with learnable parameters $W1$, $W2$ and $w3$\n",
    "+ implement and test K-headed attention\n",
    "+ compare attention-based model with baseline, e. g. LSTM\n",
    "+ replace linear layer (3) in the model with LSTM (autoregressive or not) and try to train such model\n",
    "+ ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# play with model and learn something new about attention!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}