Positional_Encoding_in_Transformer_neural_networks.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Positional Encoding\n",
        "\n",
        "This notebook will code positional encoding for Transformer neural networks with pytrch"
      ],
      "metadata": {
        "id": "NMNu7Jg5troY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "max_sequence_length = 10\n",
        "d_model = 6"
      ],
      "metadata": {
        "id": "H3iqZxn20a7m"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "$$\n",
        "PE(\\text{position}, 2i) = \\sin\\bigg( \\frac{ \\text{position} }{10000^\\frac{2i}{d_{model}}} \\bigg)\n",
        "$$\n",
        "\n",
        "$$\n",
        "PE(\\text{position}, 2i+1) = \\cos\\bigg( \\frac{ \\text{position} }{10000^\\frac{2i}{d_{model}}} \\bigg)\n",
        "$$\n",
        "\n",
        "We can rewrite these as\n",
        "\n",
        "$$\n",
        "PE(\\text{position}, i) = \\sin\\bigg( \\frac{ \\text{position} }{10000^\\frac{i}{d_{model}}} \\bigg) \\text{ when i is even}\n",
        "$$\n",
        "\n",
        "$$\n",
        "PE(\\text{position}, i) = \\cos\\bigg( \\frac{ \\text{position} }{10000^\\frac{i-1}{d_{model}}} \\bigg) \\text{ when i is odd}\n",
        "$$"
      ],
      "metadata": {
        "id": "2aUNCBlKvxew"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "even_i = torch.arange(0, d_model, 2).float()\n",
        "even_i"
      ],
      "metadata": {
        "id": "j3SWE1Nxwo-D",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5d9482c7-79ad-46cc-ce14-8e7ad7e335b2"
      },
      "execution_count": 75,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0., 2., 4.])"
            ]
          },
          "metadata": {},
          "execution_count": 75
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "even_denominator = torch.pow(10000, even_i/d_model)\n",
        "even_denominator"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h-RWR30KxdLM",
        "outputId": "123b4f15-8cba-4eb7-9f5d-faabe2dbc286"
      },
      "execution_count": 76,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([  1.0000,  21.5443, 464.1590])"
            ]
          },
          "metadata": {},
          "execution_count": 76
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "odd_i = torch.arange(1, d_model, 2).float()\n",
        "odd_i"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iITvtjTt6jO-",
        "outputId": "2a9f433f-14b6-4a42-fcb2-4b896e215de0"
      },
      "execution_count": 77,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([1., 3., 5.])"
            ]
          },
          "metadata": {},
          "execution_count": 77
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "even_denominator = torch.pow(10000, (odd_i - 1)/d_model)\n",
        "even_denominator"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AAIVnPLJ1JYC",
        "outputId": "67ca824f-82b1-41d8-fac5-84ff84999349"
      },
      "execution_count": 78,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([  1.0000,  21.5443, 464.1590])"
            ]
          },
          "metadata": {},
          "execution_count": 78
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "`even_denominator` and `odd_denominator` are the same! So we can just do one of these actions and call the resulting variable `denominator`"
      ],
      "metadata": {
        "id": "LBtyLN2NAtEC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "denominator = even_denominator"
      ],
      "metadata": {
        "id": "IyjRI1imBA3F"
      },
      "execution_count": 79,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "position = torch.arange(max_sequence_length, dtype=torch.float).reshape(max_sequence_length, 1)"
      ],
      "metadata": {
        "id": "WwxxoNSN-me9"
      },
      "execution_count": 80,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "position"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1nfvUzG8-rMK",
        "outputId": "3bc44cf5-0c38-43ae-bee2-7cd86f078601"
      },
      "execution_count": 81,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0.],\n",
              "        [1.],\n",
              "        [2.],\n",
              "        [3.],\n",
              "        [4.],\n",
              "        [5.],\n",
              "        [6.],\n",
              "        [7.],\n",
              "        [8.],\n",
              "        [9.]])"
            ]
          },
          "metadata": {},
          "execution_count": 81
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "even_PE = torch.sin(position / denominator)\n",
        "odd_PE = torch.cos(position / denominator)"
      ],
      "metadata": {
        "id": "XEm-9niG4VEl"
      },
      "execution_count": 82,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "even_PE"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Uqfc01YJ43w6",
        "outputId": "b151d8fe-ce1c-4bfc-bb57-1e3ec54aae49"
      },
      "execution_count": 83,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 0.0000,  0.0000,  0.0000],\n",
              "        [ 0.8415,  0.0464,  0.0022],\n",
              "        [ 0.9093,  0.0927,  0.0043],\n",
              "        [ 0.1411,  0.1388,  0.0065],\n",
              "        [-0.7568,  0.1846,  0.0086],\n",
              "        [-0.9589,  0.2300,  0.0108],\n",
              "        [-0.2794,  0.2749,  0.0129],\n",
              "        [ 0.6570,  0.3192,  0.0151],\n",
              "        [ 0.9894,  0.3629,  0.0172],\n",
              "        [ 0.4121,  0.4057,  0.0194]])"
            ]
          },
          "metadata": {},
          "execution_count": 83
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "even_PE.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Fjmx56D25A5T",
        "outputId": "9615f471-3bc5-445e-d229-09bcd93d239e"
      },
      "execution_count": 84,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([10, 3])"
            ]
          },
          "metadata": {},
          "execution_count": 84
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "odd_PE"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "J8TlRfY745hA",
        "outputId": "bfd9bd54-009a-4cb5-c682-88bffc6b7a4e"
      },
      "execution_count": 85,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 1.0000,  1.0000,  1.0000],\n",
              "        [ 0.5403,  0.9989,  1.0000],\n",
              "        [-0.4161,  0.9957,  1.0000],\n",
              "        [-0.9900,  0.9903,  1.0000],\n",
              "        [-0.6536,  0.9828,  1.0000],\n",
              "        [ 0.2837,  0.9732,  0.9999],\n",
              "        [ 0.9602,  0.9615,  0.9999],\n",
              "        [ 0.7539,  0.9477,  0.9999],\n",
              "        [-0.1455,  0.9318,  0.9999],\n",
              "        [-0.9111,  0.9140,  0.9998]])"
            ]
          },
          "metadata": {},
          "execution_count": 85
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "odd_PE.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Bazd5CSZ948R",
        "outputId": "396408c7-26b1-4268-b530-bcfa32a95d69"
      },
      "execution_count": 86,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([10, 3])"
            ]
          },
          "metadata": {},
          "execution_count": 86
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "stacked = torch.stack([even_PE, odd_PE], dim=2)\n",
        "stacked.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "X0n6V1nk_Pgg",
        "outputId": "0b9c0f1e-3a48-4993-a295-5634b73dd82c"
      },
      "execution_count": 87,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([10, 3, 2])"
            ]
          },
          "metadata": {},
          "execution_count": 87
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "PE = torch.flatten(stacked, start_dim=1, end_dim=2)\n",
        "PE"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LJAGzwSF_fVV",
        "outputId": "62e15b26-d4b0-433d-8fd0-fc10b5e7311b"
      },
      "execution_count": 88,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 0.0000,  1.0000,  0.0000,  1.0000,  0.0000,  1.0000],\n",
              "        [ 0.8415,  0.5403,  0.0464,  0.9989,  0.0022,  1.0000],\n",
              "        [ 0.9093, -0.4161,  0.0927,  0.9957,  0.0043,  1.0000],\n",
              "        [ 0.1411, -0.9900,  0.1388,  0.9903,  0.0065,  1.0000],\n",
              "        [-0.7568, -0.6536,  0.1846,  0.9828,  0.0086,  1.0000],\n",
              "        [-0.9589,  0.2837,  0.2300,  0.9732,  0.0108,  0.9999],\n",
              "        [-0.2794,  0.9602,  0.2749,  0.9615,  0.0129,  0.9999],\n",
              "        [ 0.6570,  0.7539,  0.3192,  0.9477,  0.0151,  0.9999],\n",
              "        [ 0.9894, -0.1455,  0.3629,  0.9318,  0.0172,  0.9999],\n",
              "        [ 0.4121, -0.9111,  0.4057,  0.9140,  0.0194,  0.9998]])"
            ]
          },
          "metadata": {},
          "execution_count": 88
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Class\n",
        "\n",
        "Let's combine all the code above into a cute class"
      ],
      "metadata": {
        "id": "Liidl3ggt0wK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class PositionalEncoding(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model, max_sequence_length):\n",
        "        super().__init__()\n",
        "        self.max_sequence_length = max_sequence_length\n",
        "        self.d_model = d_model\n",
        "\n",
        "    def forward(self):\n",
        "        even_i = torch.arange(0, self.d_model, 2).float()\n",
        "        denominator = torch.pow(10000, even_i/self.d_model)\n",
        "        position = torch.arange(self.max_sequence_length).reshape(self.max_sequence_length, 1)\n",
        "        even_PE = torch.sin(position / denominator)\n",
        "        odd_PE = torch.cos(position / denominator)\n",
        "        stacked = torch.stack([even_PE, odd_PE], dim=2)\n",
        "        PE = torch.flatten(stacked, start_dim=1, end_dim=2)\n",
        "        return PE"
      ],
      "metadata": {
        "id": "E1G1ziOa6SdZ"
      },
      "execution_count": 91,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pe = PositionalEncoding(d_model=6, max_sequence_length=10)\n",
        "pe.forward()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ENHY3b-BDgL9",
        "outputId": "6413d8c5-7fbd-48b4-dfd6-bbae0396bd0c"
      },
      "execution_count": 92,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 0.0000,  1.0000,  0.0000,  1.0000,  0.0000,  1.0000],\n",
              "        [ 0.8415,  0.5403,  0.0464,  0.9989,  0.0022,  1.0000],\n",
              "        [ 0.9093, -0.4161,  0.0927,  0.9957,  0.0043,  1.0000],\n",
              "        [ 0.1411, -0.9900,  0.1388,  0.9903,  0.0065,  1.0000],\n",
              "        [-0.7568, -0.6536,  0.1846,  0.9828,  0.0086,  1.0000],\n",
              "        [-0.9589,  0.2837,  0.2300,  0.9732,  0.0108,  0.9999],\n",
              "        [-0.2794,  0.9602,  0.2749,  0.9615,  0.0129,  0.9999],\n",
              "        [ 0.6570,  0.7539,  0.3192,  0.9477,  0.0151,  0.9999],\n",
              "        [ 0.9894, -0.1455,  0.3629,  0.9318,  0.0172,  0.9999],\n",
              "        [ 0.4121, -0.9111,  0.4057,  0.9140,  0.0194,  0.9998]])"
            ]
          },
          "metadata": {},
          "execution_count": 92
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Happy Coding!"
      ],
      "metadata": {
        "id": "VjPIbLcBt6l4"
      }
    }
  ]
}